bitvector.h

#pragma once
 
#include <pixie/bits.h>
#include <pixie/cache_line.h>
 
#include <algorithm>
#include <bit>
#include <cstdint>
#include <span>
#include <string>
#include <vector>
 
#ifdef PIXIE_DIAGNOSTICS
#include <spdlog/spdlog.h>
#endif
 
namespace pixie {

class BitVector {
 private:
  constexpr static size_t kWordSize = 64;
  constexpr static size_t kSuperBlockRankIntSize = 64;
  constexpr static size_t kBasicBlockRankIntSize = 16;
  constexpr static size_t kBasicBlockSize = 512;
  constexpr static size_t kWordsPerBlock = 8;
  constexpr static size_t kSuperBlockSize = 65536;
  constexpr static size_t kBlocksPerSuperBlock = 128;
  constexpr static size_t kSelectSampleFrequency = 16384;
 
  alignas(64) uint64_t delta_super[8];
  alignas(64) uint16_t delta_basic[32];
 
  AlignedStorage super_block_rank_;  // 64-bit global prefix sums
  AlignedStorage basic_block_rank_;  // 16-bit local prefix sums
  AlignedStorage select1_samples_;   // 64-bit global positions
  AlignedStorage select0_samples_;   // 64-bit global positions
  const size_t num_bits_;
  const size_t padded_size_;
  size_t max_rank_;
 
  std::span<const uint64_t> bits_;

  void build_rank() {
    size_t num_superblocks = 8 + (padded_size_ - 1) / kSuperBlockSize;
    // Add more blocks to ease SIMD processing
    // num_basicblocks to fully cover superblock, i.e. 128
    // This reduces branching in select
    num_superblocks = ((num_superblocks + 7) / 8) * 8;
    size_t num_basicblocks = num_superblocks * kBlocksPerSuperBlock;
    super_block_rank_.resize(num_superblocks * 64);
    basic_block_rank_.resize(num_basicblocks * 16);
 
    auto super_block_rank = super_block_rank_.As64BitInts();
    auto basic_block_rank = basic_block_rank_.As16BitInts();
 
    uint64_t super_block_sum = 0;
    uint16_t basic_block_sum = 0;
 
    for (size_t i = 0; i / kBasicBlockSize < basic_block_rank.size();
         i += kWordSize) {
      if (i % kSuperBlockSize == 0) {
        super_block_sum += basic_block_sum;
        super_block_rank[i / kSuperBlockSize] = super_block_sum;
        basic_block_sum = 0;
      }
      if (i % kBasicBlockSize == 0) {
        basic_block_rank[i / kBasicBlockSize] = basic_block_sum;
      }
      if (i / kWordSize < bits_.size()) {
        basic_block_sum += std::popcount(bits_[i / kWordSize]);
      }
    }
    max_rank_ = super_block_sum + basic_block_sum;
  }

  void build_select() {
    uint64_t milestone = kSelectSampleFrequency;
    uint64_t milestone0 = kSelectSampleFrequency;
    uint64_t rank = 0;
    uint64_t rank0 = 0;
 
    size_t num_one_samples =
        1 + (max_rank_ + kSelectSampleFrequency - 1) / kSelectSampleFrequency;
    size_t num_zero_samples =
        1 + (num_bits_ - max_rank_ + kSelectSampleFrequency - 1) /
                kSelectSampleFrequency;
 
    select1_samples_.resize(num_one_samples * 64);
    select0_samples_.resize(num_zero_samples * 64);
    auto select1_samples = select1_samples_.As64BitInts();
    auto select0_samples = select0_samples_.As64BitInts();
 
    select1_samples[0] = 0;
    select0_samples[0] = 0;
 
    size_t num_zeros = 1, num_ones = 1;
 
    for (size_t i = 0; i < bits_.size(); ++i) {
      auto ones = std::popcount(bits_[i]);
      auto zeros = 64 - ones;
      if (rank + ones >= milestone) {
        auto pos = select_64(bits_[i], milestone - rank - 1);
        // TODO: try including global rank into select samples to save
        //       a cache miss on global rank scan
        select1_samples[num_ones++] = (64 * i + pos) / kSuperBlockSize;
        milestone += kSelectSampleFrequency;
      }
      if (rank0 + zeros >= milestone0) {
        auto pos = select_64(~bits_[i], milestone0 - rank0 - 1);
        select0_samples[num_zeros++] = (64 * i + pos) / kSuperBlockSize;
        milestone0 += kSelectSampleFrequency;
      }
      rank += ones;
      rank0 += zeros;
    }
 
    for (size_t i = 0; i < 8; ++i) {
      delta_super[i] = i * kSuperBlockSize;
    }
    for (size_t i = 0; i < 32; ++i) {
      delta_basic[i] = i * kBasicBlockSize;
    }
  }

