diff --git a/rtc_base/bit_buffer.cc b/rtc_base/bit_buffer.cc index 540141fe52..222d30ca0b 100644 --- a/rtc_base/bit_buffer.cc +++ b/rtc_base/bit_buffer.cc @@ -141,10 +141,48 @@ bool BitBuffer::PeekBits(uint32_t* val, size_t bit_count) { return true; } +bool BitBuffer::PeekBits(uint64_t* val, size_t bit_count) { + // TODO(nisse): Could allow bit_count == 0 and always return success. But + // current code reads one byte beyond end of buffer in the case that + // RemainingBitCount() == 0 and bit_count == 0. + RTC_DCHECK(bit_count > 0); + if (!val || bit_count > RemainingBitCount() || bit_count > 64) { + return false; + } + const uint8_t* bytes = bytes_ + byte_offset_; + size_t remaining_bits_in_current_byte = 8 - bit_offset_; + uint64_t bits = LowestBits(*bytes++, remaining_bits_in_current_byte); + // If we're reading fewer bits than what's left in the current byte, just + // return the portion of this byte that we need. + if (bit_count < remaining_bits_in_current_byte) { + *val = HighestBits(bits, bit_offset_ + bit_count); + return true; + } + // Otherwise, subtract what we've read from the bit count and read as many + // full bytes as we can into bits. + bit_count -= remaining_bits_in_current_byte; + while (bit_count >= 8) { + bits = (bits << 8) | *bytes++; + bit_count -= 8; + } + // Whatever we have left is smaller than a byte, so grab just the bits we need + // and shift them into the lowest bits. + if (bit_count > 0) { + bits <<= bit_count; + bits |= HighestBits(*bytes, bit_count); + } + *val = bits; + return true; +} + bool BitBuffer::ReadBits(uint32_t* val, size_t bit_count) { return PeekBits(val, bit_count) && ConsumeBits(bit_count); } +bool BitBuffer::ReadBits(uint64_t* val, size_t bit_count) { + return PeekBits(val, bit_count) && ConsumeBits(bit_count); +} + bool BitBuffer::ConsumeBytes(size_t byte_count) { return ConsumeBits(byte_count * 8); } diff --git a/rtc_base/bit_buffer.h b/rtc_base/bit_buffer.h index de7bf02d56..045d369249 100644 --- a/rtc_base/bit_buffer.h +++ b/rtc_base/bit_buffer.h @@ -45,11 +45,13 @@ class BitBuffer { // Reads bit-sized values from the buffer. Returns false if there isn't enough // data left for the specified bit count. bool ReadBits(uint32_t* val, size_t bit_count); + bool ReadBits(uint64_t* val, size_t bit_count); // Peeks bit-sized values from the buffer. Returns false if there isn't enough // data left for the specified number of bits. Doesn't move the current // offset. bool PeekBits(uint32_t* val, size_t bit_count); + bool PeekBits(uint64_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)), diff --git a/rtc_base/bit_buffer_unittest.cc b/rtc_base/bit_buffer_unittest.cc index 656682c2ef..86c4f7fa5c 100644 --- a/rtc_base/bit_buffer_unittest.cc +++ b/rtc_base/bit_buffer_unittest.cc @@ -142,6 +142,38 @@ TEST(BitBufferTest, ReadBits) { EXPECT_FALSE(buffer.ReadBits(&val, 1)); } +TEST(BitBufferTest, ReadBits64) { + const uint8_t bytes[] = {0x4D, 0x32, 0xAB, 0x54, 0x00, 0xFF, 0xFE, 0x01, + 0xAB, 0xCD, 0xEF, 0x01, 0x23, 0x45, 0x67, 0x89}; + BitBuffer buffer(bytes, 16); + uint64_t val; + + // Peek and read first 33 bits. + EXPECT_TRUE(buffer.PeekBits(&val, 33)); + EXPECT_EQ(0x4D32AB5400FFFE01ull >> (64 - 33), val); + val = 0; + EXPECT_TRUE(buffer.ReadBits(&val, 33)); + EXPECT_EQ(0x4D32AB5400FFFE01ull >> (64 - 33), val); + + // Peek and read next 31 bits. + constexpr uint64_t kMask31Bits = (1ull << 32) - 1; + EXPECT_TRUE(buffer.PeekBits(&val, 31)); + EXPECT_EQ(0x4D32AB5400FFFE01ull & kMask31Bits, val); + val = 0; + EXPECT_TRUE(buffer.ReadBits(&val, 31)); + EXPECT_EQ(0x4D32AB5400FFFE01ull & kMask31Bits, val); + + // Peek and read remaining 64 bits. + EXPECT_TRUE(buffer.PeekBits(&val, 64)); + EXPECT_EQ(0xABCDEF0123456789ull, val); + val = 0; + EXPECT_TRUE(buffer.ReadBits(&val, 64)); + EXPECT_EQ(0xABCDEF0123456789ull, val); + + // Nothing more to read. + EXPECT_FALSE(buffer.ReadBits(&val, 1)); +} + TEST(BitBufferDeathTest, SetOffsetValues) { uint8_t bytes[4] = {0}; BitBufferWriter buffer(bytes, 4);