bit_span: add pop and consume

will be useful when decoding a DEFLATE stream

Change-Id: I4ba8659230270f8e571fe70444d9f5c629a888a3

huffman/src/bit_span.hpp

@@ -2,11 +2,13 @@
 #include "huffman/src/bit.hpp"
 #include "huffman/src/detail/iterator_interface.hpp"
 
+#include <bit>
 #include <bitset>
 #include <cassert>
 #include <climits>
 #include <cstddef>
 #include <cstdint>
+#include <cstring>
 #include <iterator>
 #include <limits>
 #include <ranges>
@@ -17,7 +19,8 @@ class bit_span : public std::ranges::view_interface<bit_span>
 {
   const std::byte* data_;
   std::size_t bit_size_;
-  std::uint8_t bit_offset_;  // always less than CHAR_BIT
+  std::size_t init_bit_size_;  // initial value of bit_size_
+  std::uint8_t bit_offset_;    // always less than CHAR_BIT
 
 public:
   /// An iterator over the bits in a bit_span.
@@ -82,7 +85,10 @@ class bit_span : public std::ranges::view_interface<bit_span>
   // NOLINTBEGIN(bugprone-easily-swappable-parameters)
   constexpr bit_span(
       const std::byte* data, std::size_t bit_size, std::uint8_t bit_offset = {})
-      : data_{data}, bit_size_{bit_size}, bit_offset_{bit_offset}
+      : data_{data},
+        bit_size_{bit_size},
+        init_bit_size_{bit_size},
+        bit_offset_{bit_offset}
   // NOLINTEND(bugprone-easily-swappable-parameters)
   {
     assert(
@@ -113,5 +119,59 @@ class bit_span : public std::ranges::view_interface<bit_span>
   {
     return iterator{*this, bit_offset_ + bit_size_};
   };
+
+  template <class T>
+  constexpr auto pop() -> T
+  {
+    assert(
+        bit_size_ >= sizeof(T) * CHAR_BIT and
+        "bit_span has insufficient "
+        "remaining bits to pop");
+    assert(bit_offset_ == 0 and "bit_span must be byte aligned to pop");
+    T res;
+    std::memcpy(&res, data_, sizeof(T));
+    std::advance(data_, sizeof(T));
+    bit_size_ -= sizeof(T) * CHAR_BIT;
+    if constexpr (std::endian::native == std::endian::big) {
+      res = std::byteswap(res);
+    }
+    return res;
+  }
+
+  constexpr auto pop_8() -> std::uint8_t { return pop<std::uint8_t>(); }
+
+  constexpr auto pop_16() -> std::uint16_t { return pop<std::uint16_t>(); }
+
+  /// Consumes the given number of bits. Advances the start of the view.
+  ///
+  /// @pre n <= std::ranges::size(this)
+  constexpr auto consume(std::size_t n) -> void
+  {
+    assert(n <= bit_size_);
+    bit_size_ -= n;
+    // invariant
+    assert(bit_offset_ < CHAR_BIT);
+    bit_offset_ += n % CHAR_BIT;
+    if (bit_offset_ >= CHAR_BIT) {
+      bit_offset_ -= CHAR_BIT;
+      n += CHAR_BIT;
+    }
+    std::advance(data_, n / CHAR_BIT);
+  }
+
+  /// Consumes bits until the start is aligned to a byte boundary.
+  constexpr auto consume_to_byte_boundary() -> void
+  {
+    if (bit_offset_) {
+      consume(CHAR_BIT - bit_offset_);
+    }
+  }
+
+  /// Returns the number of bits that have been consumed (or popped).
+  [[nodiscard]]
+  constexpr auto n_consumed() const -> size_t
+  {
+    return init_bit_size_ - bit_size_;
+  }
 };
 }  // namespace starflate::huffman

huffman/src/utility.hpp

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <array>
 #include <cstddef>
 
 namespace starflate::huffman {
@@ -27,4 +28,10 @@ struct symbol_bitsize_tag
 };
 inline constexpr auto symbol_bitsize = symbol_bitsize_tag{};
 
+template <class... Ts>
+constexpr auto byte_array(Ts... values)
+{
+  return std::array<std::byte, sizeof...(Ts)>{std::byte(values)...};
+}
+
 }  // namespace starflate::huffman

huffman/test/bit_span_test.cpp

@@ -1,11 +1,12 @@
 #include "huffman/huffman.hpp"
+#include "huffman/src/utility.hpp"
 
 #include <boost/ut.hpp>
 
 #include <array>
 #include <climits>
-#include <cstddef>
 #include <cstdint>
