bit_stream.h

#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) {}

  OutputBitStream& operator<<(bool bit) {
    if (size_ % 64 == 0) {
      data_.push_back(bit);
    } else if (bit) {
      data_.back() |= 1ull << (size_ % 64);
    }
    size_++;
    return *this;
  }

  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;
  }

  size_t size() const { return size_; }

  void reserve(size_t size) { data_.reserve((size + 63) / 64); }

  std::vector<uint64_t> extract() { return std::move(data_); }
};
 
}  // namespace pixie