Wavelet packing

CMakePresets.json

@@ -118,7 +118,8 @@
                 "louds_tree_tests",
                 "bp_tree_tests",
                 "dfuds_tree_tests",
-                "excess_positions_tests"
+                "excess_positions_tests",
+                "wavelet_tree_tests"
             ]
         },
         {
@@ -133,7 +134,8 @@
                 "louds_tree_tests",
                 "bp_tree_tests",
                 "dfuds_tree_tests",
-                "excess_positions_tests"
+                "excess_positions_tests",
+                "wavelet_tree_tests"
             ]
         },
         {
@@ -148,7 +150,8 @@
                 "bp_tree_benchmarks",
                 "dfuds_tree_benchmarks",
                 "alignment_comparison",
-                "excess_positions_benchmarks"
+                "excess_positions_benchmarks",
+                "wavelet_tree_benchmarks"
             ]
         },
         {
@@ -164,7 +167,8 @@
                 "bp_tree_benchmarks",
                 "dfuds_tree_benchmarks",
                 "alignment_comparison",
-                "excess_positions_benchmarks"
+                "excess_positions_benchmarks",
+                "wavelet_tree_benchmarks"
             ]
         },
         {
@@ -179,7 +183,8 @@
                 "bp_tree_benchmarks",
                 "dfuds_tree_benchmarks",
                 "alignment_comparison",
-                "excess_positions_benchmarks"
+                "excess_positions_benchmarks",
+                "wavelet_tree_benchmarks"
             ]
         },
         {
@@ -202,7 +207,8 @@
                 "louds_tree_tests",
                 "bp_tree_tests",
                 "dfuds_tree_tests",
-                "excess_positions_tests"
+                "excess_positions_tests",
+                "wavelet_tree_tests"
             ]
         },
         {
@@ -217,7 +223,8 @@
                 "louds_tree_tests",
                 "bp_tree_tests",
                 "dfuds_tree_tests",
-                "excess_positions_tests"
+                "excess_positions_tests",
+                "wavelet_tree_tests"
             ]
         }
     ]

include/pixie/bit_stream.h

@@ -0,0 +1,72 @@
+#pragma once
+
+#include <cstdint>
+#include <cstdlib>
+#include <vector>
+
+namespace pixie {
+
+class OutputBitStream {
+ private:
+  size_t size_;
+  std::vector<uint64_t> data_;
+
+ public:
+  OutputBitStream() : size_(0) {}
+
+  /**
+   * @brief Writes one bit to the stream
+   */
+  OutputBitStream& operator<<(bool bit) {
+    if (size_ % 64 == 0) {
+      data_.push_back(bit);
+    } else if (bit) {
+      data_.back() |= 1ull << (size_ % 64);
+    }
+    size_++;
+    return *this;
+  }
+
+  /**
+   * @brief Writes bits of the integral number to the stream in little-endian
+   */
+  template <std::integral T>
+  OutputBitStream& operator<<(T bits) {
+    using UT = std::make_unsigned_t<T>;
+    UT ubits = static_cast<UT>(bits);
+    constexpr size_t length = sizeof(T) * 8;
+    static_assert(length <= 64);
+    if (size_ % 64 == 0) {
+      data_.push_back(ubits);
+    } else {
+      const size_t prefix = std::min(length, 64 - (size_ % 64));
+      data_.back() |= static_cast<uint64_t>(ubits & ((1ull << prefix) - 1))
+                      << (size_ % 64);
+      if (prefix < length) {
+        data_.push_back(ubits >> prefix);
+      }
+    }
+    size_ += length;
+    return *this;
+  }
+
+  /**
+   * @brief Returns the number of written bits
+   *
+   */
+  size_t size() const { return size_; }
+
+  /**
+   * @brief Reserves memory for "size" bits
+   *
+   */
+  void reserve(size_t size) { data_.reserve((size + 63) / 64); }
+
+  /**
+   * @brief Moves vector containing written bits. There must be no operator<<
+   * after extract()
+   */
+  std::vector<uint64_t> extract() { return std::move(data_); }
+};
+
+}  // namespace pixie

include/pixie/bitvector.h

@@ -1,7 +1,9 @@
 #pragma once
 