+#include <numeric>
 #include <vector>
 
 auto main() -> int
@@ -18,7 +19,7 @@ auto main() -> int
   using namespace huffman::literals;
 
   test("basic") = [] {
-    static constexpr std::array data{std::byte{0b10101010}, std::byte{0xff}};
+    static constexpr auto data = huffman::byte_array(0b10101010, 0xff);
     // leave off the last bit of the last byte
     constexpr huffman::bit_span span{data.data(), (data.size() * CHAR_BIT) - 1};
     constexpr std::string_view expected = "010101011111111";
@@ -30,8 +31,7 @@ auto main() -> int
   };
 
   test("indexable") = [] {
-    static constexpr auto data =
-        std::array{std::byte{0b10101010}, std::byte{0xff}};
+    static constexpr auto data = huffman::byte_array(0b10101010, 0xff);
     constexpr auto bs = huffman::bit_span{data};
 
     // NOLINTBEGIN(readability-magic-numbers)
@@ -58,8 +58,7 @@ auto main() -> int
   };
 
   test("usable with non byte-aligned data") = [] {
-    static constexpr auto data =
-        std::array{std::byte{0b10101010}, std::byte{0xff}};
+    static constexpr auto data = huffman::byte_array(0b10101010, 0xff);
 
     static constexpr auto bit_size = 7;
     static constexpr auto bit_offset = 3;
@@ -84,12 +83,71 @@ auto main() -> int
   using ::boost::ut::operator|;
 
   test("aborts if bit offset too large") = [](auto bit_offset) {
-    static constexpr auto data =
-        std::array{std::byte{0b10101010}, std::byte{0xff}};
+    static constexpr auto data = huffman::byte_array(0b10101010, 0xff);
 
     expect(aborts([&] {
       static constexpr auto bit_size = 7;
       huffman::bit_span{data.begin(), bit_size, bit_offset};
     }));
   } | std::vector<std::uint8_t>{8, 9, 10};  // NOLINT(readability-magic-numbers)
+
+  std::vector<unsigned> n_to_consume(1Z + 2Z * CHAR_BIT);
+  std::iota(n_to_consume.begin(), n_to_consume.end(), 0);
+
+  test("consume") = [](auto n) {
+    static constexpr auto data = huffman::byte_array(0b10101010, 0b01010101);
+
+    static constexpr auto nth_bit = [](auto m) {
+      return huffman::bit{std::bitset<CHAR_BIT>{
+          std::to_integer<unsigned>(data[m / CHAR_BIT])}[m % CHAR_BIT]};
+    };
+
+    auto bits = huffman::bit_span{data};
+    if (n <= data.size() * CHAR_BIT) {
+      bits.consume(n);
+      expect(bits.n_consumed() == n);
+      expect(nth_bit(n) == bits[0]);
+      expect(CHAR_BIT * data.size() - n == bits.size());
+    } else {
+      expect(aborts([&] { bits.consume(n); }));
+    }
+  } | n_to_consume;
+
+  test("consume_to_byte_boundary") = [] {
+    static constexpr auto data = huffman::byte_array(0b10101010, 0b01010101);
+    huffman::bit_span span{data.data(), data.size() * CHAR_BIT};
+    expect(*span.begin() == 0_b);
+    expect(span.n_consumed() == 0);
+    // should be a no-op now.
+    span.consume_to_byte_boundary();
+    expect(*span.begin() == 0_b);
+    expect(span.n_consumed() == 0);
+
+    span.consume(1);
+
+    span.consume_to_byte_boundary();
+    expect(*span.begin() == 1_b);
+    expect(span.n_consumed() == CHAR_BIT);
+  };
+
+  test("pop") = [] {
+    // NOLINTBEGIN(readability-magic-numbers)
+    static constexpr auto data =
+        huffman::byte_array(0b10101010, 0b01010101, 0b11111111);
+    huffman::bit_span span{data.data(), data.size() * CHAR_BIT};
+    std::uint16_t got_16{span.pop_16()};
+    std::uint16_t expected_16{0b0101010110101010};
+    expect(got_16 == expected_16)
+        << "got: " << got_16 << " expected: " << expected_16;
+
+    expect(aborts([&] { span.pop_16(); }));
+
+    std::uint8_t got_8{span.pop_8()};
+    std::uint8_t expected_8{0b11111111};
+    expect(got_8 == expected_8)
+        << "got: " << got_8 << " expected: " << expected_8;
+
+    expect(aborts([&] { span.pop_8(); }));
+    // NOLINTEND(readability-magic-numbers)
+  };
 }