bitstream.hpp

#pragma once
 
#include <cassert>
#include <cstdint>
#include <type_traits>
#include <bitset>
 
namespace itk::itkpix::endec::codec {
 
template <typename T>
[[gnu::always_inline]] static inline constexpr T bit_extract(T value, std::size_t start, std::size_t len)
  requires(std::is_unsigned_v<T>)
{
  T result;
  result = (value >> start) & ((T(1) << len) - 1);
  return result;
}
 
template <typename T>
[[gnu::always_inline]] static inline constexpr void bit_insert(T value, T &dst, std::size_t start,
                                                               std::size_t len)
  requires(std::is_unsigned_v<T>)
{
  T extracted_bits = bit_extract(value, 0, len);
  T mask = ((T(1) << len) - 1) << start;
  dst &= ~(mask);
  dst |= (extracted_bits << start);
}
/*
template <class T>
concept UnsignedArray = (sizeof(typename T::value_type) == 8) &&
                        std::unsigned_integral<typename T::value_type> && requires(T a, std::size_t b) {
                          { a.size() } -> std::same_as<std::size_t>;
                          { a[b] };
                        };
 
template <class T>
concept UnsignedVector =
    (sizeof(typename T::value_type) == 8 || sizeof(typename T::value_type) == 4) &&
    std::unsigned_integral<typename T::value_type> && requires(T a, T::value_type b, std::size_t c) {
      { a.size() } -> std::same_as<std::size_t>;
      { a.push_back(b) } -> std::same_as<void>;
      { a[c] };
    };
*/
 
template <typename Buffer>
  // requires UnsignedVector<Buffer> || UnsignedArray<Buffer>
struct BitStreamReader {
 
  BitStreamReader(const BitStreamReader &other)
      : data(other.data),
        size(other.size),
        Offset(other.Offset),
        _index(Offset),
        nBitsInWord(other.nBitsInWord) {}
 
  using value_type = uint64_t;
  using result_type = std::pair<bool, value_type>;
//  using buffer_type = Buffer;
  static constexpr std::size_t nBits = 64;
  explicit BitStreamReader(const std::span<const uint64_t> data, std::size_t offset)
      : data(data), size(data.size() * nBits), Offset(offset) {}
 
  [[gnu::always_inline]] [[nodiscard]] uint64_t first_word() const { return data[0]; }
 
  [[gnu::always_inline]] [[nodiscard]] std::size_t left() const {
    return (_index > size) ? 0 : size - _index;
  }
 
  [[gnu::always_inline]] [[nodiscard]] std::size_t index() const { return _index; }
 
  [[gnu::always_inline]] bool set(std::size_t N) {
    if (N >= size) return false;
    _index = N + Offset;
    return true;
  }
 
  [[gnu::always_inline]] inline bool advance(std::size_t const N) {
    if (_index + N > size) return false;
    auto i_bit = _index % nBits;
    auto left = nBits - i_bit;
 
    if (N <= left) {
      _index += N;
      if (N == left) _index += Offset;
      return true;
    }
    auto w2 = N - left;
    if (w2 <= nBitsInWord) {
      _index += (N + Offset);
      if (w2 == nBitsInWord) _index += Offset;
      return true;
    }
    _index += (2 * Offset + N);
    return true;
  }
 
  template <bool Advance>
  [[gnu::always_inline]] inline result_type _fetch(std::size_t const N) {
    value_type value{};
    if (_index + N > size) return {false, value};
    auto i_word = _index / nBits;
    auto i_bit = _index % nBits;
    auto left = nBits - i_bit;
    if (N <= left) {
      value = bit_extract(data[i_word], left - N, N);
      if constexpr (Advance) {
        _index += N;
        if (N == left) _index += Offset;
      }
      return {true, value};
    }
    auto w2 = N - left;
    if (w2 <= nBitsInWord) {
      auto word_1 = bit_extract(data[i_word], 0, left);
      auto word_2 = bit_extract(data[i_word + 1], nBitsInWord - w2, w2);
      value = word_2 | (word_1 << w2);
      if constexpr (Advance) {
        _index += (N + Offset);
        if (w2 == nBitsInWord) _index += Offset;
      }
      return {true, value};
    }
    auto w3 = (N - left) - nBitsInWord;
    auto word_1 = bit_extract(data[i_word], 0, left);
    auto word_2 = bit_extract(data[i_word + 1], 0, nBitsInWord);
    auto word_3 = bit_extract(data[i_word + 2], nBitsInWord - w3, w3);
    value = (word_1 << i_bit) | (word_2 << (i_bit - nBitsInWord)) | word_3;
    if constexpr (Advance) {
      _index += (2 * Offset + N);
    }
    return {true, value};
  }
 