+#include <pixie/bit_stream.h>
 #include <pixie/bits.h>
 #include <pixie/cache_line.h>
+#include <pixie/mmap_storage.h>
 
 #include <algorithm>
 #include <bit>
@@ -64,7 +66,8 @@ namespace pixie {
  * not interleave data and index, favoring simpler scans. Rank metadata and
  * enabled select samples are built in one scan over the source words.
  */
-class BitVector {
+template <typename Storage = AlignedStorage>
+class BitVectorBase {
  public:
   /**
    * @brief Select directions to index during construction.
@@ -86,12 +89,12 @@ class BitVector {
   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];
+  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 select_samples_;    // 64-bit global positions
+  Storage super_block_rank_;  // 64-bit global prefix sums
+  Storage basic_block_rank_;  // 16-bit local prefix sums
+  Storage select_samples_;    // 64-bit global positions
   size_t num_bits_{};
   size_t padded_size_{};
   size_t max_rank_{};
@@ -628,11 +631,11 @@ class BitVector {
   }
 
  public:
-  BitVector() = default;
-  BitVector(const BitVector&) = default;
-  BitVector(BitVector&&) noexcept = default;
-  BitVector& operator=(const BitVector&) = default;
-  BitVector& operator=(BitVector&&) noexcept = default;
+  BitVectorBase() = default;
+  BitVectorBase(const BitVectorBase&) = default;
+  BitVectorBase(BitVectorBase&&) noexcept = default;
+  BitVectorBase& operator=(const BitVectorBase&) = default;
+  BitVectorBase& operator=(BitVectorBase&&) noexcept = default;
 
 #ifdef PIXIE_DIAGNOSTICS
   struct DiagnosticsBytes {
@@ -693,10 +696,10 @@ class BitVector {
    * @param one_count Optional exact number of 1-bits. When supplied,
    * construction can allocate select sample storage exactly.
    */
-  explicit BitVector(std::span<const uint64_t> bit_vector,
-                     size_t num_bits,
-                     SelectSupport select_support = SelectSupport::kBoth,
-                     std::optional<size_t> one_count = std::nullopt)
+  explicit BitVectorBase(std::span<const uint64_t> bit_vector,
+                         size_t num_bits,
+                         SelectSupport select_support = SelectSupport::kBoth,
+                         std::optional<size_t> one_count = std::nullopt)
       : num_bits_(std::min(num_bits, bit_vector.size() * kWordSize)),
         padded_size_(((num_bits_ + kWordSize - 1) / kWordSize) * kWordSize),
         bits_(bit_vector) {
@@ -846,8 +849,60 @@ class BitVector {
 
     return result;
   }
+
+  void serialize(pixie::OutputBitStream& bs) const {
+    bs << num_bits_ << padded_size_ << max_rank_ << select1_sample_begin_
+       << select1_sample_count_ << select0_sample_begin_
+       << select0_sample_count_ << static_cast<uint32_t>(select_support_)
+       << static_cast<uint32_t>(select0_samples_reversed_);
+    for (const uint64_t delta : delta_super) {
+      bs << delta;
+    }
+    for (const uint16_t delta : delta_basic) {
+      bs << delta;
+    }
+    super_block_rank_.serialize(bs);
+    basic_block_rank_.serialize(bs);
+    select_samples_.serialize(bs);
+  }
+
+  static BitVectorBase<MmapViewStorage> deserialize(
+      std::span<const uint64_t> source_bits,
+      std::span<const std::byte>& data) {
+    BitVectorBase result;
+    result.bits_ = source_bits;
+    auto read = [&data](auto& value) {
+      constexpr size_t length = sizeof(value);
+      std::memcpy(&value, data.data(), length);
+      data = data.subspan(length);
+    };
+    read(result.num_bits_);
+    read(result.padded_size_);
+    read(result.max_rank_);
+    read(result.select1_sample_begin_);
+    read(result.select1_sample_count_);
+    read(result.select0_sample_begin_);
+    read(result.select0_sample_count_);
+    uint32_t buf;
+    read(buf);
+    result.select_support_ = static_cast<SelectSupport>(buf);
+    read(buf);
+    result.select0_samples_reversed_ = static_cast<bool>(buf);
+    for (uint64_t& delta : result.delta_super) {
+      read(delta);
+    }
+    for (uint16_t& delta : result.delta_basic) {
+      read(delta);
+    }
+    result.super_block_rank_ = MmapViewStorage::deserialize(data);
+    result.basic_block_rank_ = MmapViewStorage::deserialize(data);
+    result.select_samples_ = MmapViewStorage::deserialize(data);
+    return result;
+  }
 };
 
+typedef BitVectorBase<AlignedStorage> BitVector;
+
 /**
  * @brief Interleaved, owning bit vector with rank and select.
  *

include/pixie/cache_line.h

@@ -6,6 +6,8 @@
 #include <span>
 #include <vector>
 
+#include "bit_stream.h"
+
 namespace pixie {
 
 inline constexpr std::size_t kAlignedStorageLineBytes = 64;
@@ -137,6 +139,16 @@ class AlignedStorage {
         reinterpret_cast<const std::uint16_t*>(data_.data()),
         data_.size() * kAlignedStorageLineWords16);
   }
+
+  /** @brief Writes bit representation of data to OutputBitStream */
+  void serialize(pixie::OutputBitStream& bs) const {
+    bs << data_.size() * sizeof(CacheLine);
+    for (const CacheLine& line : data_) {
+      for (const std::byte& byte : line.data) {
+        bs << static_cast<uint8_t>(byte);
+      }
+    }
+  }
 };
 
 }  // namespace pixie

include/pixie/mmap_storage.h

@@ -0,0 +1,85 @@
+#pragma once
+
+#include <cstring>
+#include <filesystem>
+#include <mio/mmap.hpp>
+#include <span>
+#include <stdexcept>
+
+#include "bit_stream.h"
+
+namespace pixie {
+
+class MmapFile {
+ private:
+  mio::mmap_source mmap_;
+
+ public:
+  explicit MmapFile(const std::filesystem::path& path) : mmap_(path.string()) {
+    if (!mmap_.is_open()) {
+      throw std::runtime_error("Failed to mmap file");
+    }
+  }
+
+  MmapFile(const MmapFile&) = delete;
+  MmapFile& operator=(const MmapFile&) = delete;
+
+  MmapFile(MmapFile&&) noexcept = default;
+  MmapFile& operator=(MmapFile&&) noexcept = default;
+
+  std::span<const std::byte> AsBytes() const {
+    return {reinterpret_cast<const std::byte*>(mmap_.data()), mmap_.size()};
+  }
+
+  size_t Size() const noexcept { return mmap_.size(); }
+};
+
+class MmapViewStorage {
+ private:
+  std::span<const std::byte> data_;
+
+ public:
+  MmapViewStorage() = default;
+  explicit MmapViewStorage(std::span<const std::byte> data) : data_(data) {}
+
+  void resize(size_t) {
+    throw std::logic_error("Cannot resize read-only mmap storage");
+  }
+
+  std::span<const uint64_t> AsConst64BitInts() const {
+    return {reinterpret_cast<const uint64_t*>(data_.data()), data_.size() / 8};
+  }
+
+  std::span<const std::uint16_t> AsConst16BitInts() const {
+    return {reinterpret_cast<const uint16_t*>(data_.data()), data_.size() / 2};
+  }
+
+  std::span<const std::byte> AsConstBytes() const { return data_; }
+
+  std::span<uint64_t> As64BitInts() {
+    throw std::logic_error("Read-only storage");
+  }
+  std::span<std::uint16_t> As16BitInts() {
+    throw std::logic_error("Read-only storage");
+  }
+
+  /** @brief Writes bit representation of data to OutputBitStream */
+  void serialize(pixie::OutputBitStream& bs) const {
+    bs << data_.size();
+    for (const std::byte byte : data_) {
+      bs << static_cast<uint8_t>(byte);
+    }
+  }
+
+  static MmapViewStorage deserialize(std::span<const std::byte>& data) {
+    MmapViewStorage result;
+    constexpr size_t length = sizeof(size_t);
+    size_t size;
+    std::memcpy(&size, data.data(), length);
+    result.data_ = data.subspan(length, size);
+    data = data.subspan(length + size);
+    return result;
+  };
+};
+
+}  // namespace pixie