From b25d6e47976ee19b3fe2ee79884028a9bbe46f8f Mon Sep 17 00:00:00 2001 From: vangel76 Date: Mon, 2 Nov 2026 21:36:05 +0100 Subject: [PATCH] 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 --- crates/media/src/bitstream.rs | 195 +++++++++++++++++++++++++++++++++- 1 file changed, 194 insertions(+), 1 deletion(-) diff --git a/crates/media/src/bitstream.rs b/crates/media/src/bitstream.rs index adf166ca..772e1a63 100644 --- a/crates/media/src/bitstream.rs +++ b/crates/media/src/bitstream.rs @@ -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 { - 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 { 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::(), 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 { + 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 { + 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 = (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 = (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)) {