  uint64_t find_superblock(uint64_t rank) const {
    auto select1_samples = select1_samples_.AsConst64BitInts();
    auto super_block_rank = super_block_rank_.AsConst64BitInts();
 
    uint64_t left = select1_samples[rank / kSelectSampleFrequency];
 
    while (left + 7 < super_block_rank.size()) {
      auto len = lower_bound_8x64(&super_block_rank[left], rank);
      if (len < 8) {
        return left + len - 1;
      }
      left += 8;
    }
    if (left + 3 < super_block_rank.size()) {
      auto len = lower_bound_4x64(&super_block_rank[left], rank);
      if (len < 4) {
        return left + len - 1;
      }
      left += 4;
    }
    while (left < super_block_rank.size() && super_block_rank[left] < rank) {
      left++;
    }
    return left - 1;
  }

  uint64_t find_superblock_zeros(uint64_t rank0) const {
    auto select0_samples = select0_samples_.AsConst64BitInts();
    auto super_block_rank = super_block_rank_.AsConst64BitInts();
 
    uint64_t left = select0_samples[rank0 / kSelectSampleFrequency];
 
    while (left + 7 < super_block_rank.size()) {
      auto len = lower_bound_delta_8x64(&super_block_rank[left], rank0,
                                        delta_super, kSuperBlockSize * left);
      if (len < 8) {
        return left + len - 1;
      }
      left += 8;
    }
    if (left + 3 < super_block_rank.size()) {
      auto len = lower_bound_delta_4x64(&super_block_rank[left], rank0,
                                        delta_super, kSuperBlockSize * left);
      if (len < 4) {
        return left + len - 1;
      }
      left += 4;
    }
    while (left < super_block_rank.size() &&
           kSuperBlockSize * left - super_block_rank[left] < rank0) {
      left++;
    }
    return left - 1;
  }

  uint64_t find_basicblock(uint16_t local_rank, uint64_t s_block) const {
    auto basic_block_rank = basic_block_rank_.AsConst16BitInts();
 
    for (size_t pos = 0; pos < kBlocksPerSuperBlock; pos += 32) {
      auto count = lower_bound_32x16(
          &basic_block_rank[kBlocksPerSuperBlock * s_block + pos], local_rank);
      if (count < 32) {
        return kBlocksPerSuperBlock * s_block + pos + count - 1;
      }
    }
    return kBlocksPerSuperBlock * s_block + kBlocksPerSuperBlock - 1;
  }

  uint64_t find_basicblock_zeros(uint16_t local_rank0, uint64_t s_block) const {
    auto basic_block_rank = basic_block_rank_.AsConst16BitInts();
    for (size_t pos = 0; pos < kBlocksPerSuperBlock; pos += 32) {
      auto count = lower_bound_delta_32x16(
          &basic_block_rank[kBlocksPerSuperBlock * s_block + pos], local_rank0,
          delta_basic, kBasicBlockSize * pos);
      if (count < 32) {
        return kBlocksPerSuperBlock * s_block + pos + count - 1;
      }
    }
    return kBlocksPerSuperBlock * s_block + kBlocksPerSuperBlock - 1;
  }

  uint64_t find_basicblock_is(uint16_t local_rank, uint64_t s_block) const {
    auto super_block_rank = super_block_rank_.AsConst64BitInts();
    auto basic_block_rank = basic_block_rank_.AsConst16BitInts();
 
    auto lower = super_block_rank[s_block];
    auto upper = super_block_rank[s_block + 1];
 
    uint64_t pos = kBlocksPerSuperBlock * local_rank / (upper - lower);
    pos = pos + 16 < 32 ? 0 : (pos - 16);
    pos = pos > 96 ? 96 : pos;
    while (pos < 96) {
      auto count = lower_bound_32x16(
          &basic_block_rank[kBlocksPerSuperBlock * s_block + pos], local_rank);
      if (count == 0) {
        return find_basicblock(local_rank, s_block);
      }
      if (count < 32) {
        return kBlocksPerSuperBlock * s_block + pos + count - 1;
      }
      pos += 32;
    }
    pos = 96;
    auto count = lower_bound_32x16(
        &basic_block_rank[kBlocksPerSuperBlock * s_block + pos], local_rank);
    if (count == 0) {
      return find_basicblock(local_rank, s_block);
    }
    return kBlocksPerSuperBlock * s_block + pos + count - 1;
  }

