From acf7ca1139295452a7e07a22aee320306effa806 Mon Sep 17 00:00:00 2001 From: Tommy Volk Date: Sun, 25 Oct 2026 18:12:30 -0500 Subject: [PATCH] perf(encoding): keep natural-number suffixes on the stack Replace the temporary vector with four stack entries, the maximum needed by the existing encoder on targets up to 64 bits. Check exact encodings around every power of two and at usize::MAX at all starting bit offsets. Written by OpenAI Codex. --- src/bit_encoding/encode.rs | 45 +++++++++++++++++++++++++++++++++++--- 1 file changed, 42 insertions(+), 3 deletions(-) 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}"); + } + } + } }