  [[gnu::always_inline]] inline result_type prefetch(const std::size_t sz) {
    return _fetch<false>(sz);
  }
 
  [[gnu::always_inline]] inline result_type fetch(const std::size_t sz) {
    return _fetch<true>(sz);
  }
 
  std::span<const uint64_t> data;
  const std::size_t size;
  const std::size_t Offset;
  std::size_t _index{Offset};
  const std::size_t nBitsInWord = nBits - Offset;
};
 
/*
 
template <typename Buffer>
  requires UnsignedArray<Buffer> || UnsignedVector<Buffer>
struct BitStreamWriter {
  BitStreamWriter(const BitStreamWriter &other)
      : data(other.data), _size(other._size), Offset(other.Offset), index(other.index) {}
 
  BitStreamWriter &operator=(const BitStreamWriter &other) {
    if (&other != this) {
      data = other.data;
      _size = other._size;
      Offset = other.Offset;
      index = other.index;
    }
    return *this;
  }
  using value_type = typename Buffer::value_type;
  using buffer_type = Buffer;
  static inline constexpr std::size_t nBits = sizeof(value_type) * 8;
  static inline constexpr bool has_push_back() { return UnsignedVector<Buffer>; }
  explicit BitStreamWriter(Buffer &data, const std::size_t offset)
      : data(data), _size((data.size()) * nBits), Offset(offset) {}
 
  [[gnu::always_inline]] [[nodiscard]] std::size_t left() const { return size() - index; }
 
  [[gnu::always_inline]] inline bool write(value_type value, std::size_t N) {
    if constexpr (has_push_back()) {
      if (index + N > size()) {
        data.push_back(0);
        _size += nBits;
      }
    } else {
      if (index + N > size()) return false;
    }
    auto i_word = index / nBits;
    auto i_bit = index % nBits;
    auto left = nBits - i_bit;
    if (N <= left) {
      itkpix::codec::bit_insert(value, data[i_word], left - N, N);
      index += N;
      if (N == left) index += Offset;
      return true;
    }
    auto w2 = N - left;
    if (w2 <= nBitsInWord) {
      auto word_1 = value >> w2;
      auto word_2 = value;
      itkpix::codec::bit_insert(word_1, data[i_word], 0, left);
      itkpix::codec::bit_insert(word_2, data[i_word + 1], nBitsInWord - w2, w2);
      index += (N + Offset);
      if (w2 == nBitsInWord) index += Offset;
      return true;
    }
    auto w3 = (N - left) - nBitsInWord;
    auto word_1 = value >> i_bit;
    auto word_2 = value >> (Offset + i_bit);
    auto word_3 = value;
    itkpix::codec::bit_insert(word_1, data[i_word], 0, left);
    itkpix::codec::bit_insert(word_2, data[i_word + 1], 0, nBitsInWord);
    itkpix::codec::bit_insert(word_3, data[i_word + 2], nBitsInWord - w3, w3);
    index += (2 * Offset + N);
    return true;
  }
 
  [[gnu::always_inline]] inline void setHeaders(uint32_t value) {
    value_type mask = (value_type(1) << nBitsInWord) - 1;
    for (int i = 0; i < index / nBits; i++) {
      data[i] &= mask;
      data[i] |= value_type(value) << nBitsInWord;
      data[0] |= (uint64_t(1)<<63);
    }
  }
  std::size_t wordSize() const {
    auto i_word = index / nBits;
    auto i_bit = index % nBits;
    auto offset = (i_bit) > 0 ? 1 : 0;
    return i_word + offset;
  }
 
  std::size_t constexpr size() { return _size; }
 
  Buffer &data;
  std::size_t Offset;
  std::size_t _size;
  std::size_t index{Offset};
  std::size_t nBitsInWord = (nBits - Offset);
};
*/
 
}  // itk::itkpix::endec::codec