  uint64_t find_basicblock_is_zeros(uint16_t local_rank0,
                                    uint64_t s_block) const {
    auto super_block_rank = super_block_rank_.AsConst64BitInts();
    auto basic_block_rank = basic_block_rank_.AsConst16BitInts();
 
    auto lower = kSuperBlockSize * s_block - super_block_rank[s_block];
    auto upper =
        kSuperBlockSize * (s_block + 1) - super_block_rank[s_block + 1];
 
    uint64_t pos = kBlocksPerSuperBlock * local_rank0 / (upper - lower);
    pos = pos + 16 < 32 ? 0 : (pos - 16);
    pos = pos > 96 ? 96 : pos;
    while (pos < 96) {
      auto count = lower_bound_delta_32x16(
          &basic_block_rank[kBlocksPerSuperBlock * s_block + pos], local_rank0,
          delta_basic, kBasicBlockSize * pos);
      if (count == 0) {
        return find_basicblock_zeros(local_rank0, s_block);
      }
      if (count < 32) {
        return kBlocksPerSuperBlock * s_block + pos + count - 1;
      }
      pos += 32;
    }
    pos = 96;
    auto count = lower_bound_delta_32x16(
        &basic_block_rank[kBlocksPerSuperBlock * s_block + pos], local_rank0,
        delta_basic, kBasicBlockSize * pos);
    if (count == 0) {
      return find_basicblock_zeros(local_rank0, s_block);
    }
    return kBlocksPerSuperBlock * s_block + pos + count - 1;
  }
 
 public:
#ifdef PIXIE_DIAGNOSTICS
  struct DiagnosticsBytes {
    size_t source_bitvector_bytes = 0;
    size_t super_block_rank_bytes = 0;
    size_t basic_block_rank_bytes = 0;
    size_t select1_samples_bytes = 0;
    size_t select0_samples_bytes = 0;
    size_t total_bytes = 0;
  };

  DiagnosticsBytes diagnostics_bytes() const {
    DiagnosticsBytes result;
    result.source_bitvector_bytes = (num_bits_ + 7) / 8;
    result.super_block_rank_bytes = super_block_rank_.AsConstBytes().size();
    result.basic_block_rank_bytes = basic_block_rank_.AsConstBytes().size();
    result.select1_samples_bytes = select1_samples_.AsConstBytes().size();
    result.select0_samples_bytes = select0_samples_.AsConstBytes().size();
    result.total_bytes =
        result.super_block_rank_bytes + result.basic_block_rank_bytes +
        result.select1_samples_bytes + result.select0_samples_bytes;
    return result;
  }

  void memory_report() const {
    const auto diagnostics = diagnostics_bytes();
    const double source_bytes =
        static_cast<double>(diagnostics.source_bitvector_bytes);
    const auto log_bytes = [&](std::string_view label, size_t bytes) {
      const double percentage =
          source_bytes > 0.0 ? 100.0 * static_cast<double>(bytes) / source_bytes
                             : 0.0;
      spdlog::info("BitVector {}: {} bytes ({:.2f}% of source)", label, bytes,
                   percentage);
    };
    log_bytes("source_bitvector", diagnostics.source_bitvector_bytes);
    log_bytes("super_block_rank", diagnostics.super_block_rank_bytes);
    log_bytes("basic_block_rank", diagnostics.basic_block_rank_bytes);
    log_bytes("select1_samples", diagnostics.select1_samples_bytes);
    log_bytes("select0_samples", diagnostics.select0_samples_bytes);
    log_bytes("total", diagnostics.total_bytes);
  }
#endif
  explicit BitVector(std::span<const uint64_t> bit_vector, size_t num_bits)
      : num_bits_(std::min(num_bits, bit_vector.size() * kWordSize)),
        padded_size_(((num_bits_ + kWordSize - 1) / kWordSize) * kWordSize),
        bits_(bit_vector) {
    build_rank();
    build_select();
  }

  size_t size() const { return num_bits_; }

  int operator[](size_t pos) const {
    size_t word_idx = pos / kWordSize;
    size_t bit_off = pos % kWordSize;
 
    return (bits_[word_idx] >> bit_off) & 1;
  }

