pasta.h

#pragma once
 
#ifndef PIXIE_PASTA_SUPPORT
#error "PastaRankSelectSupport requires PIXIE_THIRD_PARTY_BACKENDS"
#endif
 
#include <pixie/rank_select.h>
 
#include <algorithm>
#include <bit>
#include <cstddef>
#include <cstdint>
#include <pasta/bit_vector/bit_vector.hpp>
#include <pasta/bit_vector/support/flat_rank.hpp>
#include <pasta/bit_vector/support/flat_rank_select.hpp>
#include <pasta/bit_vector/support/l12_type.hpp>
#include <span>
 
namespace pixie {

class PastaRankSelectSupport : public RankSelectBase<PastaRankSelectSupport> {
 private:
  using PastaSupport = pasta::FlatRankSelect<>;
 
  static uint64_t logical_word(std::span<const uint64_t> source,
                               std::size_t word_index,
                               std::size_t num_bits) {
    uint64_t word = source[word_index];
    const std::size_t remaining = num_bits - word_index * 64;
    if (remaining < 64) {
      word &= (uint64_t{1} << remaining) - 1;
    }
    return word;
  }
 
  pasta::BitVector bits_;
  PastaSupport support_;
  std::size_t num_bits_ = 0;
  std::size_t max_rank_ = 0;
 
 public:
  explicit PastaRankSelectSupport(std::span<const uint64_t> source_words,
                                  std::size_t num_bits)
      : bits_(std::min(num_bits, source_words.size() * 64)),
        num_bits_(std::min(num_bits, source_words.size() * 64)) {
    std::fill(bits_.data().begin(), bits_.data().end(), 0);
    const std::size_t word_count = (num_bits_ + 63) / 64;
    for (std::size_t word_index = 0; word_index < word_count; ++word_index) {
      bits_.data()[word_index] =
          logical_word(source_words, word_index, num_bits_);
    }
    support_ = PastaSupport(bits_);
    max_rank_ = num_bits_ == 0 ? 0 : support_.rank1(num_bits_);
  }

  std::size_t size_impl() const { return num_bits_; }

  int bit_impl(std::size_t position) const {
    return static_cast<bool>(bits_[position]);
  }

  uint64_t rank_impl(std::size_t end_position) const {
    return end_position >= num_bits_ ? max_rank_ : support_.rank1(end_position);
  }

  uint64_t select_impl(std::size_t rank) const {
    if (rank == 0) {
      return 0;
    }
    return rank > max_rank_ ? num_bits_ : support_.select1(rank);
  }

  uint64_t select0_impl(std::size_t rank) const {
    if (rank == 0) {
      return 0;
    }
    return rank > num_bits_ - max_rank_ ? num_bits_ : support_.select0(rank);
  }

  bool supports_select1_impl() const { return true; }

  bool supports_select0_impl() const { return true; }

  std::size_t memory_usage_bytes_impl() const {
    return source_copy_bytes() + index_logical_bytes() +
           (sizeof(*this) - sizeof(bits_) - sizeof(support_));
  }

  std::size_t source_copy_bytes() const {
    return bits_.space_usage() - sizeof(pasta::BitVector);
  }

  std::size_t index_logical_bytes() const {
    const std::size_t l12_entries =
        bits_.data().size() / pasta::FlatRankSelectConfig::L1_WORD_SIZE + 1;
    return (support_.space_usage() - sizeof(PastaSupport)) +
           l12_entries * sizeof(pasta::BigL12Type);
  }
};
 
}  // namespace pixie