perf(encode): flush the sequence bitstream where upstream flushes it

zstd/src/encoding/blocks/compressed.rs

@@ -2280,15 +2280,15 @@ fn encode_sequences(
     // Upstream zstd-faithful sequence loop: write state diffs + extras via
     // unchecked fast-path adds with explicit `flush_bulk` calls at
     // safe burst boundaries. Per-sequence bit budget:
-    //   state diffs: of (<=8) + ml (<=9) + ll (<=9) = 26 bits → one
-    //                burst between flushes.
-    //   extras:      ll (<=16) + ml (<=16) + of (<=24) = 56 bits →
-    //                one burst between flushes.
+    //   state diffs: of (<=8) + ml (<=9) + ll (<=9) = 26 bits.
+    //   extras:      ll (<=16) + ml (<=16) + of (<=31).
     //
-    // Total per sequence: 82 bits ⇒ at least 2 flushes (one per burst).
-    // Mirrors upstream zstd `ZSTD_encodeSequences_body`
-    // (`zstd_compress_sequences.c:303-360`) which uses BIT_addBitsFast
-    // + BIT_flushBitsFast at the same burst boundaries.
+    // One flush per sequence is the common case: the two inside the body are
+    // conditional on the extras not fitting beside the diffs, which is
+    // upstream's own accounting in `ZSTD_encodeSequences_body`
+    // (`zstd_compress_sequences.c:303-367`), and the one at the end of the
+    // body is what leaves the accumulator with under a byte for the next
+    // round.
     //
     // Pre-reserve output capacity for the worst-case sequence section
     // size (~10 bytes/sequence + 32 byte slack) so the per-flush
@@ -2327,12 +2327,22 @@ fn encode_sequences(
             // per-sequence flush in this loop (≤ 16 bytes per
             // sequence, plus the 32-byte slack on top of the 64-byte
             // header reserve).
+            //
+            // What the three diffs actually put in the accumulator, tallied so
+            // the budget below is checked against the widths the FSE tables
+            // produced rather than against the ceiling its 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 diff = crate::fse::fse_encoder::transition_bits(of_state, next.num_bits);
                 unsafe {
                     writer.write_bits_64_no_check(diff as u64, next.num_bits as usize);
                 }
+                #[cfg(debug_assertions)]
+                {
+                    state_diff_bits += next.num_bits as usize;
+                }
                 of_state = next.index;
             }
             if let Some(table) = ml_table {
@@ -2341,6 +2351,10 @@ fn encode_sequences(
                 unsafe {
                     writer.write_bits_64_no_check(diff as u64, next.num_bits as usize);
                 }
+                #[cfg(debug_assertions)]
+                {
+                    state_diff_bits += next.num_bits as usize;
+                }
                 ml_state = next.index;
             }
             if let Some(table) = ll_table {
@@ -2349,35 +2363,76 @@ fn encode_sequences(
                 unsafe {
                     writer.write_bits_64_no_check(diff as u64, next.num_bits as usize);
                 }
+                #[cfg(debug_assertions)]
+                {
+                    state_diff_bits += next.num_bits as usize;
+                }
                 ll_state = next.index;
             }
-            unsafe {
-                writer.flush_bulk();
+            // The three state diffs and the three extra-bit fields share one
+            // 64-bit accumulator, and a flush here is only needed when the
+            // extras that follow would not fit beside them. Upstream asks
+            // exactly that question (`zstd_compress_sequences.c:350`):
+            // `ofBits + mlBits + llBits >= 64 - 7 - (LLFSELog + MLFSELog +
+            // OffFSELog)`, which with our accumulator logs (9 + 9 + 8 = 26) is
+            // 31. Below it the accumulator holds at most 7 leftover + 26 diff
+            // bits + 30 extra bits = 63, so nothing has to leave yet. Flushing
+            // unconditionally instead cost a second store, length commit and
+            // shift on every sequence, and a level-1 block has hundreds of
+            // thousands of them.
+            let extra_bits_total = of_num_bits + ml_num_bits + ll_num_bits;
+            // The thresholds are arithmetic on widths, so they are only right
+            // while the widths are what they were derived from. Pinned here,
+            // where the arithmetic happens, so a widened encoder fails at its
+            // cause rather than as an accumulator overflow further down or, in
+            // a release build, as a corrupted stream with nothing to point at.
+            #[cfg(debug_assertions)]
+            {
+                debug_assert!(
+                    state_diff_bits <= 26,
+                    "state diffs took {state_diff_bits} bits; the thresholds below are \
+                     derived from a 26-bit ceiling (LLFSELog 9 + MLFSELog 9 + OffFSELog 8)",
+                );
+                debug_assert!(
+                    ll_num_bits <= 16 && ml_num_bits <= 16 && of_num_bits <= 31,
+                    "extra-bit widths ll={ll_num_bits} ml={ml_num_bits} of={of_num_bits} \
+                     exceed the 16 / 16 / 31 the thresholds are derived from",
+                );
+                // The bound the branch itself rests on: without a flush, 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,
+                    "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",
+                );
+            }
+            if extra_bits_total >= 31 {
+                unsafe {
+                    writer.flush_bulk();
+                }
             }
 
-            // Extras burst: ll (≤16) + ml (≤16) + of (≤ window_log,
-            // up to 30 for our max window_log). With ≤ 7 leftover from
-            // the prior flush_bulk, total ll+ml+of+partial can exceed
-            // 64 once of_num_bits > 25. Upstream zstd handles this via
-            // `longOffsets` mode that splits high offsets across two
-            // BIT_addBits calls; we instead drain the partial after ml
-            // and write of into a fresh container. The branch matches
-            // upstream zstd's `MEM_32bits()` flush-between-each-component
-            // shape on the 32-bit build (which has the same 64-bit
-            // container constraint).
+            // Extras burst: ll (≤16) + ml (≤16) + of (≤ 31). Whether the
+            // accumulator was drained above or not, ll and ml fit: after a
+            // flush it holds ≤ 7 + 32 = 39, and without one the three extras
+            // together were under 31.
             //
             // SAFETY: `encode_literal_length` / `encode_match_len`
             // bound `*_num_bits ≤ 16` and return a clean `*_add_bits`
             // (low `num_bits` bits only). `encode_offset` bounds
-            // `of_num_bits ≤ ilog2(of)`, capped at the encoder's
-            // `window_log` ≤ 30; the conditional flush_bulk above
-            // drains the partial when of_num_bits crosses the 24-bit
-            // threshold where the sum could exceed 64.
+            // `of_num_bits ≤ ilog2(of) ≤ 31`, and the two conditionals
+            // (`>= 31` above, `> 56` below) are upstream's own accounting for
+            // when the offset field still fits beside what is already there.
             unsafe {
                 writer.write_bits_64_no_check(ll_add_bits as u64, ll_num_bits);
                 writer.write_bits_64_no_check(ml_add_bits as u64, ml_num_bits);
             }
-            if of_num_bits > 24 {
+            // Upstream `zstd_compress_sequences.c:355`: with the accumulator
+            // drained before the diffs, 7 leftover plus all three extra fields
+            // must stay under 64, so the offset needs its own container only
+            // past 56 bits of extras.
+            if extra_bits_total > 56 {
                 unsafe {
                     writer.flush_bulk();
                 }