perf(encode): trim dfast tail seeding and the optimal parser's per-node work

zstd/src/encoding/bt/mod.rs

@@ -431,31 +431,32 @@ impl BtMatcher {
     /// `encode_offset_with_reps` minus the rep-history update — used in
     /// the optimal parser's per-candidate price probe where the rep
     /// window hasn't been committed yet.
+    ///
+    /// Selects rather than branches: it runs for every candidate the search
+    /// returns, and whether an offset repeats one of three histories is data
+    /// the predictor cannot learn, so a compare chain mispredicted on a large
+    /// share of candidates. The three repeat values are picked by `lit_len`
+    /// (RFC 8878 3.1.2.5: with no literals the codes shift to `rep[1]`,
+    /// `rep[2]`, `rep[0] - 1`), then folded last to first so the first match
+    /// wins, as the chain did.
     #[inline(always)]
     pub(crate) fn encode_offset_base_with_reps(
         actual_offset: u32,
         lit_len: usize,
         reps: [u32; 3],
     ) -> u32 {
-        if lit_len > 0 {
-            if actual_offset == reps[0] {
-                1
-            } else if actual_offset == reps[1] {
-                2
-            } else if actual_offset == reps[2] {
-                3
-            } else {
-                actual_offset.saturating_add(3)
-            }
-        } else if actual_offset == reps[1] {
-            1
-        } else if actual_offset == reps[2] {
-            2
-        } else if reps[0] > 1 && actual_offset == reps[0] - 1 {
-            3
-        } else {
-            actual_offset.saturating_add(3)
-        }
+        use core::hint::select_unpredictable;
+        let ll0 = lit_len == 0;
+        // `rep[0] - 1` is a repeat code only when it is a real offset; 0
+        // never equals a match offset, which is at least 1.
+        let rep0_minus_one = select_unpredictable(reps[0] > 1, reps[0].wrapping_sub(1), 0);
+        let first = select_unpredictable(ll0, reps[1], reps[0]);
+        let second = select_unpredictable(ll0, reps[2], reps[1]);
+        let third = select_unpredictable(ll0, rep0_minus_one, reps[2]);
+        let code = actual_offset.saturating_add(3);
+        let code = select_unpredictable(actual_offset == third, 3, code);
+        let code = select_unpredictable(actual_offset == second, 2, code);
+        select_unpredictable(actual_offset == first, 1, code)
     }
 
     /// Upstream zstd parity: `ZSTD_updateStats` for one sequence the parser

zstd/src/encoding/bt/tests.rs

@@ -23,3 +23,49 @@ fn estimated_workspace_bytes_matches_retained_scratch_at_growth_bounds() {
         core::mem::size_of::<BtMatcher>() + bt.heap_size()
     );
 }
+
+/// The repeat-code probe selects instead of branching; it must return what
+/// the RFC 8878 3.1.2.5 compare chain returns for every offset and history,
+/// including colliding repeats (the first match wins), `rep[0] <= 1` with no
+/// literals (no `rep[0] - 1` code), and an offset at the saturation edge.
+#[test]
+fn offset_base_probe_matches_the_repeat_compare_chain() {
+    fn chain(offset: u32, lit_len: usize, reps: [u32; 3]) -> u32 {
+        if lit_len > 0 {
+            if offset == reps[0] {
+                1
+            } else if offset == reps[1] {
+                2
+            } else if offset == reps[2] {
+                3
+            } else {
+                offset.saturating_add(3)
+            }
+        } else if offset == reps[1] {
+            1
+        } else if offset == reps[2] {
+            2
+        } else if reps[0] > 1 && offset == reps[0] - 1 {
+            3
+        } else {
+            offset.saturating_add(3)
+        }
+    }
+    let values = [0u32, 1, 2, 3, 4, 5, u32::MAX - 3, u32::MAX - 2, u32::MAX];
+    for &r0 in &values {
+        for &r1 in &values {
+            for &r2 in &values {
+                for &offset in &values {
+                    for lit_len in [0usize, 1, 7] {
+                        let reps = [r0, r1, r2];
+                        assert_eq!(
+                            BtMatcher::encode_offset_base_with_reps(offset, lit_len, reps),
+                            chain(offset, lit_len, reps),
+                            "offset={offset} lit_len={lit_len} reps={reps:?}"
+                        );
+                    }
+                }
+            }
+        }
+    }
+}

zstd/src/encoding/dfast/mod.rs

