bintree.hpp

#pragma once
 
#include <array>
#include <cstdint>
#include <memory>
 
namespace itk::itkpix::endec::codec {
 
namespace detail {
template <bool USE_LUT = true>
class BinaryTreeLUT {
 public:
  BinaryTreeLUT() = default;
 
  ~BinaryTreeLUT() = default;
 
  static inline constexpr uint32_t LUT_LEN = 16;
 
  static inline constexpr uint32_t encode(uint32_t decoded, uint32_t &encoded) {
    if constexpr (USE_LUT) {
      uint64_t _encoded = lut_encode[decoded & 0xffff];
      encoded = _encoded & 0xffffffff;
      return _encoded >> 32;
    } else
      return encode_(decoded, encoded);
  }
 
  [[gnu::always_inline]] static inline constexpr uint32_t decode(uint32_t encoded, uint32_t &decoded) {
    if constexpr (USE_LUT) {
      auto _encoded = encoded & 0x3fffffff;
      auto [_decoded, _offset, _length] = lut_top[_encoded >> (30 - LUT_LEN)];
      uint32_t hitmap = _decoded;
      decoded = hitmap;
      if (_offset == 0xff) return _length;
      _encoded = (_encoded >> _offset) & 0x3fff;
      auto [_hitmap, length2] = lut_bottom[_encoded];
      decoded &= 0xff;
      decoded |= uint16_t(_hitmap) << 8;
      return _length + length2;
    } else
      return decode_(encoded, decoded);
  }
 private:
  using HitLutEntry = std::pair<uint8_t,uint8_t>;
  using HitLutRow = std::array<HitLutEntry, 16>;
  using HitLutType = std::array<HitLutRow, 1<<16>;
  static inline constexpr auto create_hit_lut_length_() {
    std::array<uint8_t, 1<<16> result{};
    for(uint32_t i = 0; i < (1<<16); i++)
      result[i] = std::popcount(i);
    return result;
  }
  static inline constexpr auto create_hit_lut_() {
    HitLutType result{};
    for (uint32_t hitmap = 0; hitmap < (1 << 16); ++hitmap) {
      int index{};
      for (uint8_t bit = 0; bit < 16; ++bit) {
        if((hitmap & (1<< bit)) == 0) continue;
        result[hitmap][index++] = {bit % 8, bit / 8};
      }
    }
    return result;
  }
 public:
  static inline HitLutType hit_lut {create_hit_lut_()};
  static inline std::array<uint8_t, (1<<16)> hit_lut_length{create_hit_lut_length_()};
 private:
  using LUT1_entry = std::tuple<uint16_t, uint8_t, uint8_t>;
  using LUT2_entry = std::tuple<uint8_t, uint8_t>;
  using LUT1_type = std::conditional_t<USE_LUT, std::array<LUT1_entry, 1 << LUT_LEN>, void *>;
  using LUT2_type = std::conditional_t<USE_LUT, std::array<LUT2_entry, 1 << (30 - LUT_LEN)>, void *>;
  using LUT_ENCODE_type = std::conditional_t<USE_LUT, std::array<uint64_t, (1 << 16)>, void *>;
 
  static inline constexpr auto create_lut_top_() {
    LUT1_type lut;
    for (uint32_t i = 0; i < (1 << LUT_LEN); i++) {
      uint32_t decoded{};
      auto length = decode_(i << (30 - LUT_LEN), decoded);
      uint32_t offset = 0xff;
      if (length > 16) {
        uint32_t shift = decoded & 0xff;
        uint32_t encoded;
        length = encode_(shift, encoded);
        offset = 16 - length;
      }
      lut[i] = {decoded & 0xffff, offset & 0xff, length & 0xff};
    }
    return lut;
  };
 
 public:
  static inline LUT1_type lut_top{create_lut_top_()};
 private:
  static inline constexpr auto create_lut_bottom_() {
    LUT2_type lut;
    for (uint32_t i = 0; i < 1 << (30 - LUT_LEN); i++) {
      uint32_t decoded{};
      auto length = decode_bottom_(i << LUT_LEN, decoded);
      lut[i] = {decoded >> 8, length};
    }
    return lut;
  }
 public:
  static inline LUT2_type lut_bottom{create_lut_bottom_()};
 private:
  static inline constexpr auto create_lut_encode_() {
    LUT_ENCODE_type lut;
    lut[0] = 0;
    for (uint32_t i = 1; i < 1 << 16; i++) {
      uint32_t encoded;
      uint64_t len = encode_(i, encoded);
      lut[i] = encoded | (len << 32);
    }
    return lut;
  }
 
  static inline LUT_ENCODE_type lut_encode{create_lut_encode_()};
 