  uint64_t rank(size_t pos) const {
    if (pos >= bits_.size() * kWordSize) [[unlikely]] {
      return max_rank_;
    }
 
    auto super_block_rank = super_block_rank_.AsConst64BitInts();
    auto basic_block_rank = basic_block_rank_.AsConst16BitInts();
 
    uint64_t b_block = pos / kBasicBlockSize;
    uint64_t s_block = pos / kSuperBlockSize;
    // Precomputed rank
    uint64_t result = super_block_rank[s_block] + basic_block_rank[b_block];
    // Basic block tail
    result += rank_512(&bits_[b_block * kWordsPerBlock],
                       pos - (b_block * kBasicBlockSize));
    return result;
  }

  uint64_t rank0(size_t pos) const {
    if (pos >= bits_.size() * kWordSize) [[unlikely]] {
      return bits_.size() * kWordSize - max_rank_;
    }
    return pos - rank(pos);
  }

  uint64_t select(size_t rank) const {
    if (rank > max_rank_) [[unlikely]] {
      return num_bits_;
    }
    if (rank == 0) [[unlikely]] {
      return 0;
    }
    auto super_block_rank = super_block_rank_.AsConst64BitInts();
    auto basic_block_rank = basic_block_rank_.AsConst16BitInts();
 
    uint64_t s_block = find_superblock(rank);
    rank -= super_block_rank[s_block];
    auto pos = find_basicblock_is(rank, s_block);
    rank -= basic_block_rank[pos];
    pos *= kWordsPerBlock;
 
    // Final search
    if (pos + kWordsPerBlock - 1 < kWordsPerBlock) [[unlikely]] {
      size_t ones = std::popcount(bits_[pos]);
      while (pos < bits_.size() && ones < rank) {
        rank -= ones;
        ones = std::popcount(bits_[++pos]);
      }
      return kWordSize * pos + select_64(bits_[pos], rank - 1);
    }
    return kWordSize * pos + select_512(&bits_[pos], rank - 1);
  }

  uint64_t select0(size_t rank0) const {
    if (rank0 > num_bits_ - max_rank_) [[unlikely]] {
      return num_bits_;
    }
    if (rank0 == 0) [[unlikely]] {
      return 0;
    }
    auto super_block_rank = super_block_rank_.AsConst64BitInts();
    auto basic_block_rank = basic_block_rank_.AsConst16BitInts();
 
    uint64_t s_block = find_superblock_zeros(rank0);
    rank0 -= kSuperBlockSize * s_block - super_block_rank[s_block];
    auto pos = find_basicblock_is_zeros(rank0, s_block);
    auto pos_in_super_block = pos & (kBlocksPerSuperBlock - 1);
    rank0 -= kBasicBlockSize * pos_in_super_block - basic_block_rank[pos];
    pos *= kWordsPerBlock;
 
    // Final search
    if (pos + kWordsPerBlock - 1 < kWordsPerBlock) [[unlikely]] {
      size_t zeros = std::popcount(~bits_[pos]);
      while (pos < bits_.size() && zeros < rank0) {
        rank0 -= zeros;
        zeros = std::popcount(~bits_[++pos]);
      }
      return kWordSize * pos + select_64(~bits_[pos], rank0 - 1);
    }
    return kWordSize * pos + select0_512(&bits_[pos], rank0 - 1);
  }

  std::string to_string() const {
    std::string result;
    result.reserve(num_bits_);
 
    for (size_t i = 0; i < num_bits_; i++) {
      result.push_back(operator[](i) ? '1' : '0');
    }
 
    return result;
  }
};

class BitVectorInterleaved {
 private:
  constexpr static size_t kWordSize = 64;
  constexpr static size_t kSuperBlockRankIntSize = 64;
  constexpr static size_t kBasicBlockRankIntSize = 16;
  constexpr static size_t kBasicBlockSize = 496;
  constexpr static size_t kSuperBlockSize = 63488;
  constexpr static size_t kBlocksPerSuperBlock = 128;
  constexpr static size_t kWordsPerBlock = 8;
 
  const size_t num_bits_;
  std::vector<uint64_t> bits_interleaved;
  std::vector<uint64_t> super_block_rank_;
 
  class BitReader {
    size_t iterator_64_ = 0;
    size_t offset_size_ = 0;
    size_t offset_bits_ = 0;
    std::span<const uint64_t> bits_;
 