@@ -1967,14 +1967,22 @@ impl DfastMatchGenerator {
         pos
     }
 
+    /// Index the starts of the block's dense tail that the scan did not pass,
+    /// `pos` being where it stopped, so the next block can match into them.
     pub(crate) fn seed_remaining_hashable_starts(
         &mut self,
         current_abs_start: usize,
         current_len: usize,
         pos: usize,
     ) {
         let boundary_tail_start = current_len.saturating_sub(Self::BOUNDARY_DENSE_TAIL_LEN);
-        let mut seed_pos = pos.min(current_len).min(boundary_tail_start);
+        // Never the gap the scan's step skipped before it stopped, nor what a
+        // match ending in the tail covered: upstream's loops index neither,
+        // and on incompressible input that gap is the step, hundreds of bytes
+        // by a block's end. Seeding it cost 11% of the wall-clock time of a
+        // 100 KiB random frame at level 3 (1.18x -> 1.05x of libzstd in the
+        // same run) for 43 bytes over 3.1 MB of corpus at levels 3-4.
+        let mut seed_pos = pos.min(current_len).max(boundary_tail_start);
         // Every start with a whole short key; past it nothing is hashable.
         while seed_pos + self.mls <= current_len {
             self.insert_position(current_abs_start + seed_pos);

zstd/src/encoding/hc/generator.rs

@@ -103,13 +103,13 @@ macro_rules! bt_insert_range_body {
         // The walk carries one coordinate, the stored index, as the collect body
         // does and as upstream carries `matchIndex`: the window bound becomes one
         // unsigned range test on it (HC_EMPTY, 0, decodes below the window and
-        // ends the walk), and the absolute position, the history index and the
-        // pair slot are each one add of a bias taken here, instead of reloading
-        // `position_base` / `index_shift` / `history_abs_start` through the
-        // table on every node. See the collect body for the derivation.
+        // ends the walk), the absolute position and the history index are each
+        // one add of a bias taken here, instead of reloading `position_base` /
+        // `index_shift` / `history_abs_start` through the table on every node,
+        // and the pair slot is the index itself, masked. See the collect body
+        // for the derivation.
         let abs_bias = coords.abs_bias;
         let idx_bias = coords.idx_bias;
-        let bt_bias = coords.bt_bias;
         // The window floor in stored space, fixed for the run: a node is in the
         // window when its stored index is at or above it, one compare (upstream
         // `matchIndex >= windowLow`). It is at least 1, so the empty slot 0
@@ -218,10 +218,9 @@ macro_rules! bt_insert_range_body {
                 // floor points at `closed_side`, so the final stores need no test
                 // (upstream `dummy32`).
                 let mut closed_side: u32 = 0;
-                // SAFETY: `2*((stored + bt_bias) & bt_mask) + 1 < chain_table.len()`
-                // (`stored + bt_bias == abs_pos + index_shift`, the pair index of
-                // `MatchTable::bt_pair_index_for_abs`).
-                let pair = unsafe { chain_ptr.add(2 * (stored.wrapping_add(bt_bias) & bt_mask)) };
+                // SAFETY: `2*(stored & bt_mask) + 1 < chain_table.len()`, the pair
+                // index of `MatchTable::bt_pair_index_for_stored`.
+                let pair = unsafe { chain_ptr.add(2 * (stored & bt_mask)) };
                 let mut smaller_slot: *mut u32 = pair;
                 let mut larger_slot: *mut u32 = unsafe { pair.add(1) };
                 // SAFETY: `hash` is masked to `hash_log` bits and the table is
@@ -244,8 +243,7 @@ macro_rules! bt_insert_range_body {
                     // SAFETY: the pair index is below `chain_table.len() - 1`;
                     // `chain_ptr` is the hoisted live base, not reallocated during
                     // the walk.
-                    let next_pair =
-                        unsafe { chain_ptr.add(2 * (node_stored.wrapping_add(bt_bias) & bt_mask)) };
+                    let next_pair = unsafe { chain_ptr.add(2 * (node_stored & bt_mask)) };
                     let seed_len = common_length_smaller.min(common_length_larger);
                     // Each side's common length is a match the walk did not stop
                     // at, so it is below the tail limit.
@@ -908,12 +906,11 @@ macro_rules! bt_insert_and_collect_matches_body {
         // A node is reached through its stored index alone, as upstream walks
         // one `matchIndex` (`match = base + matchIndex`, slot
         // `2*(matchIndex & btMask)`): its bytes sit at `cand_base + stored`
-        // and its pair at `(stored + bt_bias) & bt_mask`, and nothing decodes
-        // it to an absolute position unless it joins the ladder.
+        // and its pair at `stored & bt_mask`, and nothing decodes it to an
+        // absolute position unless it joins the ladder.
         // `cand_base` may lie outside the history; only the candidate it is
         // offset to is ever read, and the window gate proves that one inside.
         debug_assert!(win_low_stored >= 1);
-        let bt_bias = coords.bt_bias;
         // SAFETY: `idx + 8 <= concat.len()`, which the search's entry checks.
         let cur_ptr = unsafe { cbase.add(idx) };
         let cand_base = cbase.wrapping_add(coords.idx_bias);
@@ -926,10 +923,8 @@ macro_rules! bt_insert_and_collect_matches_body {
         let mut compares_left = $search_depth;
         let mut common_length_smaller = 0usize;
         let mut common_length_larger = 0usize;
-        // `bt_pair_index_for_abs(abs_pos)`: `stored + bt_bias` is
-        // `abs_pos + index_shift`.
-        let pair_idx = 2 * ((stored as usize).wrapping_add(bt_bias) & bt_mask);
-        debug_assert_eq!(pair_idx, $table.bt_pair_index_for_abs($abs_pos));
+        let pair_idx = 2 * (stored as usize & bt_mask);
+        debug_assert_eq!(pair_idx, $table.bt_pair_index_for_stored(stored as usize));
         // The two open sides of the tree, as pointers into the pair table
         // (upstream `smallerPtr` / `largerPtr`). A side the walk closes at the
         // tree floor is pointed at `closed_side` instead, so the final store
@@ -968,22 +963,10 @@ macro_rules! bt_insert_and_collect_matches_body {
                 node_stored <= bt_low_stored,
                 node_stored.wrapping_add(coords.abs_bias) <= bt_low,
             );
-            // `2*(candidate_abs + index_shift & bt_mask)` with `index_shift`
-            // folded away: `candidate_abs + index_shift == stored + bt_bias`.
             // SAFETY: the pair index is below `chain_table.len() - 1`;
             // `chain_ptr` is the hoisted live base, not reallocated during the
             // walk.
-            let next_pair = unsafe { chain_ptr.add(2 * (node_stored.wrapping_add(bt_bias) & bt_mask)) };
-            // Both children are read ahead of the compare, where upstream reads
-            // only the side it descends, so the next node's load can overlap
-            // the compare. Reading only the chosen side after it measured about
-            // 1% slower on z000033 10 KiB at level 19 with output unchanged,
-            // against a flat level 3 control; that is below what two builds
-            // resolve, so neither order is shown faster and the walk keeps its
-            // existing one.
-            // SAFETY: as for `next_pair`.
-            let next_smaller = unsafe { *next_pair };
-            let next_larger = unsafe { *next_pair.add(1) };
+            let next_pair = unsafe { chain_ptr.add(2 * (node_stored & bt_mask)) };
             let seed_len = common_length_smaller.min(common_length_larger);
             // Every side's common length is a match the walk did not stop at,
             // so it is below the tail limit and the compare below has bytes
@@ -1034,12 +1017,15 @@ macro_rules! bt_insert_and_collect_matches_body {
                 }
             }
 
-            if match_len >= tail_limit {
-                break;
-            }
+            // The walk enters with `best_len < tail_limit` (a repeat or hash3
+            // match reaching the block end skips the walk), and a new best that
+            // reaches it breaks above, so a match that is not a new best is
+            // shorter than `best_len` and below the limit. Upstream tests the
+            // block end only on a new best (`ip+matchLength == iLimit`).
+            debug_assert!(match_len < tail_limit);
 
             // SAFETY: the first differing bytes after a `match_len`-long
-            // prefix; `match_len < tail_limit` (break above) keeps both in
+            // prefix; `match_len < tail_limit` (asserted above) keeps both in
             // range.
             if unsafe { *cand_ptr.add(match_len) < *cur_ptr.add(match_len) } {
                 // SAFETY: an open side points at a pair slot (initial pair or
@@ -1050,9 +1036,11 @@ macro_rules! bt_insert_and_collect_matches_body {
                     smaller_slot = &raw mut closed_side;
                     break;
                 }
+                // Only the side descended is read, after the compare, as
+                // upstream does (`matchIndex = nextPtr[1]`).
                 // SAFETY: as for `next_pair`.
                 smaller_slot = unsafe { next_pair.add(1) };
-                match_stored = next_larger;
+                match_stored = unsafe { *smaller_slot };
             } else {
                 // SAFETY: as above for `larger_slot`.
                 unsafe { *larger_slot = match_stored };
@@ -1062,7 +1050,8 @@ macro_rules! bt_insert_and_collect_matches_body {
                     break;
                 }
                 larger_slot = next_pair;
-                match_stored = next_smaller;
+                // SAFETY: as for `next_pair`.
+                match_stored = unsafe { *larger_slot };
             }
         }
 

zstd/src/encoding/match_generator/tests.rs

@@ -5115,34 +5115,55 @@ fn dfast_seed_remaining_hashable_starts_seeds_last_short_hash_positions() {
     );
 }
 
-#[test]
-fn dfast_seed_remaining_hashable_starts_handles_pos_at_block_end() {
-    let mut matcher = DfastMatchGenerator::new(1 << 20);
-    let block = deterministic_high_entropy_bytes(0x7BB2_DA91_441E_C0EF, 64);
-    matcher.commit_input(&block);
-    matcher.ensure_hash_tables();
-
-    let current_len = matcher.window_blocks.back().copied().unwrap_or(0);
-    let current_abs_start = matcher.history_abs_start + matcher.window_size - current_len;
-    matcher.seed_remaining_hashable_starts(current_abs_start, current_len, current_len);
+/// The tail seed indexes only the dense tail past where the scan stopped:
+/// neither the gap a stepping scan skipped before the tail, nor tail starts a
+/// match ending at the block's end already covered, as upstream's loops leave
+/// both. Seeding the gap cost 11% of a 100 KiB random frame's wall-clock time
+/// at level 3.
+#[test]
+fn dfast_seed_remaining_hashable_starts_skips_the_gap_and_covered_starts() {
+    let short_slot_of = |matcher: &DfastMatchGenerator, abs_pos: usize| {
+        let rel = abs_pos - matcher.history_abs_start;
+        let live = matcher.live_history();
+        assert!(
+            rel + 5 <= live.len(),
+            "fixture must leave the start short-hashable"
+        );
+        matcher.tables[matcher.long_len() + matcher.short_hash_index(&live[rel..])]
+    };
+    let fresh = || {
+        let mut matcher = DfastMatchGenerator::new(1 << 20);
+        let block = deterministic_high_entropy_bytes(0x7BB2_DA91_441E_C0EF, 64);
+        matcher.commit_input(&block);
+        matcher.ensure_hash_tables();
+        let current_len = matcher.window_blocks.back().copied().unwrap_or(0);
+        let current_abs_start = matcher.history_abs_start + matcher.window_size - current_len;
+        (matcher, current_abs_start, current_len)
+    };
 
-    let target_abs_pos = current_abs_start + current_len - 5;
-    let target_rel = target_abs_pos - matcher.history_abs_start;
-    let live = matcher.live_history();
-    assert!(
-        target_rel + 5 <= live.len(),
-        "fixture must leave the last short-hash start valid"
+    // The scan stopped well before the tail: the tail is seeded, the gap is not.
+    let (mut matcher, start, len) = fresh();
+    matcher.seed_remaining_hashable_starts(start, len, 10);
+    let last = start + len - 5;
+    assert_eq!(
+        short_slot_of(&matcher, last),
+        matcher.pack_slot(last),
+        "the tail past an early stop is seeded"
     );
-    let short_hash = matcher.short_hash_index(&live[target_rel..]);
-    let target_slot = matcher.pack_slot(target_abs_pos);
+    let gap = start + len - DfastMatchGenerator::BOUNDARY_DENSE_TAIL_LEN - 20;
     assert_ne!(
-        target_slot, DFAST_EMPTY_SLOT,
-        "pack_slot must never return the empty-slot sentinel for a real position"
+        short_slot_of(&matcher, gap),
+        matcher.pack_slot(gap),
+        "the gap before the dense tail stays unindexed"
     );
-    assert_eq!(
-        matcher.tables[matcher.long_len() + short_hash],
-        target_slot,
-        "tail seeding must still include the last 5-byte-hashable start when pos is at block end"
+
+    // A match ran to the block's end: nothing is left to seed.
+    let (mut matcher, start, len) = fresh();
+    matcher.seed_remaining_hashable_starts(start, len, len);
+    assert_ne!(
+        short_slot_of(&matcher, last),
+        matcher.pack_slot(last),
+        "starts a match covered are not seeded"
     );
 }