perf(encode): keep the level-1 bit writers' flushes and FSE states off the critical path

zstd/src/bit_io/bit_writer.rs

@@ -372,29 +372,25 @@ impl<V: AsMut<Vec<u8>>> BitBurst<'_, V> {
     ///
     /// # Safety
     ///
-    /// The accumulator must have room: `bits_in_partial + num_bits <= 64`, and
-    /// `bits` must be clean above `num_bits` (`num_bits == 64 || bits >>
-    /// num_bits == 0`). Neither violation touches memory, but either corrupts
-    /// the stream with nothing to detect it, so the caller proves both from
-    /// its bit budget. Debug builds assert them.
+    /// The accumulator must keep a bit free: `bits_in_partial + num_bits < 64`
+    /// (upstream's `assert(nbBits + bitPos < 64)`), and `bits` must be clean
+    /// above `num_bits` (`bits >> num_bits == 0`). Neither violation touches
+    /// memory, but either corrupts the stream with nothing to detect it, so the
+    /// caller proves both from its bit budget. Debug builds assert them.
     #[inline(always)]
     pub(crate) unsafe fn add(&mut self, bits: u64, num_bits: usize) {
         debug_assert!(
-            self.bits_in_partial + num_bits <= 64,
+            self.bits_in_partial + num_bits < 64,
             "{num_bits} bits do not fit beside the {} held",
             self.bits_in_partial,
         );
         debug_assert!(
-            num_bits == 64 || bits >> num_bits == 0,
+            bits >> num_bits == 0,
             "value has dirty high bits beyond num_bits={num_bits}",
         );
-        // Masked so a full accumulator cannot make this a shift by 64, which is
-        // undefined for `u64`. x86 and AArch64 shifts mask the count
-        // themselves, so the mask costs nothing, where a `num_bits == 0` early
-        // return was a branch on every add. The accumulator is full only when
-        // `num_bits == 0`, and the second precondition then makes `bits == 0`,
-        // so the OR adds nothing whatever the count.
-        self.partial |= bits << (self.bits_in_partial & 63);
+        // `bits_in_partial < 64` by the first precondition, so the shift is in
+        // range.
+        self.partial |= bits << self.bits_in_partial;
         self.bits_in_partial += num_bits;
     }
 
@@ -423,13 +419,9 @@ impl<V: AsMut<Vec<u8>>> BitBurst<'_, V> {
                 .write_unaligned(self.partial.to_le());
             self.cursor = self.cursor.add(nb_bytes);
         }
-        // A full accumulator (64 bits) leaves nothing: shifting a `u64` by 64
-        // is undefined, so it is zeroed instead.
-        self.partial = if nb_bytes == 8 {
-            0
-        } else {
-            self.partial >> (nb_bytes * 8)
-        };
+        // `add` keeps the accumulator under 64 bits, so at most 7 bytes leave
+        // and the shift is in range (upstream `BIT_flushBitsFast`, unguarded).
+        self.partial >>= nb_bytes * 8;
         self.bits_in_partial &= 7;
     }
 

zstd/src/bit_io/bit_writer/tests.rs

@@ -160,28 +160,41 @@ fn add_multiple_aligned() {
     assert_eq!(vec![0xFF, 0xF0, 0x0F, 0x00], bw.dump());
 }
 
-/// A burst flushing a full accumulator writes all eight bytes and holds
-/// nothing after, so the next bits land on a clean byte rather than on bits
-/// of the flushed value.
+/// A burst flushing the fullest accumulator it allows (63 bits) writes the
+/// seven whole bytes, keeps the last 7 bits, and lands the next bits right
+/// after them rather than on bits of the flushed value.
 #[test]
