Add functions to read from/write to bitstream values with known max value

Bug: webrtc:10342
Change-Id: I701b09de574eb463daaf8d2c19008ac3452879eb
Reviewed-on: https://webrtc-review.googlesource.com/c/src/+/144033
Commit-Queue: Danil Chapovalov <danilchap@webrtc.org>
Reviewed-by: Karl Wiberg <kwiberg@webrtc.org>
Reviewed-by: Philip Eliasson <philipel@webrtc.org>
Cr-Commit-Position: refs/heads/master@{#28465}

rtc_base/bit_buffer.cc

@@ -159,6 +159,29 @@ bool BitBuffer::ConsumeBits(size_t bit_count) {
   return true;
 }
 
+bool BitBuffer::ReadNonSymmetric(uint32_t* val, uint32_t num_values) {
+  RTC_DCHECK_GT(num_values, 0);
+  RTC_DCHECK_LE(num_values, uint32_t{1} << 31);
+  size_t count_bits = CountBits(num_values);
+  uint32_t num_min_bits_values = (uint32_t{1} << count_bits) - num_values;
+
+  if (!ReadBits(val, count_bits - 1)) {
+    return false;
+  }
+
+  if (*val < num_min_bits_values) {
+    return true;
+  }
+
+  uint32_t extra_bit;
+  if (!ReadBits(&extra_bit, /*bit_count=*/1)) {
+    return false;
+  }
+
+  *val = (*val << 1) + extra_bit - num_min_bits_values;
+  return true;
+}
+
 bool BitBuffer::ReadExponentialGolomb(uint32_t* val) {
   if (!val) {
     return false;
@@ -282,6 +305,27 @@ bool BitBufferWriter::WriteBits(uint64_t val, size_t bit_count) {
   return ConsumeBits(total_bits);
 }
 
+bool BitBufferWriter::WriteNonSymmetric(uint32_t val, uint32_t num_values) {
+  RTC_DCHECK_LT(val, num_values);
+  RTC_DCHECK_LE(num_values, uint32_t{1} << 31);
+  size_t count_bits = CountBits(num_values);
+  uint32_t num_min_bits_values = (uint32_t{1} << count_bits) - num_values;
+
+  return val < num_min_bits_values
+             ? WriteBits(val, count_bits - 1)
+             : WriteBits(val + num_min_bits_values, count_bits);
+}
+
+size_t BitBufferWriter::SizeNonSymmetricBits(uint32_t val,
+                                             uint32_t num_values) {
+  RTC_DCHECK_LT(val, num_values);
+  RTC_DCHECK_LE(num_values, uint32_t{1} << 31);
+  size_t count_bits = CountBits(num_values);
+  uint32_t num_min_bits_values = (uint32_t{1} << count_bits) - num_values;
+
+  return val < num_min_bits_values ? (count_bits - 1) : count_bits;
+}
+
 bool BitBufferWriter::WriteExponentialGolomb(uint32_t val) {
   // We don't support reading UINT32_MAX, because it doesn't fit in a uint32_t
   // when encoded, so don't support writing it either.

rtc_base/bit_buffer.h

@@ -51,6 +51,18 @@ class BitBuffer {
   // offset.
   bool PeekBits(uint32_t* val, size_t bit_count);
 
+  // Reads value in range [0, num_values - 1].
+  // This encoding is similar to ReadBits(val, Ceil(Log2(num_values)),
+  // but reduces wastage incurred when encoding non-power of two value ranges
+  // Non symmetric values are encoded as:
+  // 1) n = countbits(num_values)
+  // 2) k = (1 << n) - num_values
+  // Value v in range [0, k - 1] is encoded in (n-1) bits.
+  // Value v in range [k, num_values - 1] is encoded as (v+k) in n bits.
+  // https://aomediacodec.github.io/av1-spec/#nsn
+  // Returns false if there isn't enough data left.
+  bool ReadNonSymmetric(uint32_t* val, uint32_t num_values);
+
   // Reads the exponential golomb encoded value at the current offset.
   // Exponential golomb values are encoded as:
   // 1) x = source val + 1
@@ -106,6 +118,14 @@ class BitBufferWriter : public BitBuffer {
   // room left for the specified number of bits.
   bool WriteBits(uint64_t val, size_t bit_count);
 
+  // Writes value in range [0, num_values - 1]
+  // See ReadNonSymmetric documentation for the format,
+  // Call SizeNonSymmetricBits to get number of bits needed to store the value.
+  // Returns false if there isn't enough room left for the value.
+  bool WriteNonSymmetric(uint32_t val, uint32_t num_values);
+  // Returns number of bits required to store |val| with NonSymmetric encoding.
+  static size_t SizeNonSymmetricBits(uint32_t val, uint32_t num_values);
+
   // Writes the exponential golomb encoded version of the supplied value.
   // Returns false if there isn't enough room left for the value.
   bool WriteExponentialGolomb(uint32_t val);

rtc_base/bit_buffer_unittest.cc

@@ -14,10 +14,13 @@
 
 #include "rtc_base/arraysize.h"
 #include "rtc_base/byte_buffer.h"
+#include "test/gmock.h"
 #include "test/gtest.h"
 
 namespace rtc {
 
+using ::testing::ElementsAre;
+
 TEST(BitBufferTest, ConsumeBits) {
   const uint8_t bytes[64] = {0};
   BitBuffer buffer(bytes, 32);
@@ -178,6 +181,79 @@ TEST(BitBufferTest, SetOffsetValues) {
 #endif
 }
 
+TEST(BitBufferTest, ReadNonSymmetricSameNumberOfBitsWhenNumValuesPowerOf2) {
+  const uint8_t bytes[2] = {0xf3, 0xa0};
+  BitBuffer reader(bytes, 2);
+
+  uint32_t values[4];
+  ASSERT_EQ(reader.RemainingBitCount(), 16u);
+  EXPECT_TRUE(reader.ReadNonSymmetric(&values[0], /*num_values=*/1 << 4));
+  EXPECT_TRUE(reader.ReadNonSymmetric(&values[1], /*num_values=*/1 << 4));
+  EXPECT_TRUE(reader.ReadNonSymmetric(&values[2], /*num_values=*/1 << 4));
+  EXPECT_TRUE(reader.ReadNonSymmetric(&values[3], /*num_values=*/1 << 4));
+  ASSERT_EQ(reader.RemainingBitCount(), 0u);
+
+  EXPECT_THAT(values, ElementsAre(0xf, 0x3, 0xa, 0x0));
+}
+
+TEST(BitBufferWriterTest,
+     WriteNonSymmetricSameNumberOfBitsWhenNumValuesPowerOf2) {
+  uint8_t bytes[2] = {};
+  BitBufferWriter writer(bytes, 2);
+
+  ASSERT_EQ(writer.RemainingBitCount(), 16u);
+  EXPECT_TRUE(writer.WriteNonSymmetric(0xf, /*num_values=*/1 << 4));
+  ASSERT_EQ(writer.RemainingBitCount(), 12u);
+  EXPECT_TRUE(writer.WriteNonSymmetric(0x3, /*num_values=*/1 << 4));
+  ASSERT_EQ(writer.RemainingBitCount(), 8u);
+  EXPECT_TRUE(writer.WriteNonSymmetric(0xa, /*num_values=*/1 << 4));
+  ASSERT_EQ(writer.RemainingBitCount(), 4u);
+  EXPECT_TRUE(writer.WriteNonSymmetric(0x0, /*num_values=*/1 << 4));
+  ASSERT_EQ(writer.RemainingBitCount(), 0u);
+
+  EXPECT_THAT(bytes, ElementsAre(0xf3, 0xa0));
+}
+
+TEST(BitBufferWriterTest, NonSymmetricReadsMatchesWrites) {
+  uint8_t bytes[2] = {};
+  BitBufferWriter writer(bytes, 2);
+
+  EXPECT_EQ(BitBufferWriter::SizeNonSymmetricBits(/*val=*/1, /*num_values=*/6),
+            2u);
+  EXPECT_EQ(BitBufferWriter::SizeNonSymmetricBits(/*val=*/2, /*num_values=*/6),
+            3u);
+  // Values [0, 1] can fit into two bit.
+  ASSERT_EQ(writer.RemainingBitCount(), 16u);
+  EXPECT_TRUE(writer.WriteNonSymmetric(/*val=*/0, /*num_values=*/6));
+  ASSERT_EQ(writer.RemainingBitCount(), 14u);
+  EXPECT_TRUE(writer.WriteNonSymmetric(/*val=*/1, /*num_values=*/6));
+  ASSERT_EQ(writer.RemainingBitCount(), 12u);
+  // Values [2, 5] require 3 bits.
+  EXPECT_TRUE(writer.WriteNonSymmetric(/*val=*/2, /*num_values=*/6));
+  ASSERT_EQ(writer.RemainingBitCount(), 9u);
+  EXPECT_TRUE(writer.WriteNonSymmetric(/*val=*/3, /*num_values=*/6));
+  ASSERT_EQ(writer.RemainingBitCount(), 6u);
+  EXPECT_TRUE(writer.WriteNonSymmetric(/*val=*/4, /*num_values=*/6));
+  ASSERT_EQ(writer.RemainingBitCount(), 3u);
+  EXPECT_TRUE(writer.WriteNonSymmetric(/*val=*/5, /*num_values=*/6));
+  ASSERT_EQ(writer.RemainingBitCount(), 0u);
+
+  // Bit values are
+  // 00.01.100.101.110.111 = 00011001|01110111 = 0x19|77
+  EXPECT_THAT(bytes, ElementsAre(0x19, 0x77));
+
+  rtc::BitBuffer reader(bytes, 2);
+  uint32_t values[6];
+  EXPECT_TRUE(reader.ReadNonSymmetric(&values[0], /*num_values=*/6));
+  EXPECT_TRUE(reader.ReadNonSymmetric(&values[1], /*num_values=*/6));
+  EXPECT_TRUE(reader.ReadNonSymmetric(&values[2], /*num_values=*/6));
+  EXPECT_TRUE(reader.ReadNonSymmetric(&values[3], /*num_values=*/6));
+  EXPECT_TRUE(reader.ReadNonSymmetric(&values[4], /*num_values=*/6));
+  EXPECT_TRUE(reader.ReadNonSymmetric(&values[5], /*num_values=*/6));
+
+  EXPECT_THAT(values, ElementsAre(0, 1, 2, 3, 4, 5));
+}
+
 uint64_t GolombEncoded(uint32_t val) {
   val++;
   uint32_t bit_counter = val;