index.h

#pragma once
 
#include <pixie/rank_select/support.h>
#include <pixie/wavelet_tree.h>
 
#include <cstdlib>
#include <cstring>
#include <functional>
#include <limits>
#include <numeric>
#include <queue>
#include <span>
#include <vector>
 
namespace pixie {
 
template <StorageImplementation Storage = AlignedStorage>
class WaveletTreeIndex : public WaveletTreeBase<WaveletTreeIndex<Storage>> {
 private:
  using node_index_t = size_t;
  static constexpr node_index_t npos = std::numeric_limits<node_index_t>::max();
 
  struct PreWaveletNode {
    node_index_t parent = npos;
    node_index_t left_child = npos;
    node_index_t right_child = npos;
    uint64_t middle;
    OutputBitStream stream;
    explicit PreWaveletNode(uint64_t middle) : middle(middle) {}
  };

  struct WaveletNode {
    node_index_t parent, left_child, right_child;
    uint64_t middle;
    Storage bit_vector_data;
    RankSelectSupport<Storage> data;

    static AlignedStorage align(std::vector<uint64_t>&& data) {
      AlignedStorage result(data.size() * 64);
      auto view = result.writable_words64();
      std::copy(data.begin(), data.end(), view.begin());
      return std::move(result);
    }
 
    WaveletNode() = default;
 
    WaveletNode(PreWaveletNode&& node)
      requires(std::same_as<Storage, AlignedStorage>)
        : parent(node.parent),
          left_child(node.left_child),
          right_child(node.right_child),
          middle(node.middle),
          bit_vector_data(std::move(align(node.stream.extract()))),
          data(bit_vector_data.as_words64(), node.stream.size()) {}

    void serialize(pixie::OutputBitStream& bs) const {
      bs << parent << left_child << right_child << middle;
      bit_vector_data.serialize(bs);
      data.serialize(bs);
    }

    static WaveletNode deserialize(std::span<const std::byte>& data)
      requires(std::same_as<Storage, ReadOnlyStorageView>)
    {
      WaveletNode result;
      auto read = [&data](auto& value) {
        constexpr size_t length = sizeof(value);
        std::memcpy(&value, data.data(), length);
        data = data.subspan(length);
      };
      read(result.parent);
      read(result.left_child);
      read(result.right_child);
      read(result.middle);
      result.bit_vector_data = ReadOnlyStorageView::deserialize(data);
      result.data = RankSelectSupport<Storage>::deserialize(
          result.bit_vector_data.as_words64(), data);
      return result;
    }
  };
 
  size_t alphabet_size_, data_size_;
  node_index_t root_;
  std::vector<WaveletNode> nodes_;
  std::vector<node_index_t> leaves_;
  std::vector<size_t> permutation_, inverse_permutation_;

  template <typename F>
  node_index_t build_node(size_t begin,
                          size_t end,
                          node_index_t parent,
                          const F& get_middle,
                          std::span<const size_t> prefix_sum,
                          std::vector<PreWaveletNode>& nodes)
    requires(std::same_as<Storage, AlignedStorage>)
  {
    if (end - begin == 1) {
      leaves_[begin] = parent;
      return npos;
    }
    if (prefix_sum[end] == prefix_sum[begin]) {
      for (size_t symbol = begin; symbol < end; symbol++) {
        leaves_[symbol] = parent;
      }
      return npos;
    }
 
    node_index_t result = nodes.size();
    size_t middle = get_middle(result);
    middle = begin + (middle == npos ? (end - begin) / 2 : middle);
 
    nodes.emplace_back(middle);
    nodes[result].stream.reserve(prefix_sum[end] - prefix_sum[begin]);
    nodes[result].parent = parent;
    nodes[result].left_child =
        build_node(begin, middle, result, get_middle, prefix_sum, nodes);
    nodes[result].right_child =
        build_node(middle, end, result, get_middle, prefix_sum, nodes);
 
    return result;
  }