-fn a_burst_flushes_a_full_accumulator_clean() {
+fn a_burst_flushes_63_bits_and_keeps_the_partial_byte() {
     let mut bw = BitWriter::new();
     let mut burst = bw.burst(16);
-    // SAFETY: 64 bits into an empty accumulator, then 8 into an empty one,
+    // SAFETY: 63 bits into an empty accumulator, then 1 beside the 7 kept,
     // within the 16 bytes reserved.
     unsafe {
-        burst.add(0x0123_4567_89AB_CDEF, 64);
+        burst.add(0x7123_4567_89AB_CDEF, 63);
         burst.flush();
-        burst.add(0xA5, 8);
+        burst.add(0x1, 1);
         burst.flush();
     }
     burst.finish();
     assert_eq!(
         bw.dump(),
-        vec![0xEF, 0xCD, 0xAB, 0x89, 0x67, 0x45, 0x23, 0x01, 0xA5]
+        vec![0xEF, 0xCD, 0xAB, 0x89, 0x67, 0x45, 0x23, 0xF1]
     );
 }
 
+/// An add that would fill all 64 bits breaks the burst's contract (the flush
+/// shifts by the whole bytes held, which must stay under 8); debug builds
+/// catch it at the add.
+#[test]
+#[cfg(debug_assertions)]
+#[should_panic(expected = "do not fit")]
+fn a_burst_refuses_a_full_accumulator() {
+    let mut bw = BitWriter::new();
+    let mut burst = bw.burst(16);
+    // SAFETY: the debug assertion fires before anything is written.
+    unsafe { burst.add(0x0123_4567_89AB_CDEF, 64) };
+}
+
 // #[test]
 // fn catches_more_than_in_buf() {
 //     todo!();

zstd/src/encoding/blocks/compressed.rs

@@ -2000,10 +2000,10 @@ fn fill_and_count<const OFFSETS_CODED: bool>(
         );
         // SAFETY: the codes are at most 35 (LL), 52 (ML) and 31 (OF), under the
         // histogram's `SEQUENCE_CODE_SLOTS`. The LL and ML bounds hold in release
-        // because `encode_literal_length` and `encode_match_len` index
-        // `LL_EXTRA_BITS` (36 entries) and `ML_EXTRA_BITS` (53) with a bounds
-        // check, so a larger code panics there before reaching this write; the
-        // OF code is `ilog2` of a `u32`. Those lookups must stay checked.
+        // because `encode_literal_length` and `encode_match_len` assert the
+        // length below 2^17 before deriving a code from its high bit, so a larger
+        // code panics there before reaching this write; the OF code is `ilog2`
+        // of a `u32`. Those asserts must stay.
         unsafe {
             *ll_counts.get_unchecked_mut(usize::from(ll_code)) += 1;
             *ml_counts.get_unchecked_mut(usize::from(ml_code)) += 1;
@@ -2889,6 +2889,18 @@ macro_rules! encode_sequences_body {
         // body is what leaves the accumulator with under a byte for the next
         // round.
         if sequences.len() > 1 {
+            // Streams 0 / 1 / 2 are LL / ML / OF; built only for a block that
+            // runs the loop. A local on purpose, not scratch kept on the
+            // compressor: addressed off the stack pointer it needs no base
+            // register, and the loop keeps two of the three states in
+            // registers. Held behind a pointer the same loop spilled one state
+            // more, +0.4..0.6% instructions at level 1 against the per-block
+            // clearing this costs.
+            let mut view = crate::fse::fse_encoder::SequenceStateTables::new();
+            view.set(0, ll_table);
+            view.set(1, ml_table);
+            view.set(2, of_table);
+            let view = &view;
             // At most 63 bits a sequence leave as whole bytes, under 8 bytes;
             // 12 a sequence and 64 over covers that and the up to 7 bytes the
             // burst drains from the init-state writes above as it starts,
@@ -2932,8 +2944,8 @@ macro_rules! encode_sequences_body {
                 // derivation assumed.
                 #[cfg(debug_assertions)]
                 let mut state_diff_bits = 0usize;
-                if let Some(table) = of_table {
-                    let next = table.next_state(of_code, of_state);
+                {
+                    let next = view.next_state::<2>(of_code, of_state);
                     let diff = crate::fse::fse_encoder::transition_bits(of_state, next.num_bits);
                     unsafe {
                         burst.add(diff as u64, next.num_bits as usize);
@@ -2944,8 +2956,8 @@ macro_rules! encode_sequences_body {
                     }
                     of_state = next.index;
                 }
-                if let Some(table) = ml_table {
-                    let next = table.next_state(ml_code, ml_state);
+                {
+                    let next = view.next_state::<1>(ml_code, ml_state);
                     let diff = crate::fse::fse_encoder::transition_bits(ml_state, next.num_bits);
                     unsafe {
                         burst.add(diff as u64, next.num_bits as usize);
@@ -2956,8 +2968,8 @@ macro_rules! encode_sequences_body {
                     }
                     ml_state = next.index;
                 }
-                if let Some(table) = ll_table {
-                    let next = table.next_state(ll_code, ll_state);
+                {
+                    let next = view.next_state::<0>(ll_code, ll_state);
                     let diff = crate::fse::fse_encoder::transition_bits(ll_state, next.num_bits);
                     unsafe {
                         burst.add(diff as u64, next.num_bits as usize);
@@ -3002,10 +3014,10 @@ macro_rules! encode_sequences_body {
                     // the leftover, the diffs and all three extra fields have
                     // to fit.
                     debug_assert!(
-                        extra_bits_total >= 31 || 7 + state_diff_bits + extra_bits_total <= 64,
+                        extra_bits_total >= 31 || 7 + state_diff_bits + extra_bits_total < 64,
                         "no flush at {extra_bits_total} extra bits, but 7 leftover + \
-                         {state_diff_bits} diff bits + {extra_bits_total} would overflow the \
-                         64-bit accumulator",
+                         {state_diff_bits} diff bits + {extra_bits_total} would fill the \
+                         accumulator, which keeps a bit free",
                     );
                 }
                 if extra_bits_total >= 31 {
@@ -3046,6 +3058,7 @@ macro_rules! encode_sequences_body {
             }
             burst.finish();
         }
+        // An RLE stream has no state to write.
         if let Some(table) = ml_table {
             writer.write_bits(ml_state as u64, table.table_size.ilog2() as usize);
         }
@@ -3178,6 +3191,20 @@ const ML_EXTRA_BITS: [u8; 53] = [
     1, 1, 1, 1, 2, 2, 3, 3, 4, 4, 5, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
 ];
 
+/// A short length's code in the low byte and its extra-bit width in the high
+/// byte, so one load answers both.
+const fn fuse_code_bits<const N: usize>(codes: &[u8; N], bits: &[u8]) -> [u16; N] {
+    let mut fused = [0u16; N];
+    let mut i = 0;
+    while i < N {
+        fused[i] = codes[i] as u16 | (bits[codes[i] as usize] as u16) << 8;
+        i += 1;
+    }
+    fused
+}
+const LL_CODE_BITS: [u16; 64] = fuse_code_bits(&LL_CODE, &LL_EXTRA_BITS);
+const ML_CODE_BITS: [u16; 128] = fuse_code_bits(&ML_CODE, &ML_EXTRA_BITS);
+
 /// Split a literal-length into its FSE symbol, the extra bits to append, and
 /// how many of them (upstream `ZSTD_LLcode` plus the `LL_bits` lookup its
 /// caller does).
@@ -3186,15 +3213,22 @@ const ML_EXTRA_BITS: [u8; 53] = [
 /// compiled to a chain of comparisons; upstream reaches the same answer with a
 /// table index. Every code's baseline is a multiple of its own extra-bit width,
 /// so masking off those bits is the same subtraction the ranges spelled out.
+/// From 64 up a code carries exactly `ilog2(len)` extra bits, so that branch
+/// needs no table at all.
 #[inline]
 pub(crate) fn encode_literal_length(len: u32) -> (u8, u32, usize) {
-    debug_assert!(len < 131_072, "literal length {len} out of encodable range");
-    let code = if len < 64 {
-        LL_CODE[len as usize]
+    let (code, bits) = if len < 64 {
+        let fused = LL_CODE_BITS[len as usize];
+        (fused as u8, (fused >> 8) as usize)
     } else {
-        (len.ilog2() + LL_DELTA_CODE) as u8
+        // RFC 8878 3.1.1.3.2.1.1: literal-length code 35 is the last, its
+        // baseline 65536 with 16 extra bits.
+        // No value in the message: formatting it keeps `len` on the stack
+        // across the whole calling loop.
+        assert!(len < 131_072, "literal length out of encodable range");
+        let log = len.ilog2();
+        ((log + LL_DELTA_CODE) as u8, log as usize)
     };
-    let bits = LL_EXTRA_BITS[code as usize] as usize;
     (code, len & ((1u32 << bits) - 1), bits)
 }
 
@@ -3205,17 +3239,18 @@ pub(crate) fn encode_literal_length(len: u32) -> (u8, u32, usize) {
 /// Table-driven for the same reason as [`encode_literal_length`].
 #[inline]
 pub(crate) fn encode_match_len(len: u32) -> (u8, u32, usize) {
-    debug_assert!(
-        (3..131_075).contains(&len),
-        "match length {len} out of encodable range",
-    );
+    debug_assert!(len >= 3, "match length {len} out of encodable range");
     let base = len - 3;
-    let code = if base < 128 {
-        ML_CODE[base as usize]
+    let (code, bits) = if base < 128 {
+        let fused = ML_CODE_BITS[base as usize];
+        (fused as u8, (fused >> 8) as usize)
     } else {
-        (base.ilog2() + ML_DELTA_CODE) as u8
+        // RFC 8878 3.1.1.3.2.1.1: match-length code 52 is the last, its
+        // baseline 65539 with 16 extra bits.
+        assert!(base < 131_072, "match length out of encodable range");
+        let log = base.ilog2();
+        ((log + ML_DELTA_CODE) as u8, log as usize)
     };
-    let bits = ML_EXTRA_BITS[code as usize] as usize;
     (code, base & ((1u32 << bits) - 1), bits)
 }
 

zstd/src/encoding/blocks/compressed/literal_runs.rs

@@ -92,21 +92,29 @@ macro_rules! gather_body {
         let out = unsafe { dst.as_mut_ptr().add(start) };
         let mut written = 0usize;
         let mut pos = 0usize;
+        // A run's length is a `u32`, but the tail is a matcher's unchecked sum;
+        // bounding it by the block keeps the single-compare room test below
+        // from wrapping. Once per block.
+        assert!(
+            $tail <= src_len,
+            "sequences reach past the block they describe"
+        );
         macro_rules! copy_run {
             ($len:expr) => {{
                 let len = $len;
-                // `pos <= src_len` holds on every iteration.
-                let room = src_len - pos;
                 // SAFETY (both arms): `written <= pos`, and the capacity check in
                 // `gather_literals` gives the buffer `src_len` bytes past
                 // `start`. The chunked arm reads and writes at most
-                // `max($chunk, len + $chunk - 1) <= len + WILDCOPY_SLACK <= room`
-                // bytes from `pos` / `written` (`$chunk <= WILDCOPY_SLACK`); the
-                // exact arm, `len <= room` bytes.
+                // `max($chunk, len + $chunk - 1) <= len + WILDCOPY_SLACK` bytes
+                // from `pos` / `written` (`$chunk <= WILDCOPY_SLACK`), inside the
+                // block by its guard; the exact arm, `len <= src_len - pos`.
                 unsafe {
                     let s = src.add(pos);
                     let d = out.add(written);
-                    if room >= WILDCOPY_SLACK && len <= room - WILDCOPY_SLACK {
+                    // One compare in 64 bits, where no operand can overflow
+                    // (a run is a `u32`, `pos` and the tail are bounded by the
+                    // block), rather than a room test and a length test.
+                    if pos as u64 + len as u64 + WILDCOPY_SLACK as u64 <= src_len as u64 {
                         // Upstream's shape (`ZSTD_wildcopy`): one chunk
                         // unconditionally, which is the whole run for most of
                         // them, and a loop only for the rest of a longer one.
@@ -126,7 +134,11 @@ macro_rules! gather_body {
                             core::ptr::copy_nonoverlapping(s, d, len);
                         }
                     } else {
-                        assert!(len <= room, "sequences reach past the block they describe");
+                        // `pos <= src_len` holds on every iteration.
+                        assert!(
+                            len <= src_len - pos,
+                            "sequences reach past the block they describe"
+                        );
                         if len <= 32 {
                             if len != 0 {
                                 simd_copy::copy_exact_small(s, d, len);

zstd/src/encoding/blocks/compressed/tests.rs

@@ -1246,6 +1246,21 @@ fn sequences_past_the_block_are_refused() {
     );
 }
 
+/// A trailing literal count past the block is refused before any copy, however
+/// large: a custom matcher reports it unchecked, and near `usize::MAX` the
+/// chunked copy's room test would wrap and let a copy of that length through.
+#[test]
+#[should_panic(expected = "sequences reach past the block they describe")]
+fn a_tail_past_the_block_is_refused_before_the_copy() {
+    let (_ws, mut literals) = literals_buffer(64);
+    super::literal_runs::gather_literals::<crate::encoding::fastpath::ScalarTier>(
+        &[0x5A; 64],
+        &[],
+        usize::MAX - 16,
+        &mut literals,
+    );
+}
+
 /// Every literal-gather tier this CPU can run, not only the one it selects, so
 /// a narrower tier's path is exercised on a wider machine.
 fn runnable_fastpath_kernels() -> Vec<crate::encoding::fastpath::FastpathKernel> {

zstd/src/encoding/frame_compressor.rs

@@ -889,7 +889,9 @@ fn presplit_record_fingerprint(
 /// 2-byte hashing for `hashLog >= 9`; this helper exists so the borders
 /// heuristic doesn't pay for that wider hash on its 512-byte windows.
 fn presplit_record_byte_histogram(fp: &mut PreSplitFingerprint, src: &[u8]) {
-    fp.events.fill(0);
+    // Only the byte slots: a `hash_log == 8` fingerprint is recorded, compared
+    // and read in those alone.
+    fp.events[..256].fill(0);
     for &b in src {
         fp.events[b as usize] += 1;
     }