diff --git a/src/bit_encoding/encode.rs b/src/bit_encoding/encode.rs index e7cff1f1..a5d64c91 100644 --- a/src/bit_encoding/encode.rs +++ b/src/bit_encoding/encode.rs @@ -310,7 +310,10 @@ pub fn encode_natural(mut n: usize, w: &mut BitWriter) -> io::R 8 * mem::size_of::() - n.leading_zeros() as usize - 1 } - let mut suffix = Vec::new(); + // The u64 cast below already assumes usize is at most 64 bits. + // The largest suffix lengths are therefore 63, 5, 2, 1. + let mut suffix = [(0, 0); 4]; + let mut suffix_len = 0; loop { debug_assert!(n > 0); @@ -320,12 +323,13 @@ pub fn encode_natural(mut n: usize, w: &mut BitWriter) -> io::R break; } else { w.write_bit(true)?; - suffix.push((n, len)); + suffix[suffix_len] = (n, len); + suffix_len += 1; n = len; } } - while let Some((bits, len)) = suffix.pop() { + for &(bits, len) in suffix[..suffix_len].iter().rev() { let bits = bits as u64; // Case safety: assuming 64-bit machine or lower w.write_bits_be(bits, len)?; } @@ -352,4 +356,39 @@ mod test { assert_eq!(n, m); } } + + #[test] + fn encode_natural_boundaries() { + let mut values = vec![1, usize::MAX]; + for power in 1..usize::BITS { + let n = 1usize << power; + values.extend([n - 1, n, n + 1]); + } + for n in values { + // Construct the expected bit string recursively, independently of + // the encoder's suffix storage. The decoder only supports u32. + fn expected(n: usize) -> String { + if n == 1 { + return "0".to_owned(); + } + let bits = format!("{n:b}"); + format!("1{}{}", expected(bits.len() - 1), &bits[1..]) + } + + let expected = expected(n); + for offset in 0..8 { + let mut bytes = Vec::new(); + let mut w = BitWriter::new(&mut bytes); + w.write_bits_be(0xa5, offset).unwrap(); + assert_eq!(encode_natural(n, &mut w).unwrap(), expected.len()); + w.flush_all().unwrap(); + let bits: String = + BitIter::byte_slice_window(&bytes, offset, offset + expected.len()) + .take(expected.len()) + .map(|bit| if bit { '1' } else { '0' }) + .collect(); + assert_eq!(bits, expected, "n={n}, offset={offset}"); + } + } + } }