From 2e4283c70d245adb40de60df6a434ad04a341229 Mon Sep 17 00:00:00 2001
From: Oliver Lee <oliverzlee@gmail.com>
Date: Wed, 20 Sep 2023 04:36:17 +0200
Subject: [PATCH] handle non-byte aligned data with bit_span

Change-Id: Id460151a080e8a1c55e382182e34e60eed71c78c
---
 huffman/src/bit_span.hpp       | 23 ++++++++++++++++------
 huffman/test/bit_span_test.cpp | 36 ++++++++++++++++++++++++++++++++--
 2 files changed, 51 insertions(+), 8 deletions(-)

diff --git a/huffman/src/bit_span.hpp b/huffman/src/bit_span.hpp
index eedb6fb..9213c07 100644
--- a/huffman/src/bit_span.hpp
+++ b/huffman/src/bit_span.hpp
@@ -17,6 +17,7 @@ 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
 
 public:
   /// An iterator over the bits in a bit_span.
@@ -34,7 +35,7 @@ class bit_span : public std::ranges::view_interface<bit_span>
     using pointer = void;
 
     iterator() = default;
-    constexpr iterator(const bit_span& parent, size_t offset)
+    constexpr iterator(const bit_span& parent, std::size_t offset)
         : parent_(&parent), offset_(offset)
     {
       assert(offset_ <= std::numeric_limits<difference_type>::max());
@@ -73,9 +74,19 @@ class bit_span : public std::ranges::view_interface<bit_span>
   ///
   /// @param data a pointer to the first byte of the data.
   /// @param bit_size the number of bits in the data.
-  constexpr bit_span(const std::byte* data, std::size_t bit_size)
-      : data_(data), bit_size_(bit_size)
-  {}
+  /// @param bit_offset bit offset of data, allowing a non-byte aligned range
+  ///
+  /// @pre offset < CHAR_BIT
+  ///
+  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}
+  {
+    assert(
+        bit_offset < CHAR_BIT and
+        "bit offset exceeds number of bits in a "
+        "byte");
+  }
 
   template <std::ranges::contiguous_range R>
     requires std::ranges::borrowed_range<R>
@@ -89,12 +100,12 @@ class bit_span : public std::ranges::view_interface<bit_span>
   [[nodiscard]]
   constexpr auto begin() const -> iterator
   {
-    return iterator{*this, 0};
+    return iterator{*this, bit_offset_};
   };
   [[nodiscard]]
   constexpr auto end() const -> iterator
   {
-    return iterator{*this, bit_size_};
+    return iterator{*this, bit_offset_ + bit_size_};
   };
 };
 }  // namespace starflate::huffman
diff --git a/huffman/test/bit_span_test.cpp b/huffman/test/bit_span_test.cpp
index 840f065..90d995a 100644
--- a/huffman/test/bit_span_test.cpp
+++ b/huffman/test/bit_span_test.cpp
@@ -6,6 +6,7 @@
 #include <climits>
 #include <cstddef>
 #include <cstdint>
+#include <vector>
 
 auto main() -> int
 {
@@ -17,8 +18,7 @@ auto main() -> int
   using namespace huffman::literals;
 
   test("basic") = [] {
-    static constexpr std::array<std::byte, 2> data{
-        std::byte{0b10101010}, std::byte{0xff}};
+    static constexpr std::array data{std::byte{0b10101010}, std::byte{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 = "101010101111111";
@@ -52,4 +52,36 @@ auto main() -> int
     static_assert(bs[14] == 1_b);
     static_assert(bs[15] == 1_b);
   };
+
+  test("usable with non byte-aligned data") = [] {
+    static constexpr auto data =
+        std::array{std::byte{0b10101010}, std::byte{0xff}};
+
+    static constexpr auto bit_size = 7;
+    static constexpr auto bit_offset = 3;
+    constexpr auto bs = huffman::bit_span{data.begin(), bit_size, bit_offset};
+
+    // from first byte
+    static_assert(bs[0] == 0_b);
+    static_assert(bs[1] == 1_b);
+    static_assert(bs[2] == 0_b);
+    static_assert(bs[3] == 1_b);
+    static_assert(bs[4] == 0_b);
+
+    // from second byte
+    static_assert(bs[5] == 1_b);
+    static_assert(bs[6] == 1_b);
+  };
+
+  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}};
+
+    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};
 }