perf(encode): cap HC/BT history mirror near the live window (#421)

* perf(encode): cap HC/BT history mirror near the live window

Mirror the row (#418) and dfast (#420) history-buffer fix for the shared
HashChain / binary-tree MatchTable (btlazy2 / btopt / btultra, L13-L22).
On long streams the history mirror power-of-two doubled to ~2x the window
(an 8 MiB window peaked at 16 MiB), inflating compress peak memory. Reserve
window + window/4 + one block once eviction starts and drain the dead
prefix at a quarter window so the Vec grows linearly and never reallocates
again; reserve_history's ceiling is bumped to match. Small frames keep
their tight buffer. Byte-identical (buffer management only).

* refactor(encode): plain arithmetic for history cap, not saturating_add

The reserve_history ceiling used saturating_add, which would silently
clamp on overflow and mask an upstream window_log bound violation. The
sum is max_window_size + max_window_size/4 + MAX_BLOCK_SIZE with
max_window_size = 1 << window_log and window_log <= 31, so it is at most
~2^31 + 2^29 + 2^17 < usize::MAX even on 32-bit targets — overflow is
unreachable. Use plain arithmetic with a bound comment, matching the
add_data eviction reserve.

* refactor: drop defensive saturating arithmetic where overflow is unreachable

Codebase-wide audit of saturating_add/saturating_mul against the
safe-arithmetic policy (saturating only for a deliberate min/max clamp,
plain arithmetic when the bound provably holds, so a real overflow fails
loudly at its cause instead of being silently clamped).

Converted to plain arithmetic (with a bound comment each) where overflow
is unreachable: huffman tree node counts and literal/symbol frequency
counters (<= block size), pre-split fingerprint distance/threshold and
per-block event merge, FSE/HUF cost estimates, dict+window log helper
(matches upstream's plain u64 add), FASTCOVER epoch sizing, history
pre-size hint, fast-matcher length/2x-window checks, eviction byte
accounting.

Kept (documented as deliberate, not masking): compress_bound and the
decompressed-size upper bounds (saturate to the representable ceiling),
the optimal-parser price comparison (base_cost can be the u32::MAX
"unreachable" sentinel — saturating keeps the compare correct), the
skip-step loop-break idiom (saturate then `next <= pos` detects the end),
untrusted-frame size accumulation, and abs-position/index accumulators
that rely on the rebase machinery. Behavior-identical; 838 tests pass.

* test(encode): add regression for HC/BT reserve hint overflow

driver_huge_source_hint_with_dict_does_not_overflow_hc_reserve sets a
u64::MAX source-size hint (the unknown-size sentinel) plus a positive
dictionary hint and resets at Level 16 (BtOpt → HashChain/BT storage
arm). On current code `(src as usize) + dict_hint` overflows usize
before reserve_history can clamp: debug panic / release wrap on 64-bit,
src-truncation on 32-bit. Test fails (overflow panic) until the fix.

* fix(encode): saturate HC/BT reserve hint instead of overflowing

(src as usize) + dict_hint.unwrap_or(0) overflowed usize when src is
the u64::MAX unknown-size sentinel plus a positive dictionary hint, and
truncated src on 32-bit targets. Saturate src via usize::try_from and
saturate the dict-hint addition via checked_add; reserve_history applies
the tighter window ceiling. Fixes the panic in the regression test added
in the preceding commit.

* fix(encode): avoid eviction-check overflow on 32-bit; fix stale docs

FastKernelMatcher eviction computed real_len + space.len() before the
block-size assert; at window_log = 30 both terms approach cap = 2^31, so
the sum overflows usize on 32-bit targets. Reorder to compute cap, assert
the block bound, then compare via subtraction (real_len > cap -
space.len()) which the assert proves cannot underflow. Correct the stale
window_log <= 31 comment to the enforced <= 30, and refresh the
compact_history doc to describe the quarter-window / half-mirror drain.

* fix(encode): use saturating_add for HC/BT reserve hint (clippy)

* test(encode): assert clamped HC reserve ceiling

- Strengthen driver_huge_source_hint_with_dict_does_not_overflow_hc_reserve
  to assert the post-reset history capacity reaches the clamped ceiling
  (window + window/4 + block), not just no-panic; an under-reserved mirror
  would otherwise pass.
- Correct the FastKernelMatcher eviction overflow-safety doc: the `* 2` is
  safe because dictionary priming caps max_window_size at
  MAX_PRIMED_WINDOW_SIZE = (u32::MAX - MAX_BLOCK_SIZE)/2 (so cap*2 < u32::MAX
  by construction), not because window_log <= 30 (priming widens beyond that).

Author: polaz

zstd/src/dictionary/mod.rs

@@ -232,7 +232,11 @@ pub fn create_raw_dict_from_source<R: io::Read, W: io::Write>(
     // the lowest scoring items come first
     let mut pool: BinaryHeap<Reverse<Segment>> = BinaryHeap::new();
     let (num_epochs, epoch_size_kmers) = compute_epoch_info(&params, dict_size, source_size / K);
-    let epoch_size = usize::max(K, epoch_size_kmers.saturating_mul(K));
+    // Plain `*`/`+` throughout the epoch walk below: epochs partition the
+    // training source, so `epoch_size_kmers * K`, `epoch_idx * epoch_size`, and
+    // `start + epoch_size` are all bounded by the source length (<= isize::MAX)
+    // and cannot overflow usize.
+    let epoch_size = usize::max(K, epoch_size_kmers * K);
     vprintln!("create_dict: computed epoch info, using {num_epochs} epochs of {epoch_size} bytes");
     let mut epoch_counter = 0;
     let mut ctx = Context {
@@ -242,14 +246,14 @@ pub fn create_raw_dict_from_source<R: io::Read, W: io::Write>(
     // segment for the pool. Keep exactly `num_epochs` windows to avoid
     // emitting more segments than the requested dictionary budget allows.
     for epoch_idx in 0..num_epochs {
-        let start = epoch_idx.saturating_mul(epoch_size);
+        let start = epoch_idx * epoch_size;
         if start >= all.len() {
             break;
         }
         let end = if epoch_idx + 1 == num_epochs {
             all.len()
         } else {
-            usize::min(start.saturating_add(epoch_size), all.len())
+            usize::min(start + epoch_size, all.len())
         };
         let epoch = &all[start..end];
         epoch_counter += 1;

zstd/src/encoding/blocks/compressed.rs

@@ -1603,8 +1603,10 @@ fn choose_table_from_counts<'a>(
     // wins. The estimate equals the built table's `table_header_bits()`
     // exactly, so the selection is byte-identical.
     let new_total_cost = (distinct_symbols > 1).then(|| {
+        // Plain `+`: both are bit-cost estimates bounded by the block size
+        // (<= MAX_BLOCK_SIZE * 8 bits), far under the integer's range.
         fse_header_bits_for_counts(&counts[..=max_symbol], max_log, use_low_prob_count)
-            .saturating_add(entropy_cost(counts, max_symbol, total))
+            + entropy_cost(counts, max_symbol, total)
     });
 
     let predefined_cost = cross_entropy_cost(counts, max_symbol, default_table);

zstd/src/encoding/cost_model/mod.rs

@@ -315,9 +315,10 @@ impl HcOptState {
                     self.lit_sum = Self::apply_seeded_table(&mut self.lit_freq, &seed.lit_bits, 11);
                 } else if self.literals_compressed() {
                     self.lit_freq.fill(0);
+                    // Plain `+`: lit_freq is u32 and a byte's count is bounded
+                    // by the block size (<= MAX_BLOCK_SIZE), far under u32::MAX.
                     for &byte in src {
-                        self.lit_freq[byte as usize] =
-                            self.lit_freq[byte as usize].saturating_add(1);
+                        self.lit_freq[byte as usize] += 1;
                     }
                     self.lit_sum = Self::downscale_stats(&mut self.lit_freq, 8, false);
                     if self.lit_sum == 0 {
@@ -363,9 +364,10 @@ impl HcOptState {
             } else {
                 if self.literals_compressed() {
                     self.lit_freq.fill(0);
+                    // Plain `+`: lit_freq is u32 and a byte's count is bounded
+                    // by the block size (<= MAX_BLOCK_SIZE), far under u32::MAX.
                     for &byte in src {
-                        self.lit_freq[byte as usize] =
-                            self.lit_freq[byte as usize].saturating_add(1);
+                        self.lit_freq[byte as usize] += 1;
                     }
                     self.lit_sum = Self::downscale_stats(&mut self.lit_freq, 8, false);
                     if self.lit_sum == 0 {

zstd/src/encoding/frame_compressor.rs

@@ -445,7 +445,10 @@ fn presplit_distance(lhs: &PreSplitFingerprint, rhs: &PreSplitFingerprint, hash_
     for idx in 0..slots {
         let left = lhs.events[idx] as i128 * rhs.nb_events as i128;
         let right = rhs.events[idx] as i128 * lhs.nb_events as i128;
-        distance = distance.saturating_add(left.abs_diff(right) as u64);
+        // Plain `+`: events/nb_events are per-block sample counts (<= block
+        // size), so each |left-right| <= (2^17)^2 and the sum over <= 2^hash_log
+        // slots stays far under u64::MAX — no overflow.
+        distance += left.abs_diff(right) as u64;
     }
     distance
 }
@@ -460,16 +463,21 @@ fn presplit_fingerprints_differ(
     debug_assert!(new_fp.nb_events > 0);
     let p50 = reference.nb_events as u64 * new_fp.nb_events as u64;
     let deviation = presplit_distance(reference, new_fp, hash_log);
-    let threshold = p50.saturating_mul(PRESPLIT_THRESHOLD_BASE + penalty as u64)
-        / PRESPLIT_THRESHOLD_PENALTY_RATE;
+    // Plain `*`: p50 <= (block-sample-count)^2 and the (base+penalty) factor is
+    // a small constant, so the product stays well under u64::MAX.
+    let threshold =
+        p50 * (PRESPLIT_THRESHOLD_BASE + penalty as u64) / PRESPLIT_THRESHOLD_PENALTY_RATE;
     deviation >= threshold
 }
 
 fn presplit_merge_events(acc: &mut PreSplitFingerprint, new_fp: &PreSplitFingerprint) {
+    // Plain `+`: `acc` accumulates only the chunks of a single block (caller
+    // loops within one block, <= MAX_BLOCK_SIZE), so the merged sample counts
+    // stay far under u32 / usize bounds — no overflow.
     for idx in 0..PRESPLIT_HASH_TABLE_SIZE {
-        acc.events[idx] = acc.events[idx].saturating_add(new_fp.events[idx]);
+        acc.events[idx] += new_fp.events[idx];
     }
-    acc.nb_events = acc.nb_events.saturating_add(new_fp.nb_events);
+    acc.nb_events += new_fp.nb_events;
 }
 
 fn split_block_by_chunks(block: &[u8], level: usize) -> usize {

zstd/src/encoding/match_generator.rs

@@ -792,8 +792,11 @@ fn dict_and_window_log(window_log: u8, src_size: u64, dict_size: u64) -> u32 {
         return window_log as u32;
     }
     let window_size: u64 = 1u64 << window_log;
-    let dict_and_window = dict_size.saturating_add(window_size);
-    if window_size >= dict_size.saturating_add(src_size) {
+    // Plain `+` (matches upstream zstd `ZSTD_dictAndWindowLog`): `window_size` is
+    // `1 << window_log` (window_log <= 31) and dict/src are real data sizes
+    // (<= isize::MAX), so these u64 sums cannot overflow in practice.
+    let dict_and_window = dict_size + window_size;
+    if window_size >= dict_size + src_size {
         // Window already covers source + dictionary.
         window_log as u32
     } else {
@@ -821,7 +824,9 @@ fn cdict_table_logs(
     // createCDict assumes a minSrcSize source when the real size is unknown.
     let src_size = DICT_MIN_SRC_SIZE;
     // Source-size window resize (donor caps windowLog by ceil_log2(src+dict)).
-    let tsize = src_size.saturating_add(dict_size);
+    // Plain `+`: src_size is the tiny DICT_MIN_SRC_SIZE constant and dict_size
+    // is a real dictionary length, so the u64 sum cannot overflow.
+    let tsize = src_size + dict_size;
     let resized_window_log = (window_log as u32)
         .min(source_size_ceil_log(tsize) as u32)
         .max(1);
@@ -2270,7 +2275,14 @@ impl Matcher for MatchGeneratorDriver {
                 // tables are dictionary-tier-small. Unhinted streams skip this
                 // and keep doubling growth.
                 if let Some(src) = hint {
-                    let expected = (src as usize).saturating_add(dict_hint.unwrap_or(0));
+                    // `src` is a u64 hint and may be the u64::MAX "unknown
+                    // size" sentinel, which truncates under `as usize` on
+                    // 32-bit targets and overflows when the dict hint is
+                    // added. Saturate the source size, then saturate the
+                    // dict-hint addition; `reserve_history` applies the
+                    // tighter window ceiling to the result.
+                    let src_hint = usize::try_from(src).unwrap_or(usize::MAX);
+                    let expected = src_hint.saturating_add(dict_hint.unwrap_or(0));
                     hc.table.reserve_history(expected);
                 }
             }
@@ -2779,7 +2791,9 @@ impl Matcher for MatchGeneratorDriver {
                     // when it later reuses the buffer.
                     vec_pool.push(data);
                 });
-                evicted_bytes += pre.saturating_add(space_len).saturating_sub(m.window_size);
+                // Plain `+` (the `saturating_sub` floors at 0): `pre` + one
+                // block are byte counts bounded by the window, no overflow.
+                evicted_bytes += (pre + space_len).saturating_sub(m.window_size);
             }
             MatcherStorage::Row(m) => {
                 // RowMatchGenerator::add_data recycles the *input* buffer
@@ -2799,7 +2813,9 @@ impl Matcher for MatchGeneratorDriver {
                     // `slice_size` on reuse).
                     vec_pool.push(data);
                 });
-                evicted_bytes += pre.saturating_add(space_len).saturating_sub(m.window_size);
+                // Plain `+` (the `saturating_sub` floors at 0): `pre` + one
+                // block are byte counts bounded by the window, no overflow.
+                evicted_bytes += (pre + space_len).saturating_sub(m.window_size);
             }
             MatcherStorage::HashChain(m) => {
                 // MatchTable::add_data now recycles the *incoming* buffer
@@ -2822,9 +2838,9 @@ impl Matcher for MatchGeneratorDriver {
                     // pool retains.
                     vec_pool.push(data);
                 });
-                evicted_bytes += pre
-                    .saturating_add(space_len)
-                    .saturating_sub(m.table.window_size);
+                // Plain `+` (the `saturating_sub` floors at 0): byte counts
+                // bounded by the window, no overflow.
+                evicted_bytes += (pre + space_len).saturating_sub(m.table.window_size);
             }
         }
         // Gate the second backend trim pass on actual budget
@@ -3975,6 +3991,11 @@ macro_rules! build_optimal_plan_impl_body {
             }
 
             let next_price = unsafe { nodes.get_unchecked(pos + 1).price };
+            // `saturating_add` is REQUIRED here, not a masked bug: `base_cost`
+            // is a node price that can be the `u32::MAX` "unreachable" sentinel,
+            // and saturating keeps `base_cost + margin` pinned at MAX so the
+            // comparison stays correct. Plain `+` would wrap the sentinel and
+            // flip the abort decision (a ratio bug / debug overflow panic).
             if abort_on_worse_match
                 && next_price <= base_cost.saturating_add(HC_BITCOST_MULTIPLIER / 2)
             {
@@ -8056,6 +8077,42 @@ fn driver_huge_source_hint_does_not_overflow_table_window_shift() {
     );
 }
 
+#[test]
+fn driver_huge_source_hint_with_dict_does_not_overflow_hc_reserve() {
+    // Regression: the HC/BT history-mirror pre-size adds the dictionary
+    // hint to the source-size hint before `reserve_history` clamps to the
+    // window ceiling. A `u64::MAX` pledged source size (the "unknown size"
+    // sentinel) plus any positive dictionary hint overflows `usize` in
+    // `(src as usize) + dict_hint` — debug panic / release wrap on 64-bit,
+    // and `src as usize` truncation on 32-bit targets. Level 16 (BtOpt)
+    // routes through the HashChain/BT storage arm that owns this reserve.
+    // Must size the mirror to the real window, never panic, wrap, or
+    // truncate.
+    let mut driver = MatchGeneratorDriver::new(32, 2);
+    driver.set_source_size_hint(u64::MAX);
+    driver.set_dictionary_size_hint(64 * 1024);
+    driver.reset(CompressionLevel::Level(16));
+
+    // The saturated `usize::MAX` reserve target must be clamped to the HC
+    // history ceiling, not reserved literally (which would OOM/panic). Level 16
+    // has window_log 22, so the ceiling is `window + window/4 + one block`
+    // (the `reserve_history` formula). Assert the reserve actually reached it —
+    // a no-panic-only check would also pass on an under-reserved mirror.
+    let window = 1usize << 22;
+    let expected_history_ceiling = window + (window >> 2) + crate::common::MAX_BLOCK_SIZE as usize;
+    assert!(
+        driver.hc_matcher().table.history.capacity() >= expected_history_ceiling,
+        "huge source + dict hint must reserve the clamped HC history ceiling, got {}",
+        driver.hc_matcher().table.history.capacity()
+    );
+
+    let mut space = driver.get_next_space();
+    space[..12].copy_from_slice(b"abcabcabcabc");
+    space.truncate(12);
+    driver.commit_space(space);
+    driver.skip_matching_with_hint(None);
+}
+
 #[test]
 fn driver_small_source_hint_shrinks_row_hash_tables() {
     let mut driver = MatchGeneratorDriver::new(32, 2);

zstd/src/encoding/match_table/storage.rs

@@ -724,12 +724,19 @@ impl MatchTable {
     }
 
     pub(crate) fn reserve_history(&mut self, expected_bytes: usize) {
-        // Eviction keeps the live mirror within `max_window_size`; one pending
-        // block can sit on top before `add_data` rolls it out, so the tightest
-        // sufficient capacity is `max_window_size + MAX_BLOCK_SIZE`.
-        let cap = self
-            .max_window_size
-            .saturating_add(crate::common::MAX_BLOCK_SIZE as usize);
+        // Eviction keeps the live mirror within `max_window_size`; the dead
+        // prefix is drained at a quarter window (see `compact_history`) and one
+        // pending block can sit on top, so the steady-state ceiling is
+        // `max_window_size + max_window_size/4 + MAX_BLOCK_SIZE` — matching the
+        // `add_data` eviction reserve so the two never fight over capacity.
+        // Plain arithmetic (not `saturating_*`): `max_window_size = 1 <<
+        // window_log` with `window_log <= ZSTD_WINDOWLOG_MAX` (31), so the sum
+        // is at most `2^31 + 2^29 + 2^17 < usize::MAX` even on 32-bit targets —
+        // overflow is unreachable, and a silent saturation here would only mask
+        // a window_log bound violation upstream.
+        let cap = self.max_window_size
+            + (self.max_window_size >> 2)
+            + crate::common::MAX_BLOCK_SIZE as usize;
         let want = expected_bytes.min(cap);
         if want > self.history.capacity() {
             self.history.reserve_exact(want - self.history.len());
@@ -739,6 +746,20 @@ impl MatchTable {
     pub(crate) fn add_data(&mut self, data: Vec<u8>, mut reuse_space: impl FnMut(Vec<u8>)) {
         assert!(data.len() <= self.max_window_size);
         check_stream_abs_headroom(self.history_abs_start, self.window_size, data.len());
+        if self.window_size + data.len() > self.max_window_size {
+            // Cap the history mirror near the live window once eviction starts:
+            // reserve exactly (window + window/4 + one block) so the Vec grows
+            // linearly to that ceiling instead of power-of-two doubling to ~2x
+            // window; `compact_history`'s quarter-window drain keeps len under
+            // it, so the Vec never reallocates again. Small frames that never
+            // fill the window keep their tight data-sized buffer.
+            let target = self.max_window_size
+                + (self.max_window_size >> 2)
+                + crate::common::MAX_BLOCK_SIZE as usize;
+            if target > self.history.len() && self.history.capacity() < target {
+                self.history.reserve_exact(target - self.history.len());
+            }
+        }
         while self.window_size + data.len() > self.max_window_size {
             let removed_len = self.chunk_lens.pop_front().unwrap();
             self.window_size -= removed_len;
@@ -769,13 +790,17 @@ impl MatchTable {
     }
 
     /// Drain the dead prefix of `history` (already-rolled-out bytes)
-    /// when it has grown to at least half the live region. Keeps the
-    /// contiguous mirror compact so reallocation costs stay amortised.
+    /// once it reaches a quarter window, or when it accounts for at
+    /// least half of the mirror. Keeps the contiguous mirror compact
+    /// so reallocation costs stay amortised.
     pub(crate) fn compact_history(&mut self) {
         if self.history_start == 0 {
             return;
         }
-        if self.history_start >= self.max_window_size
+        // Drain the dead prefix at a quarter window (paired with the eviction
+        // reserve in `add_data`) so the mirror stays near `window + window/4`
+        // rather than doubling to ~2x window on long streams.
+        if self.history_start >= (self.max_window_size >> 2)
             || self.history_start * 2 >= self.history.len()
         {
             self.history.drain(..self.history_start);

zstd/src/encoding/mod.rs

@@ -234,6 +234,11 @@ pub const fn compress_bound(src_size: usize) -> usize {
     } else {
         0
     };
+    // Saturating is the correct UPPER-BOUND semantic here, not a masked bug:
+    // this is a public API over an arbitrary `usize`, and the largest meaningful
+    // bound is `usize::MAX`. A real slice is at most `isize::MAX` bytes, so the
+    // `* 1.004 + margin` cannot overflow for genuine inputs; the saturation only
+    // caps a pathological caller-supplied size at the representable ceiling.
     src_size
         .saturating_add(src_size >> 8)
         .saturating_add(margin)

zstd/src/encoding/sequence_capture.rs

@@ -113,7 +113,9 @@ impl Matcher for CapturingMatcher {
         let tail_ll = self.inner.get_last_space().len() as u32;
         self.inner.skip_matching();
         self.block_tail_lengths.borrow_mut().push(tail_ll);
-        self.current_block = self.current_block.saturating_add(1);
+        // Plain `+`: diagnostic per-block index; the comparator runs on bench
+        // fixtures whose block count is nowhere near u32::MAX.
+        self.current_block += 1;
     }
 
     fn skip_matching_with_hint(&mut self, incompressible_hint: Option<bool>) {
@@ -126,7 +128,9 @@ impl Matcher for CapturingMatcher {
         let tail_ll = self.inner.get_last_space().len() as u32;
         self.inner.skip_matching_with_hint(incompressible_hint);
         self.block_tail_lengths.borrow_mut().push(tail_ll);
-        self.current_block = self.current_block.saturating_add(1);
+        // Plain `+`: diagnostic per-block index; the comparator runs on bench
+        // fixtures whose block count is nowhere near u32::MAX.
+        self.current_block += 1;
     }
 
     fn start_matching(&mut self, mut handle_sequence: impl for<'a> FnMut(Sequence<'a>)) {
@@ -164,7 +168,8 @@ impl Matcher for CapturingMatcher {
                         of: *offset as u32,
                         ml: *match_len as u32,
                     });
-                    seq_in_block = seq_in_block.saturating_add(1);
+                    // Plain `+`: sequences per block <= MAX_BLOCK_SIZE, far under u32::MAX.
+                    seq_in_block += 1;
                 }
                 Sequence::Literals { literals } => {
                     block_tail_ll = literals.len() as u32;
@@ -173,7 +178,9 @@ impl Matcher for CapturingMatcher {
             handle_sequence(seq);
         });
         self.block_tail_lengths.borrow_mut().push(block_tail_ll);
-        self.current_block = self.current_block.saturating_add(1);
+        // Plain `+`: diagnostic per-block index; the comparator runs on bench
+        // fixtures whose block count is nowhere near u32::MAX.
+        self.current_block += 1;
     }
 
     fn reset(&mut self, level: CompressionLevel) {