bitstream: one 8-byte load in BitReader::peek, direct read_bit (#516)

peek built its window from five bounds-checked byte loads on every call.
While at least 8 bytes remain it now does one big-endian 8-byte load; the
last 7 bytes and positions past the end keep the old zero-padded byte
loop in a cold function. read_bit reads its byte directly instead of
going through read_bits(1).

Behaviour is unchanged: values, errors and positions are identical,
including peeks near and past the end of the buffer. Tests compare the
reader against a verbatim copy of the previous logic over random
operation sequences, plus the exact edge positions and a bitwise
reference for peek and read_bit.

Closes #516

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This commit is contained in:
vangel76andClaude Sonnet 5.5 committed 2026-11-02 21:36:05 +01:00
1 parent ea1fbeb05b
commit b25d6e4797
1 file changed
+194 -1
+194 -1
View File
@@ -84,6 +84,16 @@ impl<'a> BitReader<'a> {
}
let byte = self.pos / 8;
let shift = self.pos % 8;
// One 8-byte load while at least 8 bytes are left (`shift` <= 7 leaves 57 valid bits).
if let Some(w) = self.data.get(byte..).and_then(<[u8]>::first_chunk::<8>) {
return ((u64::from_be_bytes(*w) << shift) >> (64 - n)) as u32;
}
self.peek_tail(n, byte, shift)
}
/// [`Self::peek`] for the last 7 bytes and past the end: byte loads, zero padded.
#[cold]
fn peek_tail(&self, n: u32, byte: usize, shift: usize) -> u32 {
let mut v: u64 = 0;
for i in 0..5 {
v = (v << 8) | *self.data.get(byte.saturating_add(i)).unwrap_or(&0) as u64;
@@ -118,7 +128,14 @@ impl<'a> BitReader<'a> {
#[inline]
pub fn read_bit(&mut self) -> Result<bool> {
self.read_bits(1).map(|b| b == 1)
if self.bits_left() == 0 {
return Err(BitError::Eof);
}
// `pos` is below `len * 8` here, so the byte exists.
let byte = self.data.get(self.pos / 8).copied().unwrap_or(0);
let bit = (byte >> (7 - self.pos % 8)) & 1;
self.pos += 1;
Ok(bit == 1)
}
pub fn read_flag(&mut self) -> Result<bool> {
self.read_bit()
@@ -531,6 +548,182 @@ mod tests {
assert_eq!(w.finish(), [0, 0, 0, 0, 0xFF, 0xFF, 0xFF, 0xFF]);
}
/// Bit `p` of `d`, MSB first, zero past the end.
fn bit_at(d: &[u8], p: usize) -> u32 {
d.get(p / 8).map_or(0, |&b| u32::from(b >> (7 - p % 8)) & 1)
}
/// Bit-by-bit reference for `peek`: the low 32 bits of the `n`-bit field at `pos`.
fn reference_peek(d: &[u8], pos: usize, n: u32) -> u32 {
let start = pos + n.saturating_sub(32) as usize;
(0..n.min(32) as usize).fold(0, |v, i| (v << 1) | bit_at(d, start + i))
}
proptest! {
#[test]
fn peek_and_read_bit_match_a_bitwise_reference(data in proptest::collection::vec(any::<u8>(), 0..24)) {
// Every position, including the last 7 bytes (byte loads) and past the end.
for pos in 0..=data.len() * 8 + 9 {
let mut r = BitReader::new(&data);
r.seek_bits(pos);
for n in 0..=40 {
prop_assert_eq!(r.peek(n), reference_peek(&data, pos, n), "pos {} n {}", pos, n);
}
let bit = r.read_bit();
if pos < data.len() * 8 {
prop_assert_eq!(bit, Ok(bit_at(&data, pos) == 1));
prop_assert_eq!(r.position(), pos + 1);
} else {
prop_assert_eq!(bit, Err(BitError::Eof));
prop_assert_eq!(r.position(), pos);
}
}
}
}
/// The reader as it was before the 8-byte load, kept verbatim as the reference.
struct OldReader<'a> {
data: &'a [u8],
pos: usize,
}
impl OldReader<'_> {
fn bits_left(&self) -> usize {
self.data.len().saturating_mul(8).saturating_sub(self.pos)
}
fn peek(&self, n: u32) -> u32 {
if n > 32 {
let r = OldReader { data: self.data, pos: self.pos.saturating_add((n - 32) as usize) };
return r.peek(32);
}
if n == 0 {
return 0;
}
let byte = self.pos / 8;
let shift = self.pos % 8;
let mut v: u64 = 0;
for i in 0..5 {
v = (v << 8) | *self.data.get(byte.saturating_add(i)).unwrap_or(&0) as u64;
}
((v << (24 + shift)) >> (64 - n)) as u32
}
fn read_bits(&mut self, n: u32) -> Result<u32> {
if n as usize > self.bits_left() {
return Err(BitError::Eof);
}
if n > 32 {
self.pos += (n - 32) as usize;
return self.read_bits(32);
}
let v = self.peek(n);
self.pos += n as usize;
Ok(v)
}
fn read_bit(&mut self) -> Result<bool> {
self.read_bits(1).map(|b| b == 1)
}
fn skip(&mut self, n: usize) -> Result<()> {
if n > self.bits_left() {
return Err(BitError::Eof);
}
self.pos += n;
Ok(())
}
}
/// Fixed-seed xorshift64*, so the differential test is reproducible.
struct Rng(u64);
impl Rng {
fn next(&mut self) -> u64 {
self.0 ^= self.0 >> 12;
self.0 ^= self.0 << 25;
self.0 ^= self.0 >> 27;
self.0.wrapping_mul(0x2545_F491_4F6C_DD1D)
}
fn below(&mut self, n: u64) -> u64 {
self.next() % n
}
}
#[test]
fn reader_matches_the_pre_change_reader_on_random_operations() {
let mut rng = Rng(0x9E37_79B9_7F4A_7C15);
for _ in 0..3000 {
// Lengths cluster around the 8-byte boundary of the fast path; some buffers are zero heavy.
let len = if rng.below(2) == 0 { rng.below(20) } else { rng.below(200) } as usize;
let zero_heavy = rng.below(3) == 0;
let data: Vec<u8> = (0..len).map(|_| if zero_heavy && rng.below(4) != 0 { 0 } else { rng.next() as u8 }).collect();
let mut new = BitReader::new(&data);
let mut old = OldReader { data: &data, pos: 0 };
for step in 0..200 {
let n = match rng.below(8) {
0 => [31, 32, 33, 63, 64, 65, u32::MAX][rng.below(7) as usize],
_ => rng.below(33) as u32,
};
let ctx = format!("len {len} step {step} pos {} n {n}", old.pos);
match rng.below(8) {
0 | 1 => assert_eq!(new.read_bits(n), old.read_bits(n), "read_bits {ctx}"),
2 => assert_eq!(new.read_bit(), old.read_bit(), "read_bit {ctx}"),
3 | 4 => assert_eq!(new.peek(n), old.peek(n), "peek {ctx}"),
5 => {
let k = [rng.below(70) as usize, len * 8 + rng.below(10) as usize, usize::MAX - rng.below(8) as usize][rng.below(3) as usize];
assert_eq!(new.skip(k), old.skip(k), "skip {k} {ctx}");
}
6 => {
let p = [
rng.below(len as u64 * 8 + 12) as usize,
len * 8,
len.saturating_sub(7) * 8 + rng.below(8) as usize,
usize::MAX - rng.below(4) as usize,
][rng.below(4) as usize];
new.seek_bits(p);
old.pos = p;
}
// `byte_align` itself overflows within 7 bits of `usize::MAX`, before and after this change.
_ if old.pos > usize::MAX - 8 => {}
_ => {
new.byte_align();
old.pos = old.pos.div_ceil(8) * 8;
}
}
assert_eq!(new.position(), old.pos, "position {ctx}");
assert_eq!(new.bits_left(), old.bits_left(), "bits_left {ctx}");
}
}
}
#[test]
fn peek_and_read_bit_at_the_edges_of_the_fast_path() {
// Exactly 8 bytes: the last position with an 8-byte load is bit 0 of byte 0.
let d: Vec<u8> = (1..=9).collect();
for len in [0, 1, 7, 8, 9] {
let data = &d[..len];
let end = len * 8;
for pos in [0, 1, 7, 8, end.saturating_sub(64), end.saturating_sub(57), end.saturating_sub(33), end.saturating_sub(1), end, end + 1] {
for n in [0, 1, 7, 8, 31, 32, 33, 40, 64, u32::MAX] {
let mut new = BitReader::new(data);
let mut old = OldReader { data, pos };
new.seek_bits(pos);
assert_eq!(new.peek(n), old.peek(n), "peek len {len} pos {pos} n {n}");
assert_eq!(new.read_bits(n), old.read_bits(n), "read_bits len {len} pos {pos} n {n}");
assert_eq!(new.position(), old.pos);
}
let mut new = BitReader::new(data);
let mut old = OldReader { data, pos };
new.seek_bits(pos);
assert_eq!(new.read_bit(), old.read_bit(), "read_bit len {len} pos {pos}");
assert_eq!(new.position(), old.pos);
}
}
// Positions near usize::MAX must neither panic nor read.
let mut r = BitReader::new(&d);
r.seek_bits(usize::MAX);
assert_eq!(r.peek(32), 0);
assert_eq!(r.read_bit(), Err(BitError::Eof));
assert_eq!(r.position(), usize::MAX);
}
proptest! {
#[test]
fn roundtrip(vals in proptest::collection::vec((0u32..u32::MAX, 1u32..=32), 0..200), ues in proptest::collection::vec(0u32..u32::MAX - 1, 0..50)) {