diff --git a/ffi-bench/Cargo.toml b/ffi-bench/Cargo.toml index 9e7a8e0c7..3c5891b7a 100644 --- a/ffi-bench/Cargo.toml +++ b/ffi-bench/Cargo.toml @@ -160,6 +160,10 @@ path = "../zstd/tests/he_level22_ratio.rs" name = "incompressibility_falsepos" path = "../zstd/tests/incompressibility_falsepos.rs" +[[example]] +name = "slide_oversized_table" +path = "../zstd/examples/slide_oversized_table.rs" + [[example]] name = "alloc_audit_decode" path = "../zstd/examples/alloc_audit_decode.rs" diff --git a/zstd/examples/slide_oversized_table.rs b/zstd/examples/slide_oversized_table.rs new file mode 100644 index 000000000..f99c4a891 --- /dev/null +++ b/zstd/examples/slide_oversized_table.rs @@ -0,0 +1,87 @@ +//! Window slides with a hash table far larger than the window. +//! +//! `hashLog` is capped at `windowLog + 1` for a frame with no dictionary, so +//! the table is at most twice the window there. The cap is lifted when a +//! dictionary is attached, because the main and dictionary tables share one +//! `hashLog` — which leaves a configuration the advanced API can ask for and +//! the level presets never produce: a table of a million entries over a window +//! of a kilobyte, sliding once per kilobyte of input. +//! +//! That is the shape where sliding the table's indices and rebuilding it from +//! the retained bytes cost very different amounts, so it is the fixture for +//! deciding between them. +//! +//! Build: cargo build --profile bench -p structured-zstd +//! --example slide_oversized_table --features hash,std,dict-builder +//! Run: ./target/release/examples/slide_oversized_table +//! + +use std::env; + +use structured_zstd::encoding::{ + CompressionLevel, CompressionParameters, FrameCompressor, Strategy, +}; + +/// Compressible but not degenerate: repeated lines with a rotating field, so +/// the matcher finds real matches across the window without the whole frame +/// collapsing to one repeat. +fn body(len: usize) -> Vec { + let mut bytes = Vec::with_capacity(len); + let mut counter = 0u32; + while bytes.len() < len { + let line = format!( + "ts=2026-03-26T21:39:{:02}Z level=INFO msg=\"flush memtable\" seq={counter} \ + tenant=demo table=orders region=eu-west\n", + counter % 60, + ); + let remaining = len - bytes.len(); + bytes.extend_from_slice(&line.as_bytes()[..line.len().min(remaining)]); + counter += 1; + } + bytes +} + +fn main() { + let args: Vec = env::args().collect(); + let window_log: u32 = args.get(1).and_then(|s| s.parse().ok()).unwrap_or(10); + let hash_log: u32 = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(20); + let frame_bytes: usize = args + .get(3) + .and_then(|s| s.parse().ok()) + .unwrap_or(4 * 1024 * 1024); + let iters: u32 = args.get(4).and_then(|s| s.parse().ok()).unwrap_or(4); + let dict_path: Option<&str> = args.get(5).map(|s| s.as_str()); + + let src = body(frame_bytes); + + let params = CompressionParameters::builder(CompressionLevel::Level(1)) + .strategy(Strategy::Fast) + .window_log(window_log) + .hash_log(hash_log) + .build() + .expect("parameters within bounds"); + + let mut cctx: FrameCompressor = FrameCompressor::new(CompressionLevel::Level(1)); + cctx.set_parameters(¶ms); + if let Some(path) = dict_path { + let dict = std::fs::read(path).expect("read dict file"); + cctx.set_dictionary_from_bytes(&dict) + .expect("dictionary should attach"); + } + + let mut out: Vec = Vec::new(); + let mut sink: usize = 0; + for _ in 0..iters { + cctx.compress_independent_frame_into(&src, &mut out); + sink = sink.wrapping_add(out.len()); + core::hint::black_box(&out); + } + + eprintln!( + "windowLog={window_log} hashLog={hash_log} frame={frame_bytes} iters={iters} \ + dict={} out={} sum={sink} heap={}", + dict_path.unwrap_or("none"), + out.len(), + cctx.heap_size(), + ); +} diff --git a/zstd/src/bit_io/bit_writer.rs b/zstd/src/bit_io/bit_writer.rs index c9abddc73..e4ad9ad07 100644 --- a/zstd/src/bit_io/bit_writer.rs +++ b/zstd/src/bit_io/bit_writer.rs @@ -157,23 +157,13 @@ impl>> BitWriter { /// pays. #[inline(always)] pub unsafe fn write_bits_64_no_check(&mut self, bits: u64, num_bits: usize) { - // num_bits == 0 short-circuit: matches upstream zstd `BIT_addBits` no-op - // semantics AND guards the `bits << self.bits_in_partial` below - // from a `<< 64` undefined-behaviour evaluation when the - // accumulator is already full (`bits_in_partial == 64`). Callers - // that legitimately drain a full container (e.g. the FSE encoder - // hitting a state-diff burst boundary) can call this with - // `num_bits = 0` as a no-op without tripping UB. - if num_bits == 0 { - return; - } debug_assert!( num_bits + self.bits_in_partial <= 64, "write_bits_64_no_check would overflow partial: would push to {} bits", num_bits + self.bits_in_partial, ); debug_assert!( - self.bits_in_partial < 64, + self.bits_in_partial < 64 || num_bits == 0, "write_bits_64_no_check called with full accumulator and num_bits>0; \ caller must flush_bulk before adding more bits", ); @@ -181,7 +171,19 @@ impl>> BitWriter { num_bits == 64 || bits >> num_bits == 0, "value has dirty high bits beyond num_bits={num_bits}", ); - self.partial |= bits << self.bits_in_partial; + // Masked so a full accumulator cannot make this a shift by 64, which is + // undefined for `u64`. It is not a guard costing anything: x86 and + // AArch64 shift instructions mask the count themselves, so the mask + // disappears. It replaces a `num_bits == 0` early return, which was a + // branch on every call, and there are six a sequence. Upstream needs + // neither, because its `BIT_addBitsFast` keeps `bitPos` strictly under + // 64 and can shift unconditionally. + // + // Correct in the one case the mask changes: the accumulator is full + // only where the caller passes `num_bits == 0`, and the second + // precondition then forces `bits == 0`, so the OR contributes nothing + // whatever the shift count. + self.partial |= bits << (self.bits_in_partial & 63); self.bits_in_partial += num_bits; } diff --git a/zstd/src/encoding/blocks/compressed.rs b/zstd/src/encoding/blocks/compressed.rs index c2f4a38f1..0d161fa50 100644 --- a/zstd/src/encoding/blocks/compressed.rs +++ b/zstd/src/encoding/blocks/compressed.rs @@ -141,10 +141,14 @@ struct EncodedBlockParts { #[derive(Default)] pub(crate) struct CompressedBlockScratch { parts: EncodedBlockParts, + /// One packed [`SequenceCode`] per sequence of the partition being + /// encoded, filled by the pass that derives the offset codes and read by + /// the bit writer. Kept here so it is allocated once for the compressor + /// rather than per block. + sequence_codes: Vec, partitions: Vec, prefix_sums: SequencePrefixSums, compressed: Vec, - estimator_sequences: Vec, /// Lazily allocated: only the block-split estimator path uses it, and /// `compress_block`'s `mem::take` constructs a throwaway `Default` /// scratch every block — an eager workspace made that default pay four @@ -178,11 +182,10 @@ impl CompressedBlockScratch { pub(crate) fn retained_heap_size(&self) -> usize { self.parts.literals.capacity() + self.parts.sequences.capacity() * core::mem::size_of::() + + self.sequence_codes.capacity() * core::mem::size_of::() + self.partitions.capacity() * core::mem::size_of::() + self.prefix_sums.heap_size() + self.compressed.capacity() - + self.estimator_sequences.capacity() - * core::mem::size_of::() + self .estimator_workspace .as_ref() @@ -247,11 +250,22 @@ impl SequencePrefixSums { } } +/// One collected sequence. +/// +/// `off_base` holds the offset the matcher found until [`fill_wire_offsets`] +/// runs over the sequence, and the wire code from then on: 1/2/3 for the repeat +/// offsets, N+3 for an explicit N. It is one field rather than two because the +/// found offset has no reader once its code exists, and a fourth word would +/// widen every sequence in the block for a value with a lifetime of one pass. +/// Upstream keeps the same single slot, filled at match time +/// (`SeqDef::offBase`, written by `ZSTD_storeSeq`); ours cannot be filled that +/// early because the code depends on a repeat-offset history that must not +/// advance across a partition the emitter ends up writing raw. #[derive(Clone, Copy)] struct RawSequence { ll: u32, ml: u32, - offset: u32, + off_base: u32, } struct EntropyOnlyMatcher; @@ -294,12 +308,12 @@ impl Matcher for EntropyOnlyMatcher { pub fn compress_block(state: &mut CompressState, output: &mut Vec) { let mut scratch = core::mem::take(&mut state.block_scratch); collect_block_parts(state, &mut scratch.parts); - let decisions = encode_block_parts_with_sequence_scratch( + let decisions = encode_block_parts( state, &scratch.parts.literals, - &scratch.parts.sequences, + &mut scratch.parts.sequences, + &mut scratch.sequence_codes, output, - &mut scratch.estimator_sequences, ); // This path writes the block it just encoded, so the tables it chose are // what the next block reads. @@ -338,14 +352,14 @@ pub(crate) fn compress_block_with_post_split( let mut emit_buffers = SingleSequenceEmitBuffers { output, compressed: &mut scratch.compressed, - sequence_scratch: &mut scratch.estimator_sequences, + codes: &mut scratch.sequence_codes, }; let emitted_raw = emit_single_sequence_block( state, last_block, source_len, &scratch.parts.literals, - &scratch.parts.sequences, + &mut scratch.parts.sequences, &mut emit_buffers, ); if emitted_raw { @@ -448,14 +462,14 @@ pub(crate) fn compress_block_with_post_split( let mut emit_buffers = SingleSequenceEmitBuffers { output, compressed: &mut scratch.compressed, - sequence_scratch: &mut scratch.estimator_sequences, + codes: &mut scratch.sequence_codes, }; let emitted_raw = emit_single_sequence_block( state, last_block && last_partition, src_size, &scratch.parts.literals[lit_start..lit_end], - &scratch.parts.sequences[seq_start..seq_end], + &mut scratch.parts.sequences[seq_start..seq_end], &mut emit_buffers, ); if emitted_raw { @@ -582,33 +596,30 @@ fn collect_block_parts(state: &mut CompressState, parts: &mut Enc parts.sequences.push(RawSequence { ll, ml: match_len as u32, - offset: offset as u32, + // The found offset. `fill_wire_offsets` replaces it with its + // code once the partition this sequence lands in is about to be + // encoded, since the code depends on a history that partition + // boundaries can rewind. + off_base: offset as u32, }); } }); } -fn encode_block_parts_with_sequence_scratch( +fn encode_block_parts( state: &mut CompressState, literals_vec: &[u8], - raw_sequences: &[RawSequence], + raw_sequences: &mut [RawSequence], + // Scratch for the packed per-sequence codes, carried by the caller so it is + // allocated once rather than per block. + codes: &mut Vec, output: &mut Vec, - sequences: &mut Vec, // What each axis decided, for the caller to apply once it knows the block // is kept. LL, ML, OF. ) -> [LastUsedTable; 3] { // A block with no sequences writes no tables, so every axis keeps what it // had whatever the caller decides. let mut decisions = [LastUsedTable::Keep; 3]; - encode_raw_sequences_into( - raw_sequences, - &mut state.offset_hist, - sequences, - matches!( - state.strategy_tag, - crate::encoding::strategy::StrategyTag::Fast - ), - ); // literals section @@ -684,10 +695,10 @@ fn encode_block_parts_with_sequence_scratch( // sequences section - if sequences.is_empty() { + if raw_sequences.is_empty() { writer.write_bits(0u8, 8); } else { - encode_seqnum(sequences.len(), &mut writer); + encode_seqnum(raw_sequences.len(), &mut writer); // Single-pass histogram of ll/ml/of codes across all sequences. // Previously did three separate `sequences.iter().map(...)` @@ -697,34 +708,36 @@ fn encode_block_parts_with_sequence_scratch( let mut ll_counts = [0usize; 256]; let mut ml_counts = [0usize; 256]; let mut of_counts = [0usize; 256]; - // Track the highest code per stream while histogramming so the table - // selector skips the full-256 reverse scan for `max_symbol` (the small - // sequence-code alphabets leave ~200 high slots permanently zero). - let mut ll_max = 0usize; - let mut ml_max = 0usize; - let mut of_max = 0usize; - for seq in sequences.iter() { - let ll_code = encode_literal_length(seq.ll).0 as usize; - let ml_code = encode_match_len(seq.ml).0 as usize; - let of_code = encode_offset(seq.of).0 as usize; - ll_counts[ll_code] += 1; - ml_counts[ml_code] += 1; - of_counts[of_code] += 1; - ll_max = ll_max.max(ll_code); - ml_max = ml_max.max(ml_code); - of_max = of_max.max(of_code); - } - let total = sequences.len(); + // The offset codes are derived in this same pass rather than by a + // walk of their own ahead of the literals section: both passes read + // every sequence, and the second one was reading back what the first + // had just written. + let counts = SequenceCodeCounts { + ll: &mut ll_counts, + ml: &mut ml_counts, + of: &mut of_counts, + }; + let (ll_max, ml_max, of_max) = if matches!( + state.strategy_tag, + crate::encoding::strategy::StrategyTag::Fast + ) { + fill_and_count::(raw_sequences, &mut state.offset_hist, counts, codes) + } else { + fill_and_count::(raw_sequences, &mut state.offset_hist, counts, codes) + }; + let raw_sequences: &[RawSequence] = raw_sequences; + let codes: &[u32] = codes; + let total = raw_sequences.len(); // Stream codes of the LAST sequence: upstream zstd codes the final symbol // of each stream via the FSE init-state and drops one occurrence of it // from the emitted table's histogram (see `build_seq_ctable`). `Some` // here because these modes are written to the frame. - let (last_ll, last_ml, last_of) = sequences.last().map_or((0, 0, 0), |seq| { + let (last_ll, last_ml, last_of) = raw_sequences.last().map_or((0, 0, 0), |seq| { ( encode_literal_length(seq.ll).0 as usize, encode_match_len(seq.ml).0 as usize, - encode_offset(seq.of).0 as usize, + encode_offset(seq.off_base).0 as usize, ) }); @@ -786,7 +799,8 @@ fn encode_block_parts_with_sequence_scratch( encode_table(&ml_mode, &mut writer); encode_sequences( - sequences, + raw_sequences, + codes, &mut writer, &ll_mode, &ml_mode, @@ -816,7 +830,7 @@ struct EstimatorWorkspace { ll_counts: Box<[usize; 256]>, ml_counts: Box<[usize; 256]>, of_counts: Box<[usize; 256]>, - sequences: Vec, + sequences: Vec, } impl EstimatorWorkspace { @@ -824,8 +838,7 @@ impl EstimatorWorkspace { /// to. All four boxes are always present once the workspace exists. fn heap_size(&self) -> usize { 4 * core::mem::size_of::<[usize; 256]>() - + self.sequences.capacity() - * core::mem::size_of::() + + self.sequences.capacity() * core::mem::size_of::() } } @@ -841,7 +854,7 @@ impl Default for EstimatorWorkspace { } } -/// Dry-run analog of [`encode_block_parts_with_sequence_scratch`]: mirrors the +/// Dry-run analog of [`encode_block_parts`]: mirrors the /// real encoder's `compress_literals` and `choose_table` decisions byte-for-byte /// (same `last_huff_table` lookup, same FSE mode selection, same /// `remember_last_used_tables` mutation), and computes the would-be output size @@ -855,10 +868,21 @@ fn estimate_block_parts_size( raw_sequences: &[RawSequence], workspace: &mut EstimatorWorkspace, ) -> usize { - encode_raw_sequences_into( - raw_sequences, - &mut state.offset_hist, + // The probe cannot fill in place: it walks sub-ranges of the block's + // sequences repeatedly, from a scratch history, while the array itself is + // borrowed immutably by the estimator for the whole search. So it keeps a + // copy — which costs nothing that matters, since block splitting only runs + // from level 11 up and never on the band this array's copy was hurting. + workspace.sequences.clear(); + if workspace.sequences.capacity() < raw_sequences.len() { + workspace + .sequences + .reserve_exact(raw_sequences.len() - workspace.sequences.len()); + } + workspace.sequences.extend_from_slice(raw_sequences); + fill_wire_offsets( &mut workspace.sequences, + &mut state.offset_hist, matches!( state.strategy_tag, crate::encoding::strategy::StrategyTag::Fast @@ -907,7 +931,7 @@ fn estimate_literals_section_bytes( weight_scratch: &mut huff0_encoder::WeightScratch, suspected_incompressible: bool, ) -> usize { - // Mirror `encode_block_parts_with_sequence_scratch` literal-mode branches + // Mirror `encode_block_parts` literal-mode branches // **in the same order**. The disabled gate (negative levels: raw literals, // no Huffman) is checked FIRST exactly as the emitter does. if lit_disabled { @@ -1068,7 +1092,7 @@ fn estimate_literals_section_bytes( } fn estimate_sequences_section_bytes( - sequences: &[crate::blocks::sequence_section::Sequence], + sequences: &[RawSequence], fse_tables: &mut FseTables, ll_counts: &mut [usize; 256], ml_counts: &mut [usize; 256], @@ -1082,7 +1106,7 @@ fn estimate_sequences_section_bytes( for seq in sequences { let (ll, _, ll_bits) = encode_literal_length(seq.ll); let (ml, _, ml_bits) = encode_match_len(seq.ml); - let (of, _, _) = encode_offset(seq.of); + let (of, _, _) = encode_offset(seq.off_base); ll_counts[ll as usize] += 1; ml_counts[ml as usize] += 1; of_counts[of as usize] += 1; @@ -1128,7 +1152,7 @@ fn estimate_sequences_section_bytes( let of_mode = choose_table( of_previous.as_ref(), of_default, - sequences.iter().map(|seq| encode_offset(seq.of).0), + sequences.iter().map(|seq| encode_offset(seq.off_base).0), 8, strategy, of_next, @@ -1160,7 +1184,7 @@ fn estimate_sequences_section_bytes( }; let stream_bytes = (bit_content + padding_bits) / 8; - // Mirror state mutation done by `encode_block_parts_with_sequence_scratch`. + // Mirror state mutation done by `encode_block_parts`. let decisions = [ into_last_used_table(ll_mode), into_last_used_table(ml_mode), @@ -1397,7 +1421,7 @@ fn compressed_literals_header_bytes(lit_size: usize) -> usize { struct SingleSequenceEmitBuffers<'a> { output: &'a mut Vec, compressed: &'a mut Vec, - sequence_scratch: &'a mut Vec, + codes: &'a mut Vec, } fn emit_single_sequence_block( @@ -1405,7 +1429,7 @@ fn emit_single_sequence_block( last_block: bool, source_len: usize, literals: &[u8], - sequences: &[RawSequence], + sequences: &mut [RawSequence], buffers: &mut SingleSequenceEmitBuffers<'_>, ) -> bool { let saved_offset_hist = state.offset_hist; @@ -1426,12 +1450,12 @@ fn emit_single_sequence_block( // the decisions below. Copying them was also what kept the built table's // handle shared, which forced a fresh one per block. buffers.compressed.clear(); - let fse_decisions = encode_block_parts_with_sequence_scratch( + let fse_decisions = encode_block_parts( state, literals, sequences, + buffers.codes, buffers.compressed, - buffers.sequence_scratch, ); let min_gain = (source_len >> 8) + 2; if buffers.compressed.len() >= source_len.saturating_sub(min_gain) { @@ -1479,45 +1503,187 @@ fn emit_single_sequence_block( } } -fn encode_raw_sequences_into( - raw_sequences: &[RawSequence], +/// One sequence's three FSE symbols and the two extra-bit widths that are not +/// derivable from a symbol alone, packed into a word. +/// +/// The derivation pass has all five in hand, and the bit writer needs all five +/// again a moment later; recomputing them there costs two table lookups with +/// their bounds checks, a bit scan and the branches that pick between the +/// small-value tables and the logarithmic form. Upstream keeps the same values +/// between the same two passes, as the three byte arrays `ZSTD_seqToCodes` +/// writes. +/// +/// The offset code is its own extra-bit width, so only two widths are stored. +/// Layout, low bits first: ll code 6, ml code 6, of code 5, ll bits 5, ml bits +/// 5 — 27 bits, and every field's range is fixed by the sequence-section +/// format. +struct SequenceCode(u32); + +impl SequenceCode { + #[inline(always)] + fn pack(ll_code: u8, ml_code: u8, of_code: u8, ll_bits: usize, ml_bits: usize) -> u32 { + debug_assert!(ll_code < 64 && ml_code < 64 && of_code < 32); + debug_assert!(ll_bits < 32 && ml_bits < 32); + ll_code as u32 + | (ml_code as u32) << 6 + | (of_code as u32) << 12 + | (ll_bits as u32) << 17 + | (ml_bits as u32) << 22 + } + + #[inline(always)] + fn ll_code(&self) -> u8 { + (self.0 & 63) as u8 + } + + #[inline(always)] + fn ml_code(&self) -> u8 { + (self.0 >> 6 & 63) as u8 + } + + #[inline(always)] + fn of_code(&self) -> u8 { + (self.0 >> 12 & 31) as u8 + } + + #[inline(always)] + fn ll_bits(&self) -> usize { + (self.0 >> 17 & 31) as usize + } + + #[inline(always)] + fn ml_bits(&self) -> usize { + (self.0 >> 22 & 31) as usize + } + + /// The offset code doubles as its own extra-bit width (upstream's + /// `ofBits = ofCode`). + #[inline(always)] + fn of_bits(&self) -> usize { + self.of_code() as usize + } +} + +/// The three sequence-code histograms, passed as one argument so the pass that +/// fills them stays under the register-pressure of six. +struct SequenceCodeCounts<'a> { + ll: &'a mut [usize; 256], + ml: &'a mut [usize; 256], + of: &'a mut [usize; 256], +} + +/// Fill each sequence's wire offset code in place, advancing the repeat-offset +/// history across the run, and histogram the three code streams while the +/// sequence is in hand. Returns the highest code seen per stream, which the +/// table selector needs and would otherwise find by scanning ~200 always-zero +/// slots. +/// +/// One pass, not two: deriving the offset codes and counting them both read +/// every sequence, and the counting pass was reading back what the filling pass +/// had just written. Upstream splits them (`ZSTD_seqToCodes` then +/// `HIST_countFast_wksp`) because its codes go to three separate byte arrays; +/// ours are already where they belong. +/// +/// `FAST_REPCODE` picks the offBase policy once per block instead of per +/// sequence. Upstream's fast matcher emits only offBase 1 (rep[0] when +/// litLength > 0, rep[1] when litLength == 0 via the secondary-position check) +/// or an explicit offset, and never 2/3; greedy and above search all three +/// repeat offsets, which is what the full `encode_offset_with_history` mirrors. +/// +/// Per PARTITION, not per block: the emitter can write a partition raw, and +/// when it does it restores the history, so the partition after it must be +/// filled from the restored one. Filling here, just before each partition is +/// encoded, is what keeps that true. +fn fill_and_count( + raw_sequences: &mut [RawSequence], + offset_hist: &mut [u32; 3], + counts: SequenceCodeCounts<'_>, + codes: &mut Vec, +) -> (usize, usize, usize) { + let SequenceCodeCounts { + ll: ll_counts, + ml: ml_counts, + of: of_counts, + } = counts; + // Written through the spare capacity rather than pushed: the length is + // known, so a push's capacity test per sequence buys nothing, and resizing + // first would zero the buffer only to overwrite all of it. + codes.clear(); + codes.reserve(raw_sequences.len()); + let code_slots = &mut codes.spare_capacity_mut()[..raw_sequences.len()]; + // The history is rotated by every sequence and read by the next one. Held + // behind the caller's reference it was three stores into the compressor per + // sequence, because the loop also writes through the sequence slice and the + // optimiser would not keep the array in registers across that. A local copy + // written back once is the same three words, moved once. + let mut hist = *offset_hist; + for (slot, seq) in code_slots.iter_mut().zip(raw_sequences.iter_mut()) { + let off_base = if FAST_REPCODE { + encode_offset_with_history_fast(seq.off_base, seq.ll, &mut hist) + } else { + encode_offset_with_history(seq.off_base, seq.ll, &mut hist) + }; + seq.off_base = off_base; + let (ll_code, _, ll_bits) = encode_literal_length(seq.ll); + let (ml_code, _, ml_bits) = encode_match_len(seq.ml); + let (of_code, _, _) = encode_offset(off_base); + slot.write(SequenceCode::pack( + ll_code, ml_code, of_code, ll_bits, ml_bits, + )); + ll_counts[ll_code as usize] += 1; + ml_counts[ml_code as usize] += 1; + of_counts[of_code as usize] += 1; + } + *offset_hist = hist; + // SAFETY: the loop wrote every one of the `raw_sequences.len()` slots it + // took from the spare capacity, which `reserve` above guaranteed. + unsafe { + codes.set_len(raw_sequences.len()); + } + ( + highest_used_code(ll_counts), + highest_used_code(ml_counts), + highest_used_code(of_counts), + ) +} + +/// The highest code with a non-zero count, which the table selector needs and +/// would otherwise find by scanning all 256 slots. +/// +/// Carried as a running maximum through the counting loop until it was three +/// compares a sequence there against one bounded scan a stream a block. The +/// bound is the format's: the three sequence alphabets end at 35, 52 and 31, so +/// nothing above 63 is ever counted. +fn highest_used_code(counts: &[usize; 256]) -> usize { + debug_assert!( + counts[64..].iter().all(|&count| count == 0), + "a sequence code above 63 was counted; the alphabets end at 35 / 52 / 31", + ); + counts[..64] + .iter() + .rposition(|&count| count != 0) + .unwrap_or(0) +} + +/// [`fill_and_count`] without the histogram, for the block-split estimator: it +/// prices sub-ranges repeatedly from a scratch history and counts them itself. +fn fill_wire_offsets( + raw_sequences: &mut [RawSequence], offset_hist: &mut [u32; 3], - out: &mut Vec, fast_repcode: bool, ) { - out.clear(); - // `reserve_exact` argument is the increment over LENGTH, not capacity — - // see `SequencePrefixSums::rebuild` for the full rationale. - if out.capacity() < raw_sequences.len() { - out.reserve_exact(raw_sequences.len() - out.len()); - } - // The strategy branch is hoisted out of the per-sequence loop so the - // offBase-policy choice is paid once per block, not per sequence. Upstream - // zstd's fast matcher emits only offBase 1 (rep[0] when litLength > 0, - // rep[1] when litLength == 0 via the secondary-position check) or an explicit - // offset — it never emits offBase 2/3. greedy+ search all three repeat - // offsets, which is what the full `encode_offset_with_history` mirrors. + // Local copy for the same reason as `fill_and_count`. + let mut hist = *offset_hist; if fast_repcode { - out.extend( - raw_sequences - .iter() - .map(|seq| crate::blocks::sequence_section::Sequence { - ll: seq.ll, - ml: seq.ml, - of: encode_offset_with_history_fast(seq.offset, seq.ll, offset_hist), - }), - ); + for seq in raw_sequences.iter_mut() { + seq.off_base = encode_offset_with_history_fast(seq.off_base, seq.ll, &mut hist); + } } else { - out.extend( - raw_sequences - .iter() - .map(|seq| crate::blocks::sequence_section::Sequence { - ll: seq.ll, - ml: seq.ml, - of: encode_offset_with_history(seq.offset, seq.ll, offset_hist), - }), - ); + for seq in raw_sequences.iter_mut() { + seq.off_base = encode_offset_with_history(seq.off_base, seq.ll, &mut hist); + } } + *offset_hist = hist; } fn clone_fse_tables(fse_tables: &FseTables) -> FseTables { @@ -2242,7 +2408,10 @@ fn commit_last_used_table( } fn encode_sequences( - sequences: &[crate::blocks::sequence_section::Sequence], + sequences: &[RawSequence], + // One packed [`SequenceCode`] per sequence, in the same order, from the + // pass that derived the offset codes. + codes: &[u32], writer: &mut BitWriter<&mut Vec>, ll_mode: &FseTableMode<'_>, ml_mode: &FseTableMode<'_>, @@ -2256,10 +2425,15 @@ fn encode_sequences( mode.as_table(default) } + debug_assert_eq!(codes.len(), sequences.len()); let sequence = sequences[sequences.len() - 1]; - let (ll_code, ll_add_bits, ll_num_bits) = encode_literal_length(sequence.ll); - let (of_code, of_add_bits, of_num_bits) = encode_offset(sequence.of); - let (ml_code, ml_add_bits, ml_num_bits) = encode_match_len(sequence.ml); + let code = SequenceCode(codes[codes.len() - 1]); + let (ll_code, ll_num_bits) = (code.ll_code(), code.ll_bits()); + let (ml_code, ml_num_bits) = (code.ml_code(), code.ml_bits()); + let (of_code, of_num_bits) = (code.of_code(), code.of_bits()); + let ll_add_bits = low_bits(sequence.ll, ll_num_bits); + let ml_add_bits = low_bits(sequence.ml - 3, ml_num_bits); + let of_add_bits = low_bits(sequence.off_base, of_num_bits); let [ll_default, ml_default, of_default] = defaults; let ll_table = mode_table(ll_mode, ll_default); let ml_table = mode_table(ml_mode, ml_default); @@ -2305,11 +2479,22 @@ fn encode_sequences( unsafe { writer.flush_bulk(); } - for sequence in (0..=sequences.len() - 2).rev() { - let sequence = sequences[sequence]; - let (ll_code, ll_add_bits, ll_num_bits) = encode_literal_length(sequence.ll); - let (of_code, of_add_bits, of_num_bits) = encode_offset(sequence.of); - let (ml_code, ml_add_bits, ml_num_bits) = encode_match_len(sequence.ml); + // Walked as a slice rather than by index: the same order, without the + // bounds check and the index arithmetic that `sequences[i]` pays on + // every sequence. The last one is coded through the FSE init states + // above, so it is not in this range. + for (&sequence, &code) in sequences[..sequences.len() - 1] + .iter() + .zip(codes[..codes.len() - 1].iter()) + .rev() + { + let code = SequenceCode(code); + let (ll_code, ll_num_bits) = (code.ll_code(), code.ll_bits()); + let (ml_code, ml_num_bits) = (code.ml_code(), code.ml_bits()); + let (of_code, of_num_bits) = (code.of_code(), code.of_bits()); + let ll_add_bits = low_bits(sequence.ll, ll_num_bits); + let ml_add_bits = low_bits(sequence.ml - 3, ml_num_bits); + let of_add_bits = low_bits(sequence.off_base, of_num_bits); // State diffs burst: max 30 bits (10+10+9 worst case for // acc_log ≤ 9 ll/ml + acc_log ≤ 8 of) + ≤ 7 leftover from @@ -2660,6 +2845,14 @@ pub(in crate::encoding) fn encode_offset_with_history_fast( actual_offset + 3 } +/// The low `bits` bits of `value`, the extra-bit field the sequence section +/// carries beside a code. Every baseline in the format is a multiple of its +/// code's field width, so masking is the subtraction of the baseline. +#[inline(always)] +fn low_bits(value: u32, bits: usize) -> u32 { + value & ((1u32 << bits) - 1) +} + fn encode_offset(len: u32) -> (u8, u32, usize) { let log = len.ilog2(); let lower = len & ((1 << log) - 1); diff --git a/zstd/src/encoding/blocks/compressed/tests.rs b/zstd/src/encoding/blocks/compressed/tests.rs index e1824121e..a09228bdf 100644 --- a/zstd/src/encoding/blocks/compressed/tests.rs +++ b/zstd/src/encoding/blocks/compressed/tests.rs @@ -270,8 +270,8 @@ fn decide_huff_reuse_prefer_repeat_forces_reuse_for_fast_band() { #[test] fn estimator_literals_section_mirrors_emit_for_short_inputs() { use super::{ - CompressedBlockScratch, EntropyOnlyMatcher, EstimatorWorkspace, - encode_block_parts_with_sequence_scratch, estimate_block_parts_size, + CompressedBlockScratch, EntropyOnlyMatcher, EstimatorWorkspace, encode_block_parts, + estimate_block_parts_size, }; // For each strategy at boundary literal lengths around `min_lits` // and across the all-identical RLE pre-check (fires for any @@ -383,13 +383,12 @@ fn estimator_literals_section_mirrors_emit_for_short_inputs() { let mut workspace = EstimatorWorkspace::default(); let est = estimate_block_parts_size(&mut est_state, &literals, &[], &mut workspace); let mut emitted: Vec = Vec::new(); - let mut scratch: Vec = Vec::new(); - encode_block_parts_with_sequence_scratch( + encode_block_parts( &mut emit_state, &literals, - &[], + &mut [], + &mut Vec::new(), &mut emitted, - &mut scratch, ); assert_eq!( est, @@ -409,8 +408,8 @@ fn estimator_literals_section_mirrors_emit_for_short_inputs() { #[test] fn a_section_with_flat_ends_costs_what_the_emitter_writes_for_it() { use super::{ - CompressedBlockScratch, EntropyOnlyMatcher, EstimatorWorkspace, - encode_block_parts_with_sequence_scratch, estimate_block_parts_size, + CompressedBlockScratch, EntropyOnlyMatcher, EstimatorWorkspace, encode_block_parts, + estimate_block_parts_size, }; // The shape the end-sample shortcut exists for, and the one where the // estimator and the emitter can disagree: both ends look random, the @@ -450,13 +449,12 @@ fn a_section_with_flat_ends_costs_what_the_emitter_writes_for_it() { let mut workspace = EstimatorWorkspace::default(); let est = estimate_block_parts_size(&mut est_state, &literals, &[], &mut workspace); let mut emitted: Vec = Vec::new(); - let mut scratch: Vec = Vec::new(); - encode_block_parts_with_sequence_scratch( + encode_block_parts( &mut emit_state, &literals, - &[], + &mut [], + &mut Vec::new(), &mut emitted, - &mut scratch, ); assert_eq!( @@ -486,6 +484,117 @@ fn encode_match_len_uses_correct_upper_range_base() { assert_eq!(encode_match_len(131074), (52, 65535, 16)); } +/// The scratch is taken and put back around a block, so every buffer it keeps +/// is retained allocation a context reports through `ZSTD_sizeof_CCtx`. The +/// per-sequence code buffer is one of them, and a caller budgeting memory sees +/// whatever this sum leaves out. +#[test] +fn retained_heap_size_counts_the_sequence_code_buffer() { + let mut scratch = super::CompressedBlockScratch::new(); + let before = scratch.retained_heap_size(); + let codes = 1024; + scratch.sequence_codes.reserve_exact(codes); + let after = scratch.retained_heap_size(); + assert!( + after - before >= codes * core::mem::size_of::(), + "reserving {codes} sequence codes grew the reported retained size by only {} bytes", + after - before, + ); +} + +/// The estimator prices a block the splitter is thinking about; the emitter +/// writes the one it chose. They walk the sequences by different routes: the +/// emitter derives each offset code, packs it with the extra-bit widths and +/// runs the FSE writer, while the estimator counts from a scratch history and +/// prices the streams from a per-symbol cost model. The parity tests above use +/// empty sequence arrays; this one carries sequences through both. +/// +/// The two agree EXACTLY on the literals section, whose Huffman code lengths +/// are known per symbol. They cannot on the sequences section, and are not +/// meant to: the FSE cost model prices a symbol at its average width from the +/// normalised probability (upstream does the same in `ZSTD_fseBitCost`), where +/// the writer pays what the state trajectory actually costs. What must hold is +/// that the two stay within that model's rounding of each other, and that they +/// walked the same sequences — a wrong offset code on either side moves the +/// histogram, the table choice and the price by far more than rounding. +#[test] +fn estimator_and_emitter_agree_on_a_block_with_sequences() { + use super::{ + CompressedBlockScratch, EntropyOnlyMatcher, EstimatorWorkspace, encode_block_parts, + estimate_block_parts_size, + }; + // Enough sequences for the section to carry real FSE tables rather than + // the degenerate single-symbol shapes, with repeats among the offsets so + // the repeat-offset codes are exercised beside the explicit ones. + let literals: Vec = (0..512u32).map(|i| (i % 251) as u8).collect(); + let sequences: Vec = (0..64u32) + .map(|i| RawSequence { + ll: i % 8, + ml: 4 + i % 13, + off_base: 1 + (i % 5) * 7, + }) + .collect(); + + for strat in [StrategyTag::Fast, StrategyTag::Lazy, StrategyTag::BtUltra2] { + let make_state = || CompressState:: { + matcher: EntropyOnlyMatcher, + copy_tier: crate::decoding::simd_copy::ExactCopyTier::resolve(), + last_huff_table: None, + huff_table_spare: None, + huff_rollback: None, + huff_weights: Default::default(), + seen_content: Default::default(), + fse_tables: FseTables::new(), + block_scratch: CompressedBlockScratch::new(), + offset_hist: [1, 4, 8], + strategy_tag: strat, + pre_split: None, + huf_optimal_search: true, + literal_compression_disabled: false, + }; + let mut est_state = make_state(); + let mut emit_state = make_state(); + + let mut workspace = EstimatorWorkspace::default(); + let est = estimate_block_parts_size(&mut est_state, &literals, &sequences, &mut workspace); + + let mut emitted: Vec = Vec::new(); + let mut to_emit = sequences.clone(); + encode_block_parts( + &mut emit_state, + &literals, + &mut to_emit, + &mut Vec::new(), + &mut emitted, + ); + + // Both sides must have advanced the repeat-offset history the same + // way. This is the assertion that catches a wrong offset code: the + // byte counts could coincide, the histories cannot. + assert_eq!( + est_state.offset_hist, emit_state.offset_hist, + "estimator and emitter disagree on the repeat-offset history at {strat:?}", + ); + // The block really did carry a sequence section — otherwise the rest + // of this test would be passing on the literals path alone. + assert!(emitted.len() > literals.len() / 2); + // One byte per sixteen sequences plus four, against a model that + // rounds each symbol's width down from a 256-scale log table. The + // measured divergence for this fixture is two bytes on every strategy; + // the bound leaves room for the rounding to accumulate without letting + // a real divergence through, which moves the price by tens of bytes. + let allowed = 4 + sequences.len() / 16; + let diff = est.abs_diff(emitted.len()); + assert!( + diff <= allowed, + "estimator priced {est} bytes against {} emitted for {} sequences at {strat:?}: \ + off by {diff}, more than the {allowed} the cost model can round away", + emitted.len(), + sequences.len(), + ); + } +} + #[test] fn raw_partition_fallback_restores_repeat_offset_history() { let mut state = CompressState { @@ -505,26 +614,26 @@ fn raw_partition_fallback_restores_repeat_offset_history() { literal_compression_disabled: false, }; let source = [0xA5; 8]; - let sequences = [RawSequence { + let mut sequences = [RawSequence { ll: 0, ml: 5, - offset: 20, + off_base: 20, }]; let mut output = Vec::new(); let mut compressed_scratch = Vec::new(); - let mut sequence_scratch = Vec::new(); + let mut code_scratch = Vec::new(); let mut emit_buffers = super::SingleSequenceEmitBuffers { output: &mut output, compressed: &mut compressed_scratch, - sequence_scratch: &mut sequence_scratch, + codes: &mut code_scratch, }; let emitted_raw = emit_single_sequence_block( &mut state, true, source.len(), &[], - &sequences, + &mut sequences, &mut emit_buffers, ); if emitted_raw { diff --git a/zstd/src/encoding/simple/fast_kernel/hash_table.rs b/zstd/src/encoding/simple/fast_kernel/hash_table.rs index 94a75d3bf..3e0bcdbf9 100644 --- a/zstd/src/encoding/simple/fast_kernel/hash_table.rs +++ b/zstd/src/encoding/simple/fast_kernel/hash_table.rs @@ -231,6 +231,41 @@ impl FastHashTable { self.bias = 0; } + /// Slide every stored position down by `drop_n`, the way upstream's + /// `ZSTD_reduceIndex` does when the window moves: the table keeps naming + /// the same bytes, and a position that named part of the dropped prefix + /// becomes the empty sentinel. + /// + /// This is what a window slide costs when the caller does not rebuild the + /// table: one pass over the table's own entries, against a pass over every + /// byte the window kept. It also leaves the table holding exactly what it + /// held before the slide, where a rebuild adds entries for positions the + /// matcher had skipped and never indexed. + /// + /// `saturating_sub` is the semantics here, not a guard: below the bias a + /// slot is already empty, and at or under `drop_n` the position it names is + /// gone. Both floor at the sentinel, and both floor at the same place, so + /// the two are one subtraction of their sum — which is what the loop does, + /// since it runs once per table entry and both terms are fixed for the + /// whole slide. + pub(crate) fn reduce_indices(&mut self, drop_n: u32) { + // Plain `+`: [`Self::advance_epoch`] keeps `bias <= u32::MAX - 2^31`, + // and `drop_n` is a length inside a history the matcher caps at + // `2 * max_window_size <= 2^31`, so the sum lands at `u32::MAX` at the + // very most. Asserted rather than assumed, because a wrap here would + // not fail — it would quietly resurrect dropped positions. + debug_assert!( + self.bias.checked_add(drop_n).is_some(), + "epoch bias {} plus a {drop_n}-byte slide overflows a stored position", + self.bias, + ); + let correction = self.bias + drop_n; + for slot in self.table.iter_mut() { + *slot = slot.saturating_sub(correction); + } + self.bias = 0; + } + /// Continue-mode frame reset (upstream zstd `ZSTD_continueCCtx` cadence): keep /// the table contents and advance the epoch bias past every entry the /// previous frame stored, so all of them read back as the empty diff --git a/zstd/src/encoding/simple/fast_matcher.rs b/zstd/src/encoding/simple/fast_matcher.rs index bbd0a7677..3e630bd21 100644 --- a/zstd/src/encoding/simple/fast_matcher.rs +++ b/zstd/src/encoding/simple/fast_matcher.rs @@ -963,16 +963,34 @@ impl FastKernelMatcher { /// 2. Reset `prefix_start_index` to `INITIAL_PREFIX_START_INDEX = 1` /// — drain re-indexes the retained tail; the sentinel-0 /// filter restores via this fixed baseline. - /// 3. Clear the hash table — entries hold pre-drain absolute - /// positions that no longer reference live bytes. + /// 3. Slide the hash table's stored positions down by `drop_n` + /// ([`FastHashTable::reduce_indices`], upstream `ZSTD_reduceIndex`) + /// so they keep naming the same bytes, with anything that named + /// the dropped prefix falling to the empty sentinel. /// 4. `saturating_sub` `last_block_start` by `drop_n`. - /// 5. Rehash retained tail starting at the sentinel-0 floor - /// ([`INITIAL_PREFIX_START_INDEX`] = 1) so block N+1 can find - /// matches against the kept bytes (without this they'd be - /// "dead history" — visible in the Vec but unlookupable). - /// Starting from index 1 instead of 0 avoids hashing a position - /// that the kernel's `match_idx >= prefix_start_index` filter - /// would reject anyway. + /// + /// Step 3 replaced a clear plus a dense rehash of the whole retained + /// tail, which was a pass over every byte the window kept — at level 1 + /// on an 8 MiB input the rehash alone was a fifth of the encode, for a + /// window that slides every `max_window_size` bytes. Sliding the indices + /// is a pass over the table instead, and it is also the more faithful + /// state: the rehash indexed every position, including the ones the + /// matcher's step had skipped and never stored, so the table came out of + /// a slide holding more than it held going in. + /// + /// The two costs scale differently — the slide with the table's entries, + /// the rehash with the window's bytes — so the choice between them would + /// matter if a table could be much larger than the window it indexes. It + /// cannot: a frame without a dictionary caps `hash_log` at + /// `window_log + 1` (upstream `ZSTD_adjustCParams_internal`), which bounds + /// the table at two entries per window byte, and a dictionary frame takes + /// its table width from the dictionary's own cParams rather than from a + /// caller's `hashLog`. Measured over `windowLog` 10 to 16 with `hashLog` + /// pinned at 20 (`examples/slide_oversized_table.rs`): identical time + /// without a dictionary, and up to twice as fast with one. So there is no + /// size-dependent choice here to make, and adding one would buy a branch + /// and two code paths for a configuration the parameter resolution does + /// not produce. fn drain_real_prefix(&mut self, drop_n: usize) { let drain_end = HISTORY_DRAIN_BASE + drop_n; self.history.drain(HISTORY_DRAIN_BASE..drain_end); @@ -991,12 +1009,11 @@ impl FastKernelMatcher { // prefix floor reverts to the windowed value (upstream zstd // `ZSTD_window_enforceMaxDist` zeroing `loadedDictEnd`). self.loaded_dict_end = 0; - self.hash_table.clear(); + // `drop_n` is a length inside `history`, which the window ceiling + // (`window_log <= 30`, retained at most `2 * max_window_size`) keeps + // under 2^31. + self.hash_table.reduce_indices(drop_n as u32); self.last_block_start = self.last_block_start.saturating_sub(drop_n); - // Skip position 0 — `prefix_start_index = 1` means the kernel - // rejects any match resolving to index 0, so populating that - // slot would just pollute the table with an unreachable entry. - self.prime_hash_table_for_range(INITIAL_PREFIX_START_INDEX as usize); } /// Internal: drain `self.pending` into `self.history`, applying @@ -1925,12 +1942,15 @@ impl FastKernelMatcher { }}; } // Dense is the ordinary path — every searched block primes through it — - // and it has to stay a plain counted loop. Running it as `step_by(1)` - // over an inclusive range keeps the iterator's stride and exhausted - // flag live, which the optimiser does not reduce back to an induction - // variable: that alone cost a factor of two on every fast-level frame. + // and it has to stay a plain counted loop. Two shapes have cost it that: + // running it as `step_by(1)` keeps the iterator's stride live, and an + // INCLUSIVE range keeps its exhausted flag live. Neither is reduced back + // to an induction variable, and both show up as `next` / `lt` frames + // inside this function's profile. The half-open range is the one that + // compiles to a counted loop. The bound cannot overflow: `last_hashable` + // is `history.len() - HASH_READ_SIZE`. if step == 1 { - for pos in range_start..=last_hashable { + for pos in range_start..last_hashable + 1 { index_at!(pos); } return; @@ -1949,11 +1969,11 @@ impl FastKernelMatcher { // avoid. const SEAM: usize = 8; let head_end = last_hashable.min(range_start + SEAM); - for pos in range_start..=head_end { + for pos in range_start..head_end + 1 { index_at!(pos); } let tail_start = last_hashable.saturating_sub(SEAM).max(range_start); - for pos in tail_start..=last_hashable { + for pos in tail_start..last_hashable + 1 { index_at!(pos); } } diff --git a/zstd/src/huff0/huff0_encoder.rs b/zstd/src/huff0/huff0_encoder.rs index c349e60f8..941529576 100644 --- a/zstd/src/huff0/huff0_encoder.rs +++ b/zstd/src/huff0/huff0_encoder.rs @@ -637,8 +637,12 @@ impl HuffmanTable { scratch: &mut WeightScratch, ) -> Self { if use_search { + // Validated by the delegate, which is an entry point of its own. + // Checking here as well would walk the histogram a second time on + // every searched build, including each splitter candidate. Self::build_from_counts(counts) } else { + assert_histogram_fits_nodes(counts); // Match upstream's cheap path: tableLog = FSE_optimalTableLog(11, // srcSize, maxSV, minus=1) (huf_compress.c:1286), height-limit to it, // not the raw natural height (11) which can cost a few bytes vs C. @@ -649,7 +653,7 @@ impl HuffmanTable { } pub fn build_from_counts(counts: &[usize]) -> Self { - assert!(counts.len() <= 256); + assert_histogram_fits_nodes(counts); let symbol_cardinality = counts.iter().filter(|&&count| count > 0).count(); if symbol_cardinality <= 1 { return Self::build_from_weights(&build_limited_weights(counts, 11)); @@ -1193,12 +1197,60 @@ fn huffman_weight_sum_is_power_of_two(weights: &[usize]) -> bool { sum.is_power_of_two() } -#[derive(Clone)] +/// A leaf or internal node of the Huffman tree, sized like upstream's +/// `nodeElt`: the table is `2 * leaf_count - 1` entries and is walked several +/// times per block, so its width is what decides whether it stays in L1 +/// alongside the histogram. In `usize` fields it was forty bytes a node, +/// twenty kilobytes for a full alphabet. +/// +/// Every field's range is bounded by the block: counts sum to the literal +/// count, symbols index a 256-entry alphabet, and node indices reach +/// `2 * 256 - 1`. [`NO_PARENT`] is the root's parent. +#[derive(Clone, Copy)] struct HuffNode { - count: usize, - symbol: usize, - parent: Option, - nb_bits: usize, + count: u32, + symbol: u16, + parent: u16, + nb_bits: u8, +} + +/// `parent` of a node with no parent yet, and of the root once the tree is +/// built. Out of range for a real node index, which is under `2 * 256 - 1`. +const NO_PARENT: u16 = u16::MAX; + +/// Refuse a histogram the tree cannot describe, before anything reads it. +/// +/// The encoder's own inputs always fit: a literals section is at most 128 KiB +/// over at most 256 symbols. The entry points are public, though, and a +/// histogram outside those bounds would not fail — it would build a wrong +/// tree. Counts wider than a node's `u32` truncate on the way into a leaf and +/// overflow at the first merge that crosses the boundary, and one of exactly +/// `u32::MAX` is indistinguishable from the sentinel marking a node the tree +/// has not built yet, which the merge loop would then take as a child. More +/// than 256 symbols overruns what the weight buffers and the node indices are +/// sized for. +/// +/// Called at the entry points rather than at the narrowing itself, because +/// what runs in between reads the histogram too: the cheap path's table-log +/// pick sums the counts in a `usize`, which on a 32-bit target overflows on +/// the same input this exists to reject, before the tree is ever built. +/// +/// The total accumulates in `u64` so the bound reads the same on 32- and +/// 64-bit targets, and saturates rather than wrapping — a total that saturates +/// is far past the bound and is refused either way. +fn assert_histogram_fits_nodes(counts: &[usize]) { + assert!( + counts.len() <= MAX_HUFFMAN_ALPHABET, + "histogram has {} symbols, more than the {MAX_HUFFMAN_ALPHABET} a Huffman table describes", + counts.len(), + ); + let total = counts + .iter() + .fold(0u64, |sum, &count| sum.saturating_add(count as u64)); + assert!( + total < u32::MAX as u64, + "symbol counts sum to {total}, which a tree node's count cannot hold", + ); } /// Build the count-sorted Huffman leaves with their natural (unlimited) code @@ -1218,6 +1270,12 @@ fn build_huffman_leaf_depths(counts: &[usize]) -> Vec { /// that builds a tree per block reuses one allocation instead of taking a fresh /// one every time. fn build_huffman_leaf_depths_into(counts: &[usize], nodes: &mut Vec) { + // Held by construction: both entry points run `assert_histogram_fits_nodes` + // before anything reads the histogram. + debug_assert!( + counts.iter().map(|&count| count as u64).sum::() < u32::MAX as u64, + "histogram reached the tree builder without passing the entry check", + ); let leaf_count = counts.iter().filter(|&&count| count > 0).count(); // Pre-size to the final node count (`2 * leaf_count - 1`) so the tree // build's resize never reallocates. @@ -1230,9 +1288,9 @@ fn build_huffman_leaf_depths_into(counts: &[usize], nodes: &mut Vec) { if leaf_count == 1 { let (symbol, &count) = counts.iter().enumerate().find(|&(_, &c)| c > 0).unwrap(); nodes.push(HuffNode { - count, - symbol, - parent: None, + count: count as u32, + symbol: symbol as u16, + parent: NO_PARENT, nb_bits: 0, }); return; @@ -1278,7 +1336,7 @@ fn build_huffman_leaf_depths_into(counts: &[usize], nodes: &mut Vec) { HuffNode { count: 0, symbol: 0, - parent: None, + parent: NO_PARENT, nb_bits: 0, }, ); @@ -1290,9 +1348,9 @@ fn build_huffman_leaf_depths_into(counts: &[usize], nodes: &mut Vec) { let pos = cursor[bucket] as usize; cursor[bucket] += 1; nodes[pos] = HuffNode { - count, - symbol, - parent: None, + count: count as u32, + symbol: symbol as u16, + parent: NO_PARENT, nb_bits: 0, }; } @@ -1309,12 +1367,16 @@ fn build_huffman_leaf_depths_into(counts: &[usize], nodes: &mut Vec) { } } + // The unbuilt internal nodes start at the maximum count so the merge loop + // below never selects one: it always takes the smaller of the next leaf and + // the next built node, and a count no real subtree can reach is what keeps + // an unbuilt slot out of that comparison. nodes.resize( 2 * leaf_count - 1, HuffNode { - count: usize::MAX, - symbol: usize::MAX, - parent: None, + count: u32::MAX, + symbol: u16::MAX, + parent: NO_PARENT, nb_bits: 0, }, ); @@ -1326,13 +1388,13 @@ fn build_huffman_leaf_depths_into(counts: &[usize], nodes: &mut Vec) { // Plain `+`: node counts are symbol frequencies whose tree-wide sum is the // block's symbol count (<= MAX_BLOCK_SIZE), so a merged parent count cannot - // overflow usize. Saturation would only mask a corrupt frequency table. + // overflow u32. Saturation would only mask a corrupt frequency table. nodes[node_nb].count = nodes[low_s as usize].count + nodes[(low_s - 1) as usize].count; nodes[node_nb].symbol = nodes[(low_s - 1) as usize] .symbol .min(nodes[low_s as usize].symbol); - nodes[low_s as usize].parent = Some(node_nb); - nodes[(low_s - 1) as usize].parent = Some(node_nb); + nodes[low_s as usize].parent = node_nb as u16; + nodes[(low_s - 1) as usize].parent = node_nb as u16; node_nb += 1; low_s -= 2; @@ -1341,7 +1403,7 @@ fn build_huffman_leaf_depths_into(counts: &[usize], nodes: &mut Vec) { let leaf_count = if low_s >= 0 { nodes[low_s as usize].count } else { - usize::MAX + u32::MAX }; let node_count = nodes[low_n].count; if leaf_count < node_count { @@ -1358,7 +1420,7 @@ fn build_huffman_leaf_depths_into(counts: &[usize], nodes: &mut Vec) { let leaf_count = if low_s >= 0 { nodes[low_s as usize].count } else { - usize::MAX + u32::MAX }; let node_count = nodes[low_n].count; if leaf_count < node_count { @@ -1374,17 +1436,17 @@ fn build_huffman_leaf_depths_into(counts: &[usize], nodes: &mut Vec) { // Plain `+`: see the leaf-merge above — counts sum to <= MAX_BLOCK_SIZE. nodes[node_nb].count = nodes[first].count + nodes[second].count; nodes[node_nb].symbol = nodes[first].symbol.min(nodes[second].symbol); - nodes[first].parent = Some(node_nb); - nodes[second].parent = Some(node_nb); + nodes[first].parent = node_nb as u16; + nodes[second].parent = node_nb as u16; node_nb += 1; } for leaf_idx in 0..leaf_count { - let mut depth = 0usize; + let mut depth = 0u8; let mut parent = nodes[leaf_idx].parent; - while let Some(parent_idx) = parent { + while parent != NO_PARENT { depth += 1; - parent = nodes[parent_idx].parent; + parent = nodes[parent as usize].parent; } nodes[leaf_idx].nb_bits = depth; } @@ -1418,7 +1480,7 @@ fn limited_weights_into( out.resize(counts_len, 0); if leaves.len() <= 1 { if let Some(leaf) = leaves.first() { - out[leaf.symbol] = 1; + out[leaf.symbol as usize] = 1; } return true; } @@ -1433,19 +1495,19 @@ fn limited_weights_into( // upstream), leaving the code under- or over-full, which projects to a // non-power-of-two weight sum the table builder rejects; detect that here // and fall back to the distributed-weight construction. - if work.iter().any(|leaf| leaf.nb_bits > max_nb_bits) { + if work.iter().any(|leaf| leaf.nb_bits as usize > max_nb_bits) { return false; } let kraft_sum = work .iter() - .map(|leaf| 1usize << (max_nb_bits - leaf.nb_bits)) + .map(|leaf| 1usize << (max_nb_bits - leaf.nb_bits as usize)) .sum::(); if kraft_sum != 1usize << max_nb_bits { return false; } for leaf in work.iter() { - out[leaf.symbol] = max_nb_bits - leaf.nb_bits + 1; + out[leaf.symbol as usize] = max_nb_bits - leaf.nb_bits as usize + 1; } true } @@ -1547,7 +1609,7 @@ fn legacy_distributed_weights(counts: &[usize]) -> Vec { } fn enforce_max_height(nodes: &mut [HuffNode], target_nb_bits: usize) { - let Some(largest_bits) = nodes.iter().map(|node| node.nb_bits).max() else { + let Some(largest_bits) = nodes.iter().map(|node| node.nb_bits as usize).max() else { return; }; if largest_bits <= target_nb_bits { @@ -1569,15 +1631,15 @@ fn enforce_max_height(nodes: &mut [HuffNode], target_nb_bits: usize) { let base_cost = 1usize << (largest_bits - target_nb_bits); let mut total_cost = 0isize; let mut n = nodes.len() - 1; - while nodes[n].nb_bits > target_nb_bits { - total_cost += (base_cost - (1usize << (largest_bits - nodes[n].nb_bits))) as isize; - nodes[n].nb_bits = target_nb_bits; + while nodes[n].nb_bits as usize > target_nb_bits { + total_cost += (base_cost - (1usize << (largest_bits - nodes[n].nb_bits as usize))) as isize; + nodes[n].nb_bits = target_nb_bits as u8; if n == 0 { break; } n -= 1; } - while n > 0 && nodes[n].nb_bits == target_nb_bits { + while n > 0 && nodes[n].nb_bits as usize == target_nb_bits { n -= 1; } total_cost >>= largest_bits - target_nb_bits; @@ -1594,10 +1656,10 @@ fn enforce_max_height(nodes: &mut [HuffNode], target_nb_bits: usize) { let mut rank_last = [NO_SYMBOL; 14]; let mut current_nb_bits = target_nb_bits; for pos in (0..=n).rev() { - if nodes[pos].nb_bits >= current_nb_bits { + if nodes[pos].nb_bits as usize >= current_nb_bits { continue; } - current_nb_bits = nodes[pos].nb_bits; + current_nb_bits = nodes[pos].nb_bits as usize; rank_last[target_nb_bits - current_nb_bits] = pos; } @@ -1636,7 +1698,7 @@ fn enforce_max_height(nodes: &mut [HuffNode], target_nb_bits: usize) { } else { let next = pos - 1; rank_last[bits_to_decrease] = - if nodes[next].nb_bits == target_nb_bits - bits_to_decrease { + if nodes[next].nb_bits as usize == target_nb_bits - bits_to_decrease { next } else { NO_SYMBOL @@ -1656,7 +1718,7 @@ fn enforce_max_height(nodes: &mut [HuffNode], target_nb_bits: usize) { while total_cost < 0 { if rank_last[1] == NO_SYMBOL { // No rank-1 symbol yet: create one from the largest rank-0 node. - while n > 0 && nodes[n].nb_bits == target_nb_bits { + while n > 0 && nodes[n].nb_bits as usize == target_nb_bits { n -= 1; } // Upstream relies on an over-sized node buffer here and reads into @@ -1664,7 +1726,7 @@ fn enforce_max_height(nodes: &mut [HuffNode], target_nb_bits: usize) { // is no real rank-0 node left to borrow, the distribution is too // degenerate to height-limit, so stop and let the caller fall back // to the distributed-weight construction. - if nodes[n].nb_bits == target_nb_bits || n + 1 >= nodes.len() { + if nodes[n].nb_bits as usize == target_nb_bits || n + 1 >= nodes.len() { break; } nodes[n + 1].nb_bits -= 1; diff --git a/zstd/src/huff0/huff0_encoder/tests.rs b/zstd/src/huff0/huff0_encoder/tests.rs index 2158fde68..58c837739 100644 --- a/zstd/src/huff0/huff0_encoder/tests.rs +++ b/zstd/src/huff0/huff0_encoder/tests.rs @@ -726,3 +726,56 @@ fn a_rejected_description_is_recorded_and_the_raw_form_written() { } assert_eq!(encoded, expected); } + +/// The tree node carries its count in a `u32`, which the encoder's own inputs +/// can never overflow: a literals section is at most 128 KiB, so its counts sum +/// to that. The entry point is public, though, and a caller handing it a +/// histogram whose counts do not fit gets a merge that overflows, a count that +/// truncates on the way into a node, and — at exactly `u32::MAX` — a leaf +/// indistinguishable from the sentinel that marks a node the tree has not built +/// yet. Say so at the boundary rather than let any of the three happen. +#[test] +#[should_panic(expected = "symbol counts sum to")] +fn build_from_counts_rejects_a_histogram_wider_than_a_node_count() { + let counts = [u32::MAX as usize, 1, 1]; + let _ = HuffmanTable::build_from_counts(&counts); +} + +/// The cheap path skips the table-log search but reaches the same tree +/// builder, so it needs the same bound. It also reads the histogram on the way +/// there — its table-log pick sums the counts in a `usize`, which overflows on +/// this input on a 32-bit target — so the bound has to be stated at the entry, +/// not at the narrowing. +#[test] +#[should_panic(expected = "symbol counts sum to")] +fn build_from_counts_gated_rejects_a_histogram_wider_than_a_node_count() { + let counts = [u32::MAX as usize, 1, 1]; + let _ = HuffmanTable::build_from_counts_gated(&counts, false); +} + +/// The alphabet bound belongs at the entry too: the search path asserted it, +/// the cheap path did not, and the weight buffers and node indices are sized +/// for 256 symbols on both. +#[test] +#[should_panic(expected = "more than the 256")] +fn build_from_counts_gated_rejects_an_alphabet_wider_than_a_huffman_table() { + let counts = alloc::vec![1usize; 257]; + let _ = HuffmanTable::build_from_counts_gated(&counts, false); +} + +/// The bound is on the SUM, and the largest histogram the encoder can produce +/// has to stay well inside it: a full 128 KiB literals section over one symbol. +#[test] +fn build_from_counts_accepts_the_largest_section_the_encoder_can_produce() { + let mut counts = [0usize; 256]; + counts[b'a' as usize] = 128 * 1024 - 2; + counts[b'b' as usize] = 1; + counts[b'c' as usize] = 1; + let table = HuffmanTable::build_from_counts(&counts); + // Three symbols so the depths can differ at all — a two-symbol alphabet + // gives both a one-bit code whatever their counts. The frequent one taking + // the shorter code is the tree having been built from these counts rather + // than from truncated ones. + assert!(table.codes[b'a' as usize].1 < table.codes[b'b' as usize].1); + assert_eq!(table.codes[b'b' as usize].1, table.codes[b'c' as usize].1); +}