  static inline constexpr uint32_t encode_(uint32_t decoded, uint32_t &encoded) {
    uint32_t b[8];
    for (int i = 0; i < 8; i++) {
      b[i] = ((decoded >> (2 * i)) & 0x1) << 1 | ((decoded >> (2 * i + 1)) & 0x1);
    }
    auto one_bit = [](uint32_t value) -> uint32_t { return value != 0 ? 1 : 0; };
    uint32_t S1 = (one_bit(b[0] | b[1] | b[2] | b[3]) << 1) | one_bit(b[4] | b[5] | b[6] | b[7]);
    uint32_t S2t = (one_bit(b[0] | b[1]) << 1) | one_bit(b[2] | b[3]);
    uint32_t S2b = (one_bit(b[4] | b[5]) << 1) | one_bit(b[6] | b[7]);
    uint32_t S3tl = (one_bit(b[0]) << 1) | one_bit(b[1]);
    uint32_t S3tr = (one_bit(b[2]) << 1) | one_bit(b[3]);
    uint32_t S3bl = (one_bit(b[4]) << 1) | one_bit(b[5]);
    uint32_t S3br = (one_bit(b[6]) << 1) | one_bit(b[7]);
 
    uint32_t pos = 0;
    encoded = 0;
 
    auto writeTwo = [&](uint32_t src) {
      if (src == 0b01) {
        encoded |= (0b0) << (28 - pos);
        pos++;
      } else {
        encoded |= (src & 0x3) << (28 - pos);
        pos += 2;
      }
    };
    for (auto val : {S1, S2t, S3tl, S3tr, b[0], b[1], b[2], b[3], S2b, S3bl, S3br, b[4], b[5], b[6], b[7]})
      if (val) writeTwo(val);
    return pos;
  }
 
  [[gnu::always_inline]] static constexpr uint32_t decode_(uint32_t encoded, uint32_t &decoded) {
    return decode_impl_<true>(encoded, decoded);
  }
 
  template <bool DECODE_ALL = true>
  [[gnu::always_inline]] static constexpr uint32_t decode_impl_(uint32_t encoded, uint32_t &decoded) {
    uint32_t pos{};
    uint32_t S1{};
    uint32_t S2t{};
    uint32_t S3tl{};
    uint32_t S3tr{};
    uint32_t S2b{};
    uint32_t S3bl{};
    uint32_t S3br{};
    uint32_t b[8]{};
 
    auto readTwo = [&](uint32_t &dst) {
      uint32_t val = (encoded >> (28 - pos)) & 0x3;
      if (val == 0b00 or val == 0b01) {
        dst = 0b01;
        pos++;
      } else {
        dst = val;
        pos += 2;
      }
    };
 
    auto read = [&](uint32_t branch, uint32_t index) {
      auto &left = b[index];
      auto &right = b[index + 1];
      switch (branch & 3) {
        case 0b10:
          readTwo(left);
          break;
        case 0b01:
          readTwo(right);
          break;
        case 0b11:
          readTwo(left);
          readTwo(right);
          break;
      }
    };
 
    auto read_bottom = [&]() {
      readTwo(S2b);
      switch (S2b & 3) {
        case 0b10:
          readTwo(S3bl);
          read(S3bl, 4);
          break;
        case 0b01:
          readTwo(S3br);
          read(S3br, 6);
          break;
        case 0b11:
          readTwo(S3bl);
          readTwo(S3br);
          read(S3bl, 4);
          read(S3br, 6);
          break;
      }
    };
 
    auto read_all = [&] {
      readTwo(S1);
      switch (S1 & 3) {
        case 0b10:
          readTwo(S2t);
          switch (S2t & 3) {
            case 0b10:
              readTwo(S3tl);
              read(S3tl, 0);
              break;
            case 0b01:
              readTwo(S3tr);
              read(S3tr, 2);
              break;
            case 0b11:
              readTwo(S3tl);
              readTwo(S3tr);
              read(S3tl, 0);
              read(S3tr, 2);
              break;
          }
          break;
        case 0b01:
          read_bottom();
          break;
        case 0b11:
          readTwo(S2t);
          switch (S2t & 3) {
            case 0b10:
              readTwo(S3tl);
              read(S3tl, 0);
              break;
            case 0b01:
              readTwo(S3tr);
              read(S3tr, 2);
              break;
            case 0b11:
              readTwo(S3tl);
              readTwo(S3tr);
              read(S3tl, 0);
              read(S3tr, 2);
              break;
          }
          read_bottom();
      }
    };
 
    if constexpr (DECODE_ALL)
      read_all();
    else
      read_bottom();
 
    decoded = 0;
    for (int i = 0; i < 8; i++) {
      decoded |= ((b[i] >> 1) & 0x1) << (i * 2);
      decoded |= ((b[i] >> 0) & 0x1) << (i * 2 + 1);
    }
    return pos;
  }
 
  static inline constexpr uint32_t decode_bottom_(uint32_t encoded, uint32_t &decoded) {
    return decode_impl_<false>(encoded, decoded);
  }
};
}  // namespace detail
using BinaryTreeLUT = detail::BinaryTreeLUT<true>;
using BinaryTreeNoLUT = detail::BinaryTreeLUT<false>;
 
}  // itk::itkpix::endec::codec