diff --git a/crates/media/src/bitstream.rs b/crates/media/src/bitstream.rs index 772e1a63..d82a7dce 100644 --- a/crates/media/src/bitstream.rs +++ b/crates/media/src/bitstream.rs @@ -150,7 +150,25 @@ impl<'a> BitReader<'a> { } /// Raw Exp-Golomb codeNum, up to 32 leading zeros (so below 2^33). + #[inline] fn read_code_num(&mut self) -> Result { + // Codes up to 31 bits (15 leading zeros) that lie wholly inside the data: one peek. + let w = self.peek(32); + let zeros = w.leading_zeros(); + if zeros < 16 { + let len = 2 * zeros + 1; + if len as usize <= self.bits_left() { + self.pos += len as usize; + // The `len`-bit field is 2^zeros + suffix; codeNum is that minus one. + return Ok(u64::from((w >> (32 - len)) - 1)); + } + } + self.read_code_num_slow() + } + + /// [`Self::read_code_num`] bit by bit: long codes, the end of the data and invalid codes. + #[cold] + fn read_code_num_slow(&mut self) -> Result { let mut zeros = 0u32; while !self.read_bit()? { zeros += 1; @@ -239,9 +257,12 @@ impl BitWriter { let x = v + 1; let len = 64 - x.leading_zeros(); let zeros = len - 1; - for _ in 0..zeros { - self.write_bit(false); + if zeros + len <= 32 { + // The leading zeros are the high bits of a (zeros + len)-bit field. + self.write_bits(x as u32, zeros + len); + return; } + self.write_bits(0, zeros); if len > 32 { self.write_bits((x >> 32) as u32, len - 32); self.write_bits(x as u32, 32); @@ -724,6 +745,143 @@ mod tests { assert_eq!(r.position(), usize::MAX); } + /// Verbatim copy of the bit-by-bit `read_code_num` that the one-peek version replaced. + fn old_read_code_num(r: &mut BitReader) -> Result { + let mut zeros = 0u32; + while !r.read_bit()? { + zeros += 1; + if zeros > 32 { + return Err(BitError::InvalidExpGolomb); + } + } + if zeros == 0 { + return Ok(0); + } + let v = r.read_bits_u64(zeros)?; + Ok((1u64 << zeros) - 1 + v) + } + + fn old_read_ue(r: &mut BitReader) -> Result { + u32::try_from(old_read_code_num(r)?).map_err(|_| BitError::InvalidExpGolomb) + } + + fn old_read_se(r: &mut BitReader) -> Result { + let k = old_read_code_num(r)? as i64; + let v = if k & 1 == 1 { (k + 1) / 2 } else { -(k / 2) }; + i32::try_from(v).map_err(|_| BitError::InvalidExpGolomb) + } + + /// Verbatim copy of the `write_ue` / `write_code_num` before the one-call version. + fn old_write_code_num(w: &mut BitWriter, v: u64) { + let x = v + 1; + let len = 64 - x.leading_zeros(); + let zeros = len - 1; + for _ in 0..zeros { + w.write_bit(false); + } + if len > 32 { + w.write_bits((x >> 32) as u32, len - 32); + w.write_bits(x as u32, 32); + } else { + w.write_bits(x as u32, len); + } + } + + #[test] + fn exp_golomb_reader_matches_the_pre_change_reader() { + let mut rng = Rng(0xD1B5_4A32_D192_ED03); + for round in 0..20_000u32 { + let len = rng.below(17) as usize; + let mut data: Vec = (0..len).map(|_| rng.next() as u8).collect(); + if round % 2 == 0 && len > 0 { + // A zero run of 28..=40 bits, then a one bit, somewhere in the buffer. + let start = rng.below(len as u64 * 8) as usize; + let zeros = 28 + rng.below(13) as usize; + for p in start..(start + zeros + 1).min(len * 8) { + data[p / 8] &= !(0x80 >> (p % 8)); + } + if start + zeros < len * 8 { + let p = start + zeros; + data[p / 8] |= 0x80 >> (p % 8); + } + } + // Every start position, including the last bits and past the end. + for pos in 0..=len * 8 + 2 { + let mut old = BitReader::new(&data); + let mut new = BitReader::new(&data); + old.seek_bits(pos); + new.seek_bits(pos); + // Several codes in a row, so the position after each result is checked too. + for step in 0..4 { + if (round + step) % 2 == 0 { + assert_eq!(new.read_ue(), old_read_ue(&mut old), "ue {data:02x?} pos {pos} step {step}"); + } else { + assert_eq!(new.read_se(), old_read_se(&mut old), "se {data:02x?} pos {pos} step {step}"); + } + assert_eq!(new.position(), old.position(), "position {data:02x?} pos {pos} step {step}"); + assert_eq!(new.bits_left(), old.bits_left()); + } + } + } + } + + #[test] + fn exp_golomb_reader_edge_codes() { + // 31-bit (15 zeros) and 32-bit (16 zeros) codes, 33-bit-wide codes, and 32 or 33 zeros. + for zeros in 0..=34usize { + for extra in 0..4usize { + let total = (2 * zeros + 1 + extra).div_ceil(8) * 8; + let mut data = vec![0u8; total / 8]; + data[zeros / 8] |= 0x80 >> (zeros % 8); + for tail in [0u8, 0xFF] { + let mut d = data.clone(); + for p in 2 * zeros + 1..total { + if tail != 0 { + d[p / 8] |= 0x80 >> (p % 8); + } + } + for cut in [d.len(), d.len().saturating_sub(1)] { + let d = &d[..cut]; + let (mut o, mut n) = (BitReader::new(d), BitReader::new(d)); + assert_eq!(n.read_ue(), old_read_ue(&mut o), "ue zeros {zeros} cut {cut}"); + assert_eq!(n.position(), o.position()); + let (mut o, mut n) = (BitReader::new(d), BitReader::new(d)); + assert_eq!(n.read_se(), old_read_se(&mut o), "se zeros {zeros} cut {cut}"); + assert_eq!(n.position(), o.position()); + } + } + } + } + } + + #[test] + fn write_ue_matches_the_pre_change_writer() { + let mut rng = Rng(0x1234_5678_9ABC_DEF1); + let mut values: Vec = vec![0, 1, 2, 3, 6, 7, 8, 65534, 65535, 65536, u32::MAX - 1, u32::MAX]; + values.extend((0..32).flat_map(|b| [(1u32 << b) - 1, 1 << b, (1u32 << b) + 1])); + values.extend((0..5000).map(|_| (rng.next() as u32) >> rng.below(32))); + for (i, &v) in values.iter().enumerate() { + for lead in [0u32, 1, 3, 7] { + let mut new = BitWriter::new(); + let mut old = BitWriter::new(); + new.write_bits(0x55, lead); + old.write_bits(0x55, lead); + new.write_ue(v); + old_write_code_num(&mut old, u64::from(v)); + assert_eq!(new.bit_len(), old.bit_len(), "value {v} lead {lead} (#{i})"); + assert_eq!(new.finish(), old.finish(), "value {v} lead {lead} (#{i})"); + } + } + // write_se goes through the same code with codeNum 2^32 at i32::MIN. + for v in [0, 1, -1, 2, -2, i32::MAX, i32::MIN, i32::MIN + 1] { + let mut new = BitWriter::new(); + new.write_se(v); + new.rbsp_trailing(); + let b = new.finish(); + assert_eq!(BitReader::new(&b).read_se(), Ok(v)); + } + } + 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)) {