bit_dict.hpp

#pragma once
 
#include <climits>
#include <cstdint>
#include <memory>
#include <vector>
 
#include "cplib/port/bit.hpp"
#include "cplib/utils/bit.hpp"
 
namespace cplib {
 
namespace impl {
 
class alignas(64) BitDictBlock {
 public:
  using size_type = std::size_t;
  using word_type = std::uint64_t;
  static constexpr size_type words_per_block = 6;
  static constexpr size_type word_size = std::numeric_limits<word_type>::digits;
  static constexpr size_type block_size = word_size * words_per_block;
 
  BitDictBlock() : data_({0}) {}
 
  size_type build(size_type block_rank) {
    block_rank_ = block_rank;
    wrap_word_ = words_per_block;
    size_type acc_rank = 0;
    for (int i = 0; i < words_per_block; i++) {
      if (acc_rank >= 256 && wrap_word_ == words_per_block) {
        wrap_word_ = i;
      }
      word_rank_[i] = acc_rank;
      acc_rank += port::popcount(data_[i]);
    }
    return block_rank_ + acc_rank;
  }
 
  bool get(size_type idx) const { return (data_[idx / word_size] >> (idx % word_size)) & 1; }
 
  void set(size_type idx) { data_[idx / word_size] |= word_type(1) << (idx % word_size); }
 
  size_type rank(size_type idx) const {
    size_type word_idx = idx / word_size;
    return block_rank_ + get_word_rank_(word_idx) + impl::popcount_low(data_[word_idx], idx % word_size);
  }
 
 private:
  size_type get_word_rank_(size_type idx) const { return word_rank_[idx] + (idx >= wrap_word_) * 256; }
 
  std::array<word_type, words_per_block> data_;
  size_type block_rank_;
  std::array<std::uint8_t, words_per_block> word_rank_;
  std::uint8_t wrap_word_;
};
 
static_assert(sizeof(BitDictBlock) == 64);
 
}  // namespace impl
 
class BitDict {
 public:
  using size_type = impl::BitDictBlock::size_type;
 
  BitDict() : size_(0), zeros_(0) {}
 
  BitDict(size_type num_bits) : size_(num_bits), blocks_(num_bits / block_size + 1) {}
 
  size_type size() const { return size_; }
 
  size_type zeros() const { return zeros_; }
 
  size_type ones() const { return size_ - zeros_; }
 
  bool get(size_type idx) const { return blocks_[idx / block_size].get(idx % block_size); }
 
  void set(size_type idx) { blocks_[idx / block_size].set(idx % block_size); }
 
  template <typename BitGenerator>
  void fill_with_bit_generator(BitGenerator& gen) {
    for (size_type i = 0; i < size_; i++) {
      if (gen()) {
        set(i);
      }
    }
  }
 
  void build() {
    size_type acc_rank = 0;
    for (auto& block : blocks_) {
      acc_rank = block.build(acc_rank);
    }
    zeros_ = size_ - acc_rank;
  }
 
  size_type rank1(size_type idx) const { return blocks_[idx / block_size].rank(idx % block_size); }
 
  size_type rank0(size_type idx) const { return idx - rank1(idx); }
 
  size_type rank_to_child(size_type idx, bool bit) const { return bit ? zeros() + rank1(idx) : rank0(idx); }
 
 private:
  static constexpr size_type block_size = impl::BitDictBlock::block_size;
  std::vector<impl::BitDictBlock> blocks_;
  size_type size_, zeros_;
};
 
}  // namespace cplib