BitArray.h

/**
 * @file    src/BitArray.h
 * @author  MohammadJavad Rezaei Seraji <mjrezaei (at) ce.sharif.edu>
 *
 * @section LICENCE
 *
 *
 * Copyright (C) 2017-2020
 *   MohammadJavad Rezaei Seraji <mjrezaei (at) ce.sharif.edu>
 *	 Seyed Abolfazl Motahari <motahari (at) sharif.edu
 *
 * Permission is hereby granted, free of charge, to any person
 * obtaining a copy of this software and associated documentation
 * files (the "Software"), to deal in the Software without
 * restriction, including without limitation the rights to use,
 * copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the
 * Software is furnished to do so, subject to the following
 * conditions:
 *
 * The above copyright notice and this permission notice shall be
 * included in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
 * OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
 * HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
 * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
 * OTHER DEALINGS IN THE SOFTWARE.
 **/

#include <string>
#include <algorithm>
#include <stdlib.h>
#include <stdio.h>
#include <iostream>
#define SA_SAMPLERATE	128
#define BT_SAMPLERATE	8

using namespace std;

struct BitArray{
	static const int blockLog2 = 6;
	static const int block = ( 1 << blockLog2 ) ;
	uint32_t smRate;
    uint64_t* arr;
	uint8_t* pc;
    uint32_t* sum;
    uint32_t sz;
	uint64_t needMask[64];
	BitArray(){
		smRate = 0;
		needMask[0] = 1;
		arr = NULL, sum = NULL, pc = NULL;
		for( int i = 1; i < 64; i++ )
			needMask[i] = needMask[i - 1] + ( 1ULL << i );
		sz = 0;
	}
	
	int myPopcount(unsigned long long x)
	{
		x = (x & 0x5555555555555555ULL) + ((x >> 1) & 0x5555555555555555ULL);
		x = (x & 0x3333333333333333ULL) + ((x >> 2) & 0x3333333333333333ULL);
		x = (x & 0x0F0F0F0F0F0F0F0FULL) + ((x >> 4) & 0x0F0F0F0F0F0F0F0FULL);
		return (x * 0x0101010101010101ULL) >> 56;
	}
	
    void reset( uint64_t aS ){
        sz = ( aS + block - 1 ) / block;
        arr = new uint64_t[sz];
        for( uint64_t i = 0; i < sz; i++ )
			arr[i] = 0;
    }
	
    void setBit( uint32_t pos ){
        uint32_t idx = pos >> blockLog2;
        arr[idx] |= ( 1ULL << ( pos & ( block - 1 ) ) );
    }
	
    void setSum(){
		smRate = BT_SAMPLERATE;
		unsigned int smSize = ( sz + smRate - 1 ) / smRate;
        sum = new unsigned int[smSize];
		pc = new uint8_t[sz];
		for( uint32_t i = 0; i < sz; i++ )
			pc[i] = myPopcount( arr[i] );
        for( uint64_t i = 0; i < smSize; i++ )
			sum[i] = 0;
        sum[0] = myPopcount(arr[0]);
        for( long long i = smRate; i < sz; i += smRate ){
			uint32_t idx = i / smRate;
			sum[idx] = sum[idx - 1];
			for( unsigned int j = ( idx - 1 ) * smRate + 1; j <= i; j++ )
				sum[idx] += myPopcount(arr[j]);
		}
    }
	
    int getPos( uint32_t pos ){
        uint32_t idx = pos >> blockLog2;
        return ( arr[idx] & ( 1ULL << ( pos & ( block - 1 ) ) ) ) ? 1 : 0;
    }
	
    inline uint32_t getRank( uint32_t pos ){
        int idx = ( pos >> blockLog2 );
		if( sz == 0 )
			return 0;
		uint32_t ret;
		try{
			ret = myPopcount( arr[idx] & needMask[( pos & ( block - 1 ) )] );
		}catch( const std::exception& e ){
			 std::cout << e.what() << endl;
		}
		if( idx == 0 )
			return ret;
		int pre = idx - 1;
		while( pre >= 0 && ( pre & ( smRate - 1 ) ) ){
			//ret += myPopcount( arr[pre] );
			//cout << pre << ' ' << ret << endl;
			ret += pc[pre];
			pre--;
		}
		//cout << "fuck " << pre << ' ' << ret << endl;
		if( pre >= 0 && ( pre & ( smRate - 1 ) ) == 0 ){
			ret += sum[pre / smRate];
		}
        return ret;
    }
	
	void save( FILE* fout ){
		fwrite((const void*)& sz,sizeof(uint32_t),1,fout);
		if( sz > 0 ){
			for( uint32_t i = 0; i < sz; i++ )
				fwrite((const void*)& arr[i],sizeof(uint64_t),1,fout);
			fwrite((const void*)& smRate,sizeof(uint32_t),1,fout);
			if( smRate > 0 ){
				unsigned int smSize = ( sz + smRate - 1 ) / smRate;
				for( uint32_t i = 0; i < smSize; i++ )
					fwrite((const void*)& sum[i],sizeof(uint32_t),1,fout);
			}
		}
	}
	
	void load( FILE* fin ){
		pc = NULL;
		fread( &sz, sizeof( uint32_t ), 1, fin );
		if( sz > 0 ){
			arr = new uint64_t[sz];
			fread( arr, sizeof( uint64_t ), sz, fin );
			fread( &smRate, sizeof( uint32_t ), 1, fin );
			if( smRate > 0 ){
				unsigned int smSize = ( sz + smRate - 1 ) / smRate;
				sum = new uint32_t[smSize];
				fread( sum, sizeof( uint32_t ), smSize, fin );
				pc = new uint8_t[sz];
				for( uint32_t i = 0; i < sz; i++ )
					pc[i] = myPopcount( arr[i] );
			}
		}
	}
	
	long long getBytes(){
		long long ret = sz / 8;
		if( sum != NULL )
			ret += ( sz / smRate ) * 4;
		return ret;
	}
};