   public:
    BitReader(std::span<const uint64_t> bits) : bits_(bits) {}
    uint64_t ReadBits64(size_t num_bits) {
      if (num_bits > 64) {
        num_bits = 64;
      }
      uint64_t result = offset_bits_ & first_bits_mask(num_bits);
      if (offset_size_ >= num_bits) {
        offset_bits_ >>= num_bits;
        offset_size_ -= num_bits;
        return result;
      }
      uint64_t next = iterator_64_ < bits_.size() ? bits_[iterator_64_++] : 0;
      result ^= (next & first_bits_mask(num_bits - offset_size_))
                << offset_size_;
      offset_bits_ = (num_bits - offset_size_ == 64)
                         ? 0
                         : next >> (num_bits - offset_size_);
      offset_size_ = 64 - (num_bits - offset_size_);
      return result;
    }
  };
 
 public:
  explicit BitVectorInterleaved(std::span<const uint64_t> bit_vector,
                                size_t num_bits)
      : num_bits_(std::min(num_bits, bit_vector.size() * kWordSize)) {
    build_rank_interleaved(bit_vector, num_bits);
  }

  static inline uint64_t first_bits_mask(size_t num) {
    return num >= 64 ? UINT64_MAX : ((1llu << num) - 1);
  }

  size_t size() const { return num_bits_; }

  int operator[](size_t pos) const {
    size_t block_id = pos / kBasicBlockSize;
    size_t block_bit = pos - block_id * kBasicBlockSize;
    size_t word_id = block_id * kWordsPerBlock + block_bit / kWordSize;
    size_t word_bit = block_bit % kWordSize;
    kWordSize;
 
    return (bits_interleaved[word_id] >> word_bit) & 1;
  }

  void build_rank_interleaved(std::span<const uint64_t> bits, size_t num_bits) {
    size_t num_superblocks = 1 + (num_bits_ - 1) / kSuperBlockSize;
    super_block_rank_.resize(num_superblocks);
    size_t num_basicblocks = 1 + (num_bits_ - 1) / kBasicBlockSize;
    bits_interleaved.resize(num_basicblocks * (512 / kWordSize));
 
    uint64_t super_block_sum = 0;
    uint16_t basic_block_sum = 0;
    auto bit_reader = BitReader(bits);
 
    for (size_t i = 0; i * kBasicBlockSize < num_bits; ++i) {
      if (i % (kSuperBlockSize / kBasicBlockSize) == 0) {
        super_block_sum += basic_block_sum;
        super_block_rank_[i / (kSuperBlockSize / kBasicBlockSize)] =
            super_block_sum;
        basic_block_sum = 0;
      }
      bits_interleaved[i * (kWordsPerBlock) + 7] =
          static_cast<uint64_t>(basic_block_sum) << 48;
 
      for (size_t j = 0; j < 7 && kWordSize * (i + j) < num_bits; ++j) {
        bits_interleaved[i * (kWordsPerBlock) + j] =
            bit_reader.ReadBits64(std::min<uint64_t>(
                64ull, num_bits - i * kBasicBlockSize + j * kWordSize));
        basic_block_sum +=
            std::popcount(bits_interleaved[i * (kWordsPerBlock) + j]);
      }
      if ((i + 7) * kWordSize < num_bits) {
        auto v = bit_reader.ReadBits64(std::min<uint64_t>(
            48ull, num_bits - (i * kBasicBlockSize + 7 * kWordSize)));
        bits_interleaved[i * (kWordsPerBlock) + 7] ^= v;
        basic_block_sum += std::popcount(v);
      }
    }
  }

  uint64_t rank(size_t pos) const {
    // Multiplication/devisions
    uint64_t b_block = pos / kBasicBlockSize;
    uint64_t s_block = b_block / kBlocksPerSuperBlock;
    uint64_t b_block_pos = b_block * kWordsPerBlock;
    // Super block rank
    uint64_t result = super_block_rank_[s_block];
    // __builtin_prefetch(&bits_interleaved[b_block_pos]);
    result += rank_512(&bits_interleaved[b_block_pos],
                       pos - (b_block * kBasicBlockSize));
    result += bits_interleaved[b_block_pos + 7] >> 48;
    return result;
  }

  std::string to_string() const {
    std::string result;
    result.reserve(num_bits_);
 
    for (size_t i = 0; i < num_bits_; i++) {
      result.push_back(operator[](i) ? '1' : '0');
    }
 
    return result;
  }
};
 
}  // namespace pixie