  void copy_segment_content(node_index_t node,
                            size_t begin,
                            size_t end,
                            std::span<uint64_t> dst,
                            std::span<uint64_t> tmp) const {
    if (begin == end) {
      return;
    }
    const size_t rank = nodes_[node].data.rank(begin), rank0 = begin - rank;
    const size_t right = nodes_[node].data.rank(end) - rank,
                 left = (end - begin) - right;
 
    if (nodes_[node].left_child == npos) {
      std::fill_n(tmp.begin(), static_cast<long long>(left),
                  inverse_permutation_[nodes_[node].middle - 1]);
    } else {
      copy_segment_content(nodes_[node].left_child, rank0, rank0 + left,
                           tmp.subspan(0, left), dst.subspan(0, left));
    }
    if (nodes_[node].right_child == npos) {
      std::fill(tmp.begin() + static_cast<long long>(left), tmp.end(),
                inverse_permutation_[nodes_[node].middle]);
    } else {
      copy_segment_content(nodes_[node].right_child, rank, rank + right,
                           tmp.subspan(left, right), dst.subspan(left, right));
    }
 
    size_t j = 0, k = left;
    const auto& bit_vector = nodes_[node].bit_vector_data.as_words64();
    for (size_t i = begin; i < end; i++) {
      if ((bit_vector[i / 64] >> (i % 64)) & 1) {
        dst[i - begin] = tmp[k++];
      } else {
        dst[i - begin] = tmp[j++];
      }
    }
  }
 
  WaveletTreeIndex() = default;
 
 public:
  WaveletTreeIndex(
      size_t alphabet_size,
      std::span<const uint64_t> data,
      const WaveletTreeBuildType build_type = WaveletTreeBuildType::Standard)
    requires(std::same_as<Storage, AlignedStorage>)
      : alphabet_size_(alphabet_size),
        data_size_(data.size()),
        leaves_(alphabet_size_, npos) {
    if (alphabet_size == 0) {
      root_ = npos;
      return;
    }
    std::vector<PreWaveletNode> nodes;
    nodes.reserve(alphabet_size_);
    std::vector<size_t> nodes_structure;
    nodes_structure.reserve(alphabet_size_);
 
    if (build_type == WaveletTreeBuildType::Standard) {
      permutation_.resize(alphabet_size);
      inverse_permutation_.resize(alphabet_size);
      std::iota(permutation_.begin(), permutation_.end(), 0);
      std::iota(inverse_permutation_.begin(), inverse_permutation_.end(), 0);
      nodes_structure.resize(alphabet_size_, npos);
    } else {
      struct Node {
        size_t size, left, right;
      };
      std::vector<Node> huffman_nodes(alphabet_size_, {0, 0, 0});
      for (auto symb : data) {
        huffman_nodes[symb].size++;
      }
 
      using elem_t = std::pair<size_t, size_t>;
      std::priority_queue<elem_t, std::vector<elem_t>, std::greater<>> queue;
      for (size_t i = 0; i < alphabet_size_; i++) {
        queue.emplace(huffman_nodes[i].size, i);
      }
      while (queue.size() >= 2) {
        auto right = queue.top().second;
        queue.pop();
        auto left = queue.top().second;
        queue.pop();
        huffman_nodes.push_back(
            {huffman_nodes[left].size + huffman_nodes[right].size, left + 1,
             right + 1});
        queue.emplace(huffman_nodes.back().size, huffman_nodes.size() - 1);
      }
 
      std::function<size_t(size_t)> enumerate = [&](size_t index) -> size_t {
        const auto& [size, left, right] = huffman_nodes[index];
        if (left == 0 || right == 0) {
          permutation_[index] = inverse_permutation_.size();
          inverse_permutation_.push_back(index);
          return 1;
        }
        size_t ind = nodes_structure.size(), subtree = 0;
        if (size > 0) {
          nodes_structure.push_back(0);
        }
        subtree += enumerate(left - 1);
        if (size > 0) {
          nodes_structure[ind] = subtree;
        }
        subtree += enumerate(right - 1);
        return subtree;
      };
 
      permutation_.resize(alphabet_size_);
      inverse_permutation_.reserve(alphabet_size_);
      enumerate(huffman_nodes.size() - 1);
    }
 
    std::vector<size_t> prefix_sum(alphabet_size + 1);
    for (auto symbol : data) {
      prefix_sum[permutation_[symbol] + 1]++;
    }
    for (size_t i = 0; i < alphabet_size_; i++) {
      prefix_sum[i + 1] += prefix_sum[i];
    }
 
    root_ = build_node(
        0, alphabet_size_, npos,
        [&](node_index_t node) { return nodes_structure[node]; }, prefix_sum,
        nodes);
    for (auto symbol : data) {
      auto index = permutation_[symbol];
      for (node_index_t current = root_; current != npos;) {
        auto& node = nodes[current];
        bool go_right = index >= node.middle;
        node.stream << go_right;
        if (go_right) {
          current = node.right_child;
        } else {
          current = node.left_child;
        }
      }
    }
    nodes_.reserve(nodes.size());
    for (auto& node : nodes) {
      nodes_.emplace_back(std::move(node));
    }
  }

  size_t rank_impl(uint64_t symbol, size_t pos) const {
    if (symbol >= alphabet_size_) [[unlikely]] {
      return 0;
    }
    symbol = permutation_[symbol];
    for (node_index_t current = root_; current != npos;) {
      const WaveletNode& node = nodes_[current];
      if (symbol < node.middle) {
        pos = node.data.rank0(pos);
        current = node.left_child;
      } else {
        pos = node.data.rank(pos);
        current = node.right_child;
      }
    }
    return pos;
  }

  size_t select_impl(uint64_t symbol, size_t rank) const {
    if (symbol >= alphabet_size_ || data_size_ == 0) [[unlikely]] {
      return data_size_;
    }
    symbol = permutation_[symbol];
    node_index_t current = leaves_[symbol];
    for (; current != npos; current = nodes_[current].parent) {
      const WaveletNode& node = nodes_[current];
      if (symbol < node.middle) {
        rank = node.data.select0(rank) + 1;
      } else {
        rank = node.data.select(rank) + 1;
      }
    }
    return rank - 1;
  }

  std::vector<uint64_t> get_segment_impl(size_t begin, size_t end) const {
    if (alphabet_size_ == 0 || data_size_ == 0 || begin >= end) [[unlikely]] {
      return {};
    }
    auto length = static_cast<long long>(end - begin);
    std::vector<uint64_t> result(2 * length);
    copy_segment_content(root_, begin, end,
                         std::span{result.begin(), result.begin() + length},
                         std::span{result.begin() + length, result.end()});
    result.resize(length);
    return result;
  }

  size_t size_impl() const { return data_size_; }

  void serialize(pixie::OutputBitStream& bs) const {
    bs << alphabet_size_ << data_size_ << root_ << nodes_.size();
    for (const WaveletNode& node : nodes_) {
      node.serialize(bs);
    }
    for (const node_index_t leaf : leaves_) {
      bs << leaf;
    }
    for (const size_t idx : permutation_) {
      bs << idx;
    }
  }
 
  static WaveletTreeIndex<ReadOnlyStorageView> deserialize(
      std::span<const std::byte>& data) {
    WaveletTreeIndex<ReadOnlyStorageView> result;
    auto read = [&data](auto& value) {
      constexpr size_t length = sizeof(value);
      std::memcpy(&value, data.data(), length);
      data = data.subspan(length);
    };
    read(result.alphabet_size_);
    read(result.data_size_);
    read(result.root_);
    size_t size;
    read(size);
    result.nodes_.resize(size);
    for (auto& node : result.nodes_) {
      node = WaveletNode::deserialize(data);
    }
    result.leaves_.resize(result.alphabet_size_);
    for (node_index_t& leaf : result.leaves_) {
      read(leaf);
    }
    result.permutation_.resize(result.alphabet_size_);
    for (size_t& index : result.permutation_) {
      read(index);
    }
    result.inverse_permutation_.resize(result.alphabet_size_);
    for (size_t i = 0; i < result.alphabet_size_; i++) {
      result.inverse_permutation_[result.permutation_[i]] = i;
    }
    return result;
  }
};
 
using WaveletTree = WaveletTreeIndex<AlignedStorage>;
using WaveletTreeView = WaveletTreeIndex<ReadOnlyStorageView>;
 
}  // namespace pixie