diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4c193e61..e7bc73dd 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -353,7 +353,9 @@ jobs: # read as docs-only. needs: [changes, lint] if: ${{ !cancelled() && needs.changes.result == 'success' && needs.lint.result != 'failure' }} - timeout-minutes: 15 + # The lazy search is monomorphised per (row log, key width) and SIMD tier, + # so the crate's test build is serial-bound; Windows builds it three times. + timeout-minutes: 25 strategy: # Each OS hits its own runner-image regressions (macos-latest currently # ships Homebrew rustup without `cargo`/`rustc` shims under diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index f3b0205d..a7a2b609 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -41,6 +41,12 @@ jobs: languages: ${{ matrix.language }} # Match Default setup without running contributor-controlled builds. build-mode: none + # The codec's hash salts are fixed mixing constants, not secrets, and + # nothing here is cryptographic, so this query only reports them. + config: | + query-filters: + - exclude: + id: rust/hard-coded-cryptographic-value - name: Perform CodeQL Analysis uses: github/codeql-action/analyze@f3c49fd1ef3281abc41069b3dd22807a60c963f6 # v4 with: diff --git a/ffi-bench/Cargo.toml b/ffi-bench/Cargo.toml index e07f7552..3175aac6 100644 --- a/ffi-bench/Cargo.toml +++ b/ffi-bench/Cargo.toml @@ -135,6 +135,10 @@ name = "dictionary_ffi" path = "tests/dictionary_ffi.rs" required-features = ["bench-internals", "dict-builder"] +[[test]] +name = "dictionary_interop" +path = "tests/dictionary_interop.rs" + [[test]] name = "legacy_trainer_ffi" path = "tests/legacy_trainer_ffi.rs" diff --git a/ffi-bench/tests/dictionary_interop.rs b/ffi-bench/tests/dictionary_interop.rs new file mode 100644 index 00000000..f4e69c57 --- /dev/null +++ b/ffi-bench/tests/dictionary_interop.rs @@ -0,0 +1,246 @@ +//! Dictionary frames against libzstd, across every way this crate makes one. +//! +//! A dictionary frame is only right when encoder and decoder agree on where +//! every dictionary byte sits, and the encoder takes the dictionary in pieces +//! sized by the window, which the API path decides: a one-shot compressor, a +//! stream with or without its size pledged or hinted, a context reused for +//! frame after frame. So the matrix crosses those paths with dictionary +//! lengths just past a whole number of pieces, inputs that reach the +//! dictionary's start and its end, and a representative level of every +//! strategy, and checks each frame through the C decoder and ours. The other +//! direction, libzstd's frames through our decoders, runs over the same +//! dictionaries. + +use std::io::{Read, Write}; + +use structured_zstd::decoding::{DictionaryHandle, FrameDecoder, StreamingDecoder}; +use structured_zstd::encoding::{ + CompressionContext, CompressionLevel, EncoderDictionary, FrameCompressor, StreamingEncoder, +}; + +/// One level per strategy: fast, dfast (two tiers), greedy, lazy, lazy2, +/// btlazy2, btopt, btultra, btultra2 and the top of the scale. +const LEVELS: [i32; 11] = [1, 2, 3, 5, 6, 8, 13, 16, 18, 19, 22]; + +/// Compressible bytes over a small alphabet with runs, so matches fire both +/// within the input and into a dictionary drawn the same way. +fn bytes(seed: u64, len: usize) -> Vec { + let mut state = seed | 1; + let mut out = Vec::with_capacity(len); + while out.len() < len { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + let symbol = b'a' + (state % 40) as u8; + let run = 1 + (state >> 40) as usize % 4; + for _ in 0..run.min(len - out.len()) { + out.push(symbol); + } + } + out +} + +/// The inputs a dictionary is tried with: its first bytes (matches reach the +/// farthest back a frame can look), its last bytes (matches reach the bytes a +/// short last piece holds), and a stream of several blocks that mixes +/// dictionary stretches with fresh bytes. +fn inputs(dict: &[u8]) -> Vec<(&'static str, Vec)> { + let small = dict.len().min(11_000); + let mut mixed = Vec::new(); + let mut seed = 0x51EDu64; + while mixed.len() < 120_000 { + let at = (seed as usize * 7919) % dict.len(); + let take = (dict.len() - at).min(3_000); + mixed.extend_from_slice(&dict[at..at + take]); + mixed.extend_from_slice(&bytes(seed, 1_500)); + seed += 1; + } + vec![ + ("start", dict[..small].to_vec()), + ("end", dict[dict.len() - small..].to_vec()), + ("mixed", mixed), + ] +} + +/// Ways a frame is made here. +#[derive(Clone, Copy, Debug)] +enum Path { + OneShot, + Streamed, + Pledged, + Hinted, + /// Three frames on one context, each with its size pledged; the second and + /// third reuse the dictionary already in place. + Reused, +} + +const PATHS: [Path; 5] = [ + Path::OneShot, + Path::Streamed, + Path::Pledged, + Path::Hinted, + Path::Reused, +]; + +fn stream_frame(context: &mut CompressionContext, data: &[u8], pledge: bool) -> Vec { + let mut encoder = StreamingEncoder::with_context(Vec::new(), context); + if pledge { + encoder.set_pledged_content_size(data.len() as u64).unwrap(); + } + encoder.write_all(data).unwrap(); + encoder.finish().unwrap() +} + +/// The frames `path` makes of `data` at `level` with `dictionary`. +fn compress(path: Path, level: i32, dictionary: &EncoderDictionary, data: &[u8]) -> Vec> { + let level = CompressionLevel::from_level(level); + if let Path::OneShot = path { + let mut compressor: FrameCompressor = FrameCompressor::new(level); + compressor + .set_encoder_dictionary(dictionary.clone()) + .unwrap(); + return vec![compressor.compress_independent_frame(data)]; + } + let mut context = CompressionContext::new(level); + context.set_content_checksum(true).unwrap(); + context.set_encoder_dictionary(dictionary.clone()).unwrap(); + match path { + Path::Streamed => vec![stream_frame(&mut context, data, false)], + Path::Pledged => vec![stream_frame(&mut context, data, true)], + Path::Hinted => { + context.set_source_size_hint(data.len() as u64).unwrap(); + vec![stream_frame(&mut context, data, false)] + } + Path::Reused => (0..3) + .map(|_| stream_frame(&mut context, data, true)) + .collect(), + Path::OneShot => unreachable!("handled above"), + } +} + +/// Decode `frame` through libzstd and through both of our decoders; a +/// description of the first that does not give back `data`. +fn check(frame: &[u8], dict_raw: &[u8], data: &[u8]) -> Option { + let mut c_decoder = zstd::bulk::Decompressor::with_dictionary(dict_raw).unwrap(); + match c_decoder.decompress(frame, data.len()) { + Ok(decoded) if decoded == data => {} + Ok(_) => return Some("libzstd decoded other bytes".into()), + Err(err) => return Some(format!("libzstd refused: {err}")), + } + let handle = DictionaryHandle::from_serialized_or_raw_content(dict_raw).unwrap(); + let mut decoder = FrameDecoder::new(); + let mut out = vec![0u8; data.len()]; + match decoder.decode_all_with_dict_handle(frame, &mut out, &handle) { + Ok(n) if out[..n] == *data => {} + Ok(_) => return Some("FrameDecoder decoded other bytes".into()), + Err(err) => return Some(format!("FrameDecoder refused: {err:?}")), + } + let mut stream = + StreamingDecoder::new_with_decoder_and_dictionary_handle(frame, &mut decoder, &handle) + .unwrap_or_else(|_| panic!("the dictionary is usable")); + let mut streamed = Vec::new(); + match stream.read_to_end(&mut streamed) { + Ok(_) if streamed == data => None, + Ok(_) => Some("StreamingDecoder decoded other bytes".into()), + Err(err) => Some(format!("StreamingDecoder refused: {err}")), + } +} + +/// Every path, input and level for `dict_raw`, ours to libzstd and libzstd to +/// ours; the failures, labelled. +fn matrix(name: &str, dict_raw: &[u8]) -> Vec { + let dictionary = EncoderDictionary::from_serialized_or_raw_content(dict_raw).unwrap(); + let mut failures = Vec::new(); + for (input, data) in inputs(dict_raw) { + for level in LEVELS { + for path in PATHS { + for (frame_no, frame) in + compress(path, level, &dictionary, &data).iter().enumerate() + { + if let Some(why) = check(frame, dict_raw, &data) { + failures.push(format!( + "{name} {input} level {level} {path:?} frame {frame_no}: {why}" + )); + } + } + } + let mut c_encoder = zstd::bulk::Compressor::with_dictionary(level, dict_raw).unwrap(); + c_encoder + .set_parameter(zstd::zstd_safe::CParameter::ChecksumFlag(true)) + .unwrap(); + let theirs = c_encoder.compress(&data).unwrap(); + if let Some(why) = check(&theirs, dict_raw, &data) { + failures.push(format!("{name} {input} level {level} libzstd frame: {why}")); + } + } + } + failures +} + +fn assert_clean(failures: Vec) { + assert!( + failures.is_empty(), + "{} failures:\n{}", + failures.len(), + failures.join("\n") + ); +} + +/// A raw-content dictionary `base + tail` bytes long. +fn raw_dictionary(base: usize, tail: usize) { + let dict = bytes(0xD1C7 ^ (base + tail) as u64, base + tail); + assert_clean(matrix(&format!("raw {base}+{tail}"), &dict)); +} + +/// One test per dictionary length, so the matrix runs in parallel. The bases +/// are whole numbers of every piece size up to them, except 110 KiB (the +/// command line's default dictionary size), a whole number of 1 KiB and 2 KiB +/// pieces only. +macro_rules! raw_dictionary_tests { + ($($name:ident: $base:expr, $tail:expr;)*) => { + $( + #[test] + fn $name() { + raw_dictionary($base, $tail); + } + )* + }; +} + +raw_dictionary_tests! { + raw_4_kib_plus_0: 4 * 1024, 0; + raw_4_kib_plus_1: 4 * 1024, 1; + raw_4_kib_plus_2: 4 * 1024, 2; + raw_4_kib_plus_3: 4 * 1024, 3; + raw_4_kib_plus_4: 4 * 1024, 4; + raw_16_kib_plus_0: 16 * 1024, 0; + raw_16_kib_plus_1: 16 * 1024, 1; + raw_16_kib_plus_2: 16 * 1024, 2; + raw_16_kib_plus_3: 16 * 1024, 3; + raw_16_kib_plus_4: 16 * 1024, 4; + raw_64_kib_plus_0: 64 * 1024, 0; + raw_64_kib_plus_1: 64 * 1024, 1; + raw_64_kib_plus_2: 64 * 1024, 2; + raw_64_kib_plus_3: 64 * 1024, 3; + raw_64_kib_plus_4: 64 * 1024, 4; + raw_110_kib_plus_0: 112_640, 0; + raw_110_kib_plus_1: 112_640, 1; + raw_110_kib_plus_2: 112_640, 2; + raw_110_kib_plus_3: 112_640, 3; + raw_110_kib_plus_4: 112_640, 4; + raw_128_kib_plus_0: 128 * 1024, 0; + raw_128_kib_plus_1: 128 * 1024, 1; + raw_128_kib_plus_2: 128 * 1024, 2; + raw_128_kib_plus_3: 128 * 1024, 3; + raw_128_kib_plus_4: 128 * 1024, 4; +} + +/// A dictionary libzstd trained, so its content length is whatever training +/// left, with entropy tables and repeat offsets of its own. +#[test] +fn trained_dictionary_interoperates() { + let corpus = bytes(0x7EA1, 400_000); + let samples: Vec<&[u8]> = corpus.chunks(1_000).collect(); + let dict = zstd::dict::from_samples(&samples, 100_000).unwrap(); + assert_clean(matrix("trained", &dict)); +} diff --git a/zstd/src/encoding/blocks/compressed.rs b/zstd/src/encoding/blocks/compressed.rs index e3c953d4..194d819e 100644 --- a/zstd/src/encoding/blocks/compressed.rs +++ b/zstd/src/encoding/blocks/compressed.rs @@ -165,11 +165,6 @@ pub(crate) struct CompressedBlockScratch { /// first block-split). The estimator uses `EntropyOnlyMatcher` and /// never re-splits, so this nesting is one level deep. estimator_inner: Option>, - /// Persistent slot for `compress_block_encoded`'s pre-block entropy - /// rollback snapshot. `clone_from` into this slot reuses its `Vec` - /// buffers across blocks; a fresh `.clone()` per block paid a - /// malloc + free pair on both Huffman code containers every block. - pub(crate) huff_rollback: Option, } impl CompressedBlockScratch { @@ -192,10 +187,6 @@ impl CompressedBlockScratch { .estimator_workspace .as_ref() .map_or(0, EstimatorWorkspace::heap_size) - + self - .huff_rollback - .as_ref() - .map_or(0, |table| table.heap_size()) + self.estimator_inner.as_ref().map_or(0, |inner| { core::mem::size_of::() + inner.retained_heap_size() }) @@ -395,21 +386,25 @@ impl Matcher for EntropyOnlyMatcher { } } -/// A block of [`crate::blocks::block::BlockType::Compressed`] -pub fn compress_block(state: &mut CompressState, output: &mut Vec) { +/// A block of [`crate::blocks::block::BlockType::Compressed`]. Its tables are +/// returned for the caller to [`commit_block_tables`] if it keeps the block, +/// or [`discard_block_tables`] if it writes the block raw instead. +#[must_use = "the block's tables have to be committed or discarded"] +pub(crate) fn compress_block( + state: &mut CompressState, + output: &mut Vec, +) -> BlockTables { let mut scratch = core::mem::take(&mut state.block_scratch); collect_block_parts(state, &mut scratch.parts); - let decisions = encode_block_parts( + let tables = encode_block_parts( state, &scratch.parts.literals, &mut scratch.parts.sequences, &mut scratch.sequence_codes, output, ); - // This path writes the block it just encoded, so the tables it chose are - // what the next block reads. - remember_last_used_tables(&mut state.fse_tables, decisions); state.block_scratch = scratch; + tables } pub(crate) fn compress_block_with_post_split( @@ -500,7 +495,6 @@ pub(crate) fn compress_block_with_post_split( // `built_huff`; this slot stays empty. last_huff_table: None, huff_table_spare: None, - huff_rollback: None, // Lent, not created: the estimator builds a table per split // candidate, and a fresh scratch here would take its buffers again // on every post-split block and drop them at the end of it. Handed @@ -710,9 +704,9 @@ fn encode_block_parts( // allocated once rather than per block. codes: &mut RegionVec, output: &mut Vec, - // What each axis decided, for the caller to apply once it knows the block - // is kept. LL, ML, OF. -) -> [LastUsedTable; 3] { + // What the block decided for its tables, for the caller to commit once it + // knows the block is kept. +) -> BlockTables { // 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]; @@ -754,7 +748,10 @@ fn encode_block_parts( // small-input framing economy (single-segment header, store-raw fallback // for non-shrinking blocks) lives in `append_frame_header` / // `compress_block_encoded`, not in this literals path. - if state.literal_compression_disabled { + // What the section does to the Huffman table is returned, not applied: the + // table the block started from stays in place until the caller keeps the + // block, so a block that loses to a raw one has nothing to put back. + let huff = if state.literal_compression_disabled { // Upstream zstd `ZSTD_literalsCompressionIsDisabled` (auto mode: // `strategy == ZSTD_fast && targetLength > 0`, i.e. the negative levels): // emit RAW literals, skipping the Huffman pass entirely @@ -762,10 +759,10 @@ fn encode_block_parts( // C's negative-band frames byte-for-byte (the literals section there is // Raw, not Compressed/RLE). raw_literals(literals_vec, &mut writer); - state.clear_huff_table(); + HuffOutcome::Clear } else if !literals_vec.is_empty() && all_bytes_identical(literals_vec) { rle_literals(literals_vec, &mut writer); - state.clear_huff_table(); + HuffOutcome::Clear } else if literals_vec.len() >= min_lits { match compress_literals( literals_vec, @@ -776,18 +773,14 @@ fn encode_block_parts( &mut state.huff_weights, literals_suspected_incompressible(literals_vec.len(), raw_sequences.len()), ) { - HuffmanTableUpdate::New(table) => { - state.replace_huff_table(table); - } - HuffmanTableUpdate::Reused => {} - HuffmanTableUpdate::Cleared => { - state.clear_huff_table(); - } + HuffmanTableUpdate::New(table) => HuffOutcome::New(table), + HuffmanTableUpdate::Reused => HuffOutcome::Keep, + HuffmanTableUpdate::Cleared => HuffOutcome::Clear, } } else { raw_literals(literals_vec, &mut writer); - state.clear_huff_table(); - } + HuffOutcome::Clear + }; // sequences section @@ -801,9 +794,9 @@ fn encode_block_parts( // passes; folded into one loop here saves the per-element // closure overhead (profile #220 round 3: `Map::fold` + // `call_mut` accounted for ~5% of total bench CPU). - let mut ll_counts = [0usize; 256]; - let mut ml_counts = [0usize; 256]; - let mut of_counts = [0usize; 256]; + let mut ll_counts: SequenceHistogram = [0; SEQUENCE_CODE_SLOTS]; + let mut ml_counts: SequenceHistogram = [0; SEQUENCE_CODE_SLOTS]; + let mut of_counts: SequenceHistogram = [0; SEQUENCE_CODE_SLOTS]; // 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 @@ -913,16 +906,49 @@ fn encode_block_parts( ]; } writer.flush(); - decisions + BlockTables { + fse: decisions, + huff, + } +} + +/// What a block decided for the tables the next block may repeat, held until +/// the caller knows whether the block is kept. Upstream builds into +/// `nextCBlock` and makes it the previous state only on confirmation +/// (`ZSTD_blockState_confirmRepcodesAndEntropyTables`); a block that goes out +/// raw is simply not confirmed, so nothing the block started from is ever +/// copied to be put back. +pub(crate) struct BlockTables { + fse: [LastUsedTable; 3], + huff: HuffOutcome, +} + +/// Make a kept block's tables the ones the next block reads. +pub(crate) fn commit_block_tables(state: &mut CompressState, tables: BlockTables) { + remember_last_used_tables(&mut state.fse_tables, tables.fse); + match tables.huff { + HuffOutcome::Keep => {} + HuffOutcome::Clear => state.clear_huff_table(), + HuffOutcome::New(table) => state.replace_huff_table(table), + } +} + +/// Drop a discarded block's tables, keeping a Huffman table it built as a +/// buffer for the next build; the FSE tables it built stay in their `*_next` +/// slots for the same reason. +pub(crate) fn discard_block_tables(state: &mut CompressState, tables: BlockTables) { + if let HuffOutcome::New(table) = tables.huff { + state.park_huff_table(table); + } } /// Workspace shared across estimator probes so per-probe cost computation never /// allocates. Counts are zeroed at the top of every probe. struct EstimatorWorkspace { lit_counts: Box<[usize; 256]>, - ll_counts: Box<[usize; 256]>, - ml_counts: Box<[usize; 256]>, - of_counts: Box<[usize; 256]>, + ll_counts: Box, + ml_counts: Box, + of_counts: Box, /// FSE table handles no probe state holds any more, for the next probe to /// build into. Every one is unique, so a build writes it in place. spare_fse: Vec, @@ -937,7 +963,8 @@ impl EstimatorWorkspace { /// the pooled tables with their reference counts, and the room kept for /// the probes' Huffman tables. fn heap_size(&self) -> usize { - 4 * core::mem::size_of::<[usize; 256]>() + core::mem::size_of::<[usize; 256]>() + + 3 * core::mem::size_of::() + self.built_huff.capacity() * core::mem::size_of::() + self.spare_fse.capacity() * core::mem::size_of::() + self.spare_fse.len() @@ -965,9 +992,9 @@ impl Default for EstimatorWorkspace { fn default() -> Self { Self { lit_counts: Box::new([0; 256]), - ll_counts: Box::new([0; 256]), - ml_counts: Box::new([0; 256]), - of_counts: Box::new([0; 256]), + ll_counts: Box::new([0; SEQUENCE_CODE_SLOTS]), + ml_counts: Box::new([0; SEQUENCE_CODE_SLOTS]), + of_counts: Box::new([0; SEQUENCE_CODE_SLOTS]), spare_fse: Vec::new(), built_huff: Vec::new(), } @@ -1281,11 +1308,13 @@ fn estimate_sequences_section_bytes( } = workspace; let (ll_counts, ml_counts, of_counts) = (&mut **ll_counts, &mut **ml_counts, &mut **of_counts); debug_assert_eq!(codes.ll.len(), sequences.len()); - let histogram = |codes: &[u8], counts: &mut [usize; 256]| { + let histogram = |codes: &[u8], counts: &mut SequenceHistogram| { counts.fill(0); let mut max = 0u8; for &code in codes { - counts[code as usize] += 1; + debug_assert!(usize::from(code) < SEQUENCE_CODE_SLOTS); + // SAFETY: a length code is at most 52, under the histogram's slots. + unsafe { *counts.get_unchecked_mut(usize::from(code)) += 1 }; max = max.max(code); } max as usize @@ -1298,7 +1327,9 @@ fn estimate_sequences_section_bytes( let mut hist = *offset_hist; let mut count_offset = |off_base: u32| { let (of, _, _) = encode_offset(off_base); - of_counts[of as usize] += 1; + debug_assert!(usize::from(of) < SEQUENCE_CODE_SLOTS); + // SAFETY: an offset code is at most 31, under the histogram's slots. + unsafe { *of_counts.get_unchecked_mut(usize::from(of)) += 1 }; of_max = of_max.max(of as usize); // Upstream zstd: OF code's value equals its additional-bits width. of_bits += of as usize; @@ -1412,7 +1443,7 @@ fn estimate_sequences_section_bytes( /// `encode_sequences` would emit: FSE state transitions + final state flush. fn fse_section_bits_for_mode( mode: &FseTableMode<'_>, - counts: &[usize; 256], + counts: &SequenceHistogram, default: &FSETable, ) -> usize { let max_symbol = counts.iter().rposition(|&c| c > 0).unwrap_or_default(); @@ -1633,24 +1664,10 @@ fn emit_single_sequence_block( buffers: &mut SingleSequenceEmitBuffers<'_>, ) -> bool { let saved_offset_hist = state.offset_hist; - // Copy into the rollback slot rather than a fresh `Option`: the slot keeps - // its buffers between blocks, so this reuses them instead of taking two - // `Vec`s per block. `had_huff_table` carries what an `Option` copy would - // have carried, without discarding the buffer when there is nothing to - // copy. - let had_huff_table = state.last_huff_table.is_some(); - if let Some(src) = &state.last_huff_table { - match &mut state.huff_rollback { - Some(dst) => dst.clone_from(src), - slot => *slot = Some(src.clone()), - } - } - // The FSE tables need no copy: the block builds into the `*_next` slots and - // the previous ones are only read, so rolling back is simply not applying - // the decisions below. Copying them was also what kept the built table's - // handle shared, which forced a fresh one per block. + // The tables are only read here and the block's decisions come back to be + // committed, so writing the partition raw leaves nothing to restore. buffers.compressed.clear(); - let fse_decisions = encode_block_parts( + let tables = encode_block_parts( state, literals, sequences, @@ -1660,27 +1677,7 @@ fn emit_single_sequence_block( let min_gain = (source_len >> 8) + 2; if buffers.compressed.len() >= source_len.saturating_sub(min_gain) { state.offset_hist = saved_offset_hist; - if had_huff_table { - // Swap, not assign: the table this block built goes back into the - // rollback slot and becomes the next block's copy buffer. - core::mem::swap(&mut state.last_huff_table, &mut state.huff_rollback); - } else { - // Nothing was carried in, so whatever this block built is dropped - // from the state; park it for the next build rather than freeing. - state.clear_huff_table(); - } - // A partition that built a table and then chose raw or RLE literals - // parked it on the way there, so the swap above had nothing to give - // back: the slot would go into the next partition empty, which is the - // allocation of both code buffers it exists to avoid — on exactly the - // run of partitions that keeps failing the size test. Take the parked - // table; the slot is asked for one per partition, the spare's other - // reader once per frame. - if state.huff_rollback.is_none() { - state.huff_rollback = state.huff_table_spare.take(); - } - // The FSE decisions are simply not applied, so the previous tables are - // whatever the last kept block left. + discard_block_tables(state, tables); let header = BlockHeader { last_block, block_type: BlockType::Raw, @@ -1691,7 +1688,7 @@ fn emit_single_sequence_block( } else { // The block is kept, so its tables become what the next one reads. // Upstream's confirm step, at the same point in the decision. - remember_last_used_tables(&mut state.fse_tables, fse_decisions); + commit_block_tables(state, tables); let header = BlockHeader { last_block, block_type: BlockType::Compressed, @@ -1764,12 +1761,22 @@ impl SequenceCode { } } +/// Slots in a sequence-code histogram: enough for the widest of the three +/// alphabets, the 53 match-length codes (RFC 8878 3.1.1.3.2.1.1). Upstream's +/// `HIST_countFast_wksp` likewise clears and counts only up to the last code; +/// a byte-wide histogram here zeroed 6 KiB a block for three streams that use +/// under 128 slots between them. +const SEQUENCE_CODE_SLOTS: usize = 64; + +/// Occurrences of each code of one sequence stream. +type SequenceHistogram = [usize; SEQUENCE_CODE_SLOTS]; + /// 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], + ll: &'a mut SequenceHistogram, + ml: &'a mut SequenceHistogram, + of: &'a mut SequenceHistogram, } /// Fill each sequence's wire offset code in place, advancing the repeat-offset @@ -1835,9 +1842,21 @@ fn fill_and_count( 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; + debug_assert!( + usize::from(ll_code.max(ml_code).max(of_code)) < SEQUENCE_CODE_SLOTS, + "the code functions stay inside their alphabets", + ); + // SAFETY: the codes are at most 35 (LL), 52 (ML) and 31 (OF), under the + // histogram's `SEQUENCE_CODE_SLOTS`. The LL and ML bounds hold in release + // because `encode_literal_length` and `encode_match_len` index + // `LL_EXTRA_BITS` (36 entries) and `ML_EXTRA_BITS` (53) with a bounds + // check, so a larger code panics there before reaching this write; the + // OF code is `ilog2` of a `u32`. Those lookups must stay checked. + unsafe { + *ll_counts.get_unchecked_mut(usize::from(ll_code)) += 1; + *ml_counts.get_unchecked_mut(usize::from(ml_code)) += 1; + *of_counts.get_unchecked_mut(usize::from(of_code)) += 1; + } } if !OFFSETS_CODED { *offset_hist = hist; @@ -1883,7 +1902,7 @@ fn fill_and_count( /// scan starts at its own alphabet's last code (`MAX_CODE`: 35, 52 or 31), /// where upstream's `HIST_count_simple` starts its own, so nothing past the /// alphabet is read. -fn highest_used_code(counts: &[usize; 256]) -> usize { +fn highest_used_code(counts: &SequenceHistogram) -> usize { debug_assert!( counts[MAX_CODE + 1..].iter().all(|&count| count == 0), "a code past the alphabet's last ({MAX_CODE}) was counted", @@ -2219,7 +2238,7 @@ impl FseTableMode<'_> { } } -fn entropy_cost(counts: &[usize; 256], max_symbol: usize, total: usize) -> usize { +fn entropy_cost(counts: &SequenceHistogram, max_symbol: usize, total: usize) -> usize { let mut cost = 0usize; for &count in counts.iter().take(max_symbol + 1) { if count == 0 { @@ -2234,7 +2253,11 @@ fn entropy_cost(counts: &[usize; 256], max_symbol: usize, total: usize) -> usize cost >> 8 } -fn cross_entropy_cost(counts: &[usize; 256], max_symbol: usize, table: &FSETable) -> Option { +fn cross_entropy_cost( + counts: &SequenceHistogram, + max_symbol: usize, + table: &FSETable, +) -> Option { let acc_log = table.acc_log(); if acc_log > 8 { return None; @@ -2259,7 +2282,7 @@ fn cross_entropy_cost(counts: &[usize; 256], max_symbol: usize, table: &FSETable Some(cost >> 8) } -fn fse_bit_cost(counts: &[usize; 256], max_symbol: usize, table: &FSETable) -> Option { +fn fse_bit_cost(counts: &SequenceHistogram, max_symbol: usize, table: &FSETable) -> Option { let table_log = table.acc_log() as usize; let table_size = 1usize << table_log; let mut cost = 0usize; @@ -2316,7 +2339,7 @@ fn choose_table_from_counts<'a>( // last symbol in place for its normalize, and restore it (see // `build_seq_ctable`) — no per-table copy. Every selection-time read // re-borrows it immutably; the value is unchanged on return. - counts: &mut [usize; 256], + counts: &mut SequenceHistogram, total: usize, // The highest symbol code with a non-zero count, tracked by the caller as it // builds `counts`. The sequence-code alphabets are tiny (LL <= 35, ML <= 52, diff --git a/zstd/src/encoding/blocks/compressed/tests.rs b/zstd/src/encoding/blocks/compressed/tests.rs index f4522701..44e07b33 100644 --- a/zstd/src/encoding/blocks/compressed/tests.rs +++ b/zstd/src/encoding/blocks/compressed/tests.rs @@ -1,7 +1,8 @@ use super::{ - FseTableMode, LastUsedTable, RawSequence, choose_table_from_counts, emit_single_sequence_block, - encode_literal_length, encode_match_len, encode_offset_with_history, min_gain, - min_literals_to_compress, previous_table, remember_last_used_tables, + FseTableMode, LastUsedTable, RawSequence, SEQUENCE_CODE_SLOTS, SequenceHistogram, + choose_table_from_counts, emit_single_sequence_block, encode_literal_length, encode_match_len, + encode_offset_with_history, min_gain, min_literals_to_compress, previous_table, + remember_last_used_tables, }; use crate::encoding::frame_compressor::{ CompressState, FseTables, PreviousFseTable, SharedFseTable, @@ -22,7 +23,7 @@ fn choose_table<'a>( strategy: StrategyTag, next: &'a mut Option, ) -> FseTableMode<'a> { - let mut counts = [0usize; 256]; + let mut counts: SequenceHistogram = [0; SEQUENCE_CODE_SLOTS]; let mut total = 0usize; let mut max_symbol = 0usize; for symbol in data { @@ -388,7 +389,6 @@ fn estimator_literals_section_mirrors_emit_for_short_inputs() { copy_kernel: crate::encoding::fastpath::select_kernel(), last_huff_table: seed_table.clone(), huff_table_spare: None, - huff_rollback: None, huff_weights: Default::default(), seen_content: Default::default(), fse_tables: FseTables::new(), @@ -406,7 +406,6 @@ fn estimator_literals_section_mirrors_emit_for_short_inputs() { copy_kernel: crate::encoding::fastpath::select_kernel(), last_huff_table: seed_table, huff_table_spare: None, - huff_rollback: None, huff_weights: Default::default(), seen_content: Default::default(), fse_tables: FseTables::new(), @@ -472,7 +471,6 @@ fn a_section_with_flat_ends_costs_what_the_emitter_writes_for_it() { copy_kernel: crate::encoding::fastpath::select_kernel(), last_huff_table: None, huff_table_spare: None, - huff_rollback: None, huff_weights: Default::default(), seen_content: Default::default(), fse_tables: FseTables::new(), @@ -538,7 +536,6 @@ fn retained_heap_size_counts_the_block_buffers() { copy_kernel: crate::encoding::fastpath::select_kernel(), last_huff_table: None, huff_table_spare: None, - huff_rollback: None, huff_weights: Default::default(), seen_content: Default::default(), fse_tables: FseTables::new(), @@ -568,6 +565,88 @@ fn retained_heap_size_counts_the_block_buffers() { ); } +/// The splitter's workspace survives the block that created it, so its count +/// tables are retained allocation the context must report. +#[test] +fn retained_heap_size_counts_the_estimator_workspace() { + let mut scratch = super::CompressedBlockScratch::new(); + let before = scratch.retained_heap_size(); + scratch.estimator_workspace = Some(Default::default()); + let tables = + core::mem::size_of::<[usize; 256]>() + 3 * core::mem::size_of::(); + assert!( + scratch.retained_heap_size() - before >= tables, + "the estimator's {tables} bytes of count tables went unreported", + ); +} + +/// A block written raw is not confirmed: the tables the next block repeats stay +/// those it started from, and a Huffman table it built is kept as a buffer +/// rather than becoming the carried table. +#[test] +fn discarded_block_keeps_the_entry_tables() { + let mut state = CompressState { + matcher: super::EntropyOnlyMatcher, + copy_kernel: crate::encoding::fastpath::select_kernel(), + last_huff_table: Some(huff0_encoder::HuffmanTable::build_from_data(b"aaaabbc")), + huff_table_spare: None, + huff_weights: Default::default(), + seen_content: Default::default(), + fse_tables: FseTables::new(), + block_scratch: super::CompressedBlockScratch::new(), + workspace: Workspace::new(), + offset_hist: [1, 4, 8], + strategy_tag: StrategyTag::Fast, + pre_split: None, + huf_optimal_search: true, + literal_compression_disabled: false, + post_split_allowed: true, + }; + let entry_fse = [(); 3].map(|()| SharedFseTable::new(FSETable::blank())); + state.fse_tables.ll_previous = Some(PreviousFseTable::Custom(entry_fse[0].clone())); + state.fse_tables.ml_previous = Some(PreviousFseTable::Custom(entry_fse[1].clone())); + state.fse_tables.of_previous = Some(PreviousFseTable::Custom(entry_fse[2].clone())); + let entry_codes = state + .last_huff_table + .as_ref() + .unwrap() + .packed_codes() + .to_vec(); + let built = huff0_encoder::HuffmanTable::build_from_data(b"xyzzzzzzzzzz"); + let built_codes = built.packed_codes().to_vec(); + assert_ne!(entry_codes, built_codes); + super::discard_block_tables( + &mut state, + super::BlockTables { + fse: [ + LastUsedTable::Default, + LastUsedTable::Rle(3), + LastUsedTable::Default, + ], + huff: super::HuffOutcome::New(built), + }, + ); + assert_eq!( + state.last_huff_table.as_ref().unwrap().packed_codes(), + entry_codes + ); + assert_eq!( + state.huff_table_spare.as_ref().unwrap().packed_codes(), + built_codes + ); + let previous = [ + &state.fse_tables.ll_previous, + &state.fse_tables.ml_previous, + &state.fse_tables.of_previous, + ]; + for (slot, entry) in previous.into_iter().zip(&entry_fse) { + assert!( + matches!(slot, Some(PreviousFseTable::Custom(t)) if SharedFseTable::ptr_eq(t, entry)), + "a discarded block replaced the FSE table the next block repeats", + ); + } +} + /// A buffer for `count` sequence codes, carved from `ws`, which must outlive it. fn code_buffer(ws: &mut Workspace, count: usize) -> RegionVec { // The block target stands for the trailing part's size in bytes. @@ -688,7 +767,6 @@ fn estimator_and_emitter_agree_on_a_block_with_sequences() { copy_kernel: crate::encoding::fastpath::select_kernel(), last_huff_table: None, huff_table_spare: None, - huff_rollback: None, huff_weights: Default::default(), seen_content: Default::default(), fse_tables: FseTables::new(), @@ -753,7 +831,6 @@ fn raw_partition_fallback_restores_repeat_offset_history() { copy_kernel: crate::encoding::fastpath::select_kernel(), last_huff_table: None, huff_table_spare: None, - huff_rollback: None, huff_weights: Default::default(), seen_content: Default::default(), fse_tables: FseTables::new(), @@ -836,51 +913,6 @@ fn a_predefined_or_rle_block_parks_the_custom_table_it_replaces() { } } -/// A block that loses to a raw one must leave the axis exactly as it found -/// it: the restored table unique, and a buffer free to build into. -/// -/// The caller holds a CLONE of what `previous` was, so restoring only that -/// leaves the same table in both slots — `previous` restored and `next` where -/// confirmation displaced it — and the next block finds its buffer shared and -/// allocates a fresh eleven-kilobyte table. On a run of blocks that keep -/// falling back to raw that is a per-axis allocation per block, which is the -/// cost this pair of slots exists to remove. -#[test] -fn a_block_that_falls_back_to_raw_leaves_its_axis_reusable() { - use crate::encoding::frame_compressor::SharedFseTable; - use crate::fse::fse_encoder::FSETable; - - let mut fse_tables = FseTables::new(); - let before = SharedFseTable::new(FSETable::blank()); - fse_tables.ll_previous = Some(PreviousFseTable::Custom(before)); - // What the caller keeps across the attempt. - let saved = fse_tables.ll_previous.clone(); - - // The attempt builds into the axis's spare slot, then confirms it, which - // displaces the old table into that slot. - fse_tables.ll_next = Some(SharedFseTable::new(FSETable::blank())); - super::commit_last_used_table( - &mut fse_tables.ll_previous, - &mut fse_tables.ll_next, - LastUsedTable::Encoded, - ); - - fse_tables.roll_back_confirmation([saved, None, None]); - - let Some(PreviousFseTable::Custom(restored)) = fse_tables.ll_previous.as_ref() else { - panic!("the restored table must be the custom one the axis started with"); - }; - assert_eq!( - SharedFseTable::strong_count(restored), - 1, - "the restored table must be unique, or the next block cannot build into its slot", - ); - assert!( - fse_tables.ll_next.is_some(), - "the discarded table is exactly the buffer the next block wants", - ); -} - /// A frame start replaces the previous slots, and the table that was there is /// the next frame's buffer — dropping it costs a reused compressor an /// allocation per axis on every frame that builds a custom table. @@ -989,7 +1021,7 @@ fn fast_band_strategies_prefer_repeat_fse_table() { PreviousFseTable::Custom(crate::encoding::frame_compressor::SharedFseTable::new(prev)); let fse_tables = FseTables::new(); // Distribution over symbols {0,1}, both covered by `previous`. - let mut counts = [0usize; 256]; + let mut counts: SequenceHistogram = [0; SEQUENCE_CODE_SLOTS]; counts[0] = 4; counts[1] = 6; let total = 10; diff --git a/zstd/src/encoding/frame_compressor.rs b/zstd/src/encoding/frame_compressor.rs index 9cb1b18a..032f4a93 100644 --- a/zstd/src/encoding/frame_compressor.rs +++ b/zstd/src/encoding/frame_compressor.rs @@ -509,65 +509,6 @@ pub(crate) struct FseTables { } impl FseTables { - /// Undo a block's table confirmation, keeping the tables it built as the - /// next block's buffers. - /// - /// A block that loses to a raw one has already had its tables confirmed - /// into the previous slots, and the caller holds a clone of what they - /// were. Restoring only that clone leaves the axis holding the SAME table - /// in both slots — the restored one in `previous` and the one confirmation - /// displaced into `next` — so the next block finds its buffer shared and - /// allocates eleven kilobytes instead of writing into it, on every block - /// of a run that keeps falling back. - /// - /// Handing the discarded table to `next` fixes both halves at once: the - /// restored table is unique again, and the buffer the block just filled is - /// exactly what the next one wants to build into. - pub(crate) fn roll_back_confirmation(&mut self, saved: [Option; 3]) { - let [saved_ll, saved_ml, saved_of] = saved; - for (previous, next, restored) in [ - (&mut self.ll_previous, &mut self.ll_next, saved_ll), - (&mut self.ml_previous, &mut self.ml_next, saved_ml), - (&mut self.of_previous, &mut self.of_next, saved_of), - ] { - let discarded = core::mem::replace(previous, restored); - // The table the block built, and it is not the one just restored: - // whatever `next` held was the outgoing previous, which is what was - // restored, so it is a duplicate and this one is free and unique. - // - // The exception is `RepeatLast`, where confirmation never replaced - // `previous` at all — the discarded handle IS the restored one. - // Parking it in `next` then leaves the two sharing a handle, and the - // next custom build cannot write into a table `previous` still - // holds, so it allocates another one. - let built = match discarded { - Some(PreviousFseTable::Custom(built)) - if !matches!(previous.as_ref(), Some(PreviousFseTable::Custom(back)) - if SharedFseTable::ptr_eq(&built, back)) => - { - Some(built) - } - _ => None, - }; - match built { - Some(built) => *next = Some(built), - // Nothing new to park. The block settled on a predefined or RLE - // table, or repeated the last one, so confirmation may have - // parked the outgoing custom table here — and if that is the - // table now back in `previous`, holding it here too is what - // keeps the next build from writing into it. - None => { - if let (Some(spare), Some(PreviousFseTable::Custom(back))) = - (next.as_ref(), previous.as_ref()) - && SharedFseTable::ptr_eq(spare, back) - { - *next = None; - } - } - } - } - } - /// Keep the frame's tables as buffers before the next frame overwrites /// them. /// @@ -1345,14 +1286,6 @@ pub(crate) struct CompressState { /// allocations. Without this, every dict-seeded frame whose last block /// ended raw/RLE paid a fresh two-Vec table clone per frame. pub(crate) huff_table_spare: Option, - /// Where a block copies `last_huff_table` before encoding, so it can be put - /// back if the compressed form loses to a raw block. - /// - /// A slot rather than a local because the copy is per block: cloning into a - /// fresh `Option` took two `Vec`s every time, while cloning into one that - /// already holds a table reuses them. Restoring is a swap, which also hands - /// the discarded table back here as the next block's buffer. - pub(crate) huff_rollback: Option, /// The Huffman weight builder's three buffers, kept across blocks and /// frames. The cheap build path takes a tree and two weight buffers per /// call, and it runs once per block plus once per split candidate wherever @@ -1489,19 +1422,14 @@ impl CompressState { } /// Heap bytes the compressor keeps between blocks and frames beyond the - /// match finder: the FSE tables both slots of each axis hold, the rollback - /// slot the emit paths copy a Huffman table into before a block that may not - /// be kept, the block scratch with everything it holds (the splitter's - /// workspace and the nested estimator scratch), and the context workspace - /// the literal and sequence buffers are carved from. + /// match finder: the FSE tables both slots of each axis hold, the block + /// scratch with everything it holds (the splitter's workspace and the + /// nested estimator scratch), and the context workspace the literal and + /// sequence buffers are carved from. /// /// All of it survives a frame, so a caller sizing a context has to see it. pub(crate) fn retained_scratch_heap_size(&self) -> usize { self.fse_tables.heap_size() - + self - .huff_rollback - .as_ref() - .map_or(0, |table| table.heap_size()) + self.block_scratch.retained_heap_size() + self.workspace.heap_bytes() } @@ -1531,7 +1459,7 @@ impl CompressState { /// dropped the previous table on every block, so the builds that followed /// allocated their buffers again. #[inline] - fn park_huff_table(&mut self, table: crate::huff0::huff0_encoder::HuffmanTable) { + pub(crate) fn park_huff_table(&mut self, table: crate::huff0::huff0_encoder::HuffmanTable) { if self.huff_table_spare.is_none() { self.huff_table_spare = Some(table); } else { @@ -1772,7 +1700,6 @@ impl FrameCompressor { copy_kernel: crate::encoding::fastpath::select_kernel(), last_huff_table: None, huff_table_spare: None, - huff_rollback: None, huff_weights: Default::default(), seen_content: Default::default(), fse_tables: FseTables::new(), @@ -2212,7 +2139,6 @@ impl FrameCompressor { copy_kernel: crate::encoding::fastpath::select_kernel(), last_huff_table: None, huff_table_spare: None, - huff_rollback: None, huff_weights: Default::default(), seen_content: Default::default(), fse_tables: FseTables::new(), diff --git a/zstd/src/encoding/levels/fastest.rs b/zstd/src/encoding/levels/fastest.rs index 24ef1af5..3bc259d8 100644 --- a/zstd/src/encoding/levels/fastest.rs +++ b/zstd/src/encoding/levels/fastest.rs @@ -4,7 +4,10 @@ use crate::{ encoding::{ CompressionLevel, Matcher, block_header::BlockHeader, - blocks::{compress_block, compress_block_with_post_split}, + blocks::{ + commit_block_tables, compress_block, compress_block_with_post_split, + discard_block_tables, + }, frame_compressor::CompressState, incompressible::{ block_looks_incompressible, block_looks_incompressible_strict, @@ -223,37 +226,17 @@ pub(crate) fn compress_block_encoded( sink.push(precomputed_checksum.expect("checksum is hashed whenever a sink exists")); } - // Keep rollback snapshots for the oversize fallback path below: - // `compress_block` can mutate entropy/history state before we know - // whether the compressed payload fits `MAX_BLOCK_SIZE`. + // The offset history is the only state `compress_block` advances in + // place; its tables come back to be committed once the size test below + // keeps the block, so a raw fallback has no entropy state to restore. let saved_offset_hist = state.offset_hist; - // Snapshot the Huffman table into the scratch's persistent rollback - // slot: `clone_from` reuses the slot's buffers across blocks (a - // fresh `.clone()` paid a malloc + free pair on both code containers - // every block). FSE previous tables are `SharedFseTable` handles — - // their clone is a refcount bump, no slot needed. - let mut saved_huff_table = core::mem::take(&mut state.block_scratch.huff_rollback); - // Only when there IS a table to snapshot. `Option::clone_from` from a - // `None` source drops what the slot holds, which is every buffer this - // persistent slot exists to keep: a run of blocks that build a table and - // then lose the size test would free and rebuild it every time. - let had_prior_huff_table = state.last_huff_table.is_some(); - if let Some(prior) = state.last_huff_table.as_ref() { - match saved_huff_table.as_mut() { - Some(slot) => slot.clone_from(prior), - None => saved_huff_table = Some(prior.clone()), - } - } - let saved_ll_previous = state.fse_tables.ll_previous.clone(); - let saved_ml_previous = state.fse_tables.ml_previous.clone(); - let saved_of_previous = state.fse_tables.of_previous.clone(); // Compress directly into `output`: reserve the fixed 3-byte block // header, append the payload after it, then backfill the header in // place once its length is known — no temp `Vec`, no extend-copy. let hdr_off = output.len(); output.extend_from_slice(&[0u8; 3]); let payload_off = output.len(); - compress_block(state, output); + let tables = compress_block(state, output); let payload_len = output.len() - payload_off; // Fall back to a raw block when the compressed payload is not // smaller than the source (`payload >= block_size`) or exceeds the @@ -266,34 +249,11 @@ pub(crate) fn compress_block_encoded( // already look incompressible, so small inputs that slip past it but // fail to shrink still need this post-hoc store-raw. if payload_len >= MAX_BLOCK_SIZE as usize || payload_len >= block_size as usize { - // Roll back the payload + reserved header and the entropy state. + // Drop the payload + reserved header; the block's tables are not + // committed, so the next block reads the ones this one started from. output.truncate(hdr_off); state.offset_hist = saved_offset_hist; - // Swap (not move) so the slot keeps owning a reusable table - // allocation for the next block's snapshot. With no prior table - // there is nothing to restore, and the table this block built goes - // into the slot rather than being dropped. - if had_prior_huff_table { - core::mem::swap(&mut state.last_huff_table, &mut saved_huff_table); - } else if let Some(built) = state.last_huff_table.take() { - saved_huff_table = Some(built); - } - // A block that built a table and then chose raw or RLE literals - // parked it on its way there, so the swap above had nothing to give - // back and the slot would go into the next block empty — which is - // the per-block allocation of both code buffers this slot exists to - // avoid, on exactly the run of blocks that keeps failing the size - // test. Take the parked table instead: the slot is asked for one - // per block, the spare's other reader once per frame. - if saved_huff_table.is_none() { - saved_huff_table = state.huff_table_spare.take(); - } - state.fse_tables.roll_back_confirmation([ - saved_ll_previous, - saved_ml_previous, - saved_of_previous, - ]); - state.block_scratch.huff_rollback = saved_huff_table; + discard_block_tables(state, tables); let header = BlockHeader { last_block, block_type: BlockType::Raw, @@ -304,9 +264,7 @@ pub(crate) fn compress_block_encoded( output.extend_from_slice(state.matcher.get_last_space()); BlockType::Raw } else { - // Return the snapshot to its slot so the next block's - // `clone_from` reuses the allocation. - state.block_scratch.huff_rollback = saved_huff_table; + commit_block_tables(state, tables); let header = BlockHeader { last_block, block_type: BlockType::Compressed, @@ -462,24 +420,9 @@ pub(crate) fn compress_block_encoded_borrowed( // below, so hashing `block` is both correct and order-safe. sink.push(crate::encoding::frame_compressor::xxh64_block_low32(block)); } + // As in the owned path: only the offset history is restored by hand; + // the block's tables are committed once it is kept. let saved_offset_hist = state.offset_hist; - // Persistent rollback slot — same allocation-reuse rationale as the - // owned `compress_block_encoded` snapshot above. - let mut saved_huff_table = core::mem::take(&mut state.block_scratch.huff_rollback); - // Only when there IS a table to snapshot. `Option::clone_from` from a - // `None` source drops what the slot holds, which is every buffer this - // persistent slot exists to keep: a run of blocks that build a table and - // then lose the size test would free and rebuild it every time. - let had_prior_huff_table = state.last_huff_table.is_some(); - if let Some(prior) = state.last_huff_table.as_ref() { - match saved_huff_table.as_mut() { - Some(slot) => slot.clone_from(prior), - None => saved_huff_table = Some(prior.clone()), - } - } - let saved_ll_previous = state.fse_tables.ll_previous.clone(); - let saved_ml_previous = state.fse_tables.ml_previous.clone(); - let saved_of_previous = state.fse_tables.of_previous.clone(); // Compress directly into `output`: reserve the fixed 3-byte block // header, append the payload after it, then backfill the header in // place once its length is known. Avoids the per-block temp `Vec` @@ -488,42 +431,18 @@ pub(crate) fn compress_block_encoded_borrowed( let hdr_off = output.len(); output.extend_from_slice(&[0u8; 3]); let payload_off = output.len(); - compress_block(state, output); + let tables = compress_block(state, output); let payload_len = output.len() - payload_off; if payload_len >= MAX_BLOCK_SIZE as usize || payload_len >= block_size as usize { // Incompressible (compressed payload not smaller than the source, - // or over the max block size): roll back the payload + reserved - // header and the entropy state, then emit a stored Raw block. A + // or over the max block size): drop the payload + reserved header + // and the block's tables, then emit a stored Raw block. A // non-shrinking compressed block wastes bytes and can reference // past a single-segment frame's window (== content size); storing // raw matches the upstream post-hoc fallback. output.truncate(hdr_off); state.offset_hist = saved_offset_hist; - // Swap (not move) so the slot keeps owning a reusable table - // allocation for the next block's snapshot. With no prior table - // there is nothing to restore, and the table this block built goes - // into the slot rather than being dropped. - if had_prior_huff_table { - core::mem::swap(&mut state.last_huff_table, &mut saved_huff_table); - } else if let Some(built) = state.last_huff_table.take() { - saved_huff_table = Some(built); - } - // A block that built a table and then chose raw or RLE literals - // parked it on its way there, so the swap above had nothing to give - // back and the slot would go into the next block empty — which is - // the per-block allocation of both code buffers this slot exists to - // avoid, on exactly the run of blocks that keeps failing the size - // test. Take the parked table instead: the slot is asked for one - // per block, the spare's other reader once per frame. - if saved_huff_table.is_none() { - saved_huff_table = state.huff_table_spare.take(); - } - state.fse_tables.roll_back_confirmation([ - saved_ll_previous, - saved_ml_previous, - saved_of_previous, - ]); - state.block_scratch.huff_rollback = saved_huff_table; + discard_block_tables(state, tables); let header = BlockHeader { last_block, block_type: BlockType::Raw, @@ -533,9 +452,7 @@ pub(crate) fn compress_block_encoded_borrowed( output.extend_from_slice(block); BlockType::Raw } else { - // Return the snapshot to its slot so the next block's - // `clone_from` reuses the allocation. - state.block_scratch.huff_rollback = saved_huff_table; + commit_block_tables(state, tables); let header = BlockHeader { last_block, block_type: BlockType::Compressed, diff --git a/zstd/src/encoding/levels/fastest/tests.rs b/zstd/src/encoding/levels/fastest/tests.rs index 510bdb35..3a7f5b7f 100644 --- a/zstd/src/encoding/levels/fastest/tests.rs +++ b/zstd/src/encoding/levels/fastest/tests.rs @@ -72,7 +72,6 @@ fn rle_branch_passes_compressible_hint_to_skip_matching() { copy_kernel: crate::encoding::fastpath::select_kernel(), last_huff_table: None, huff_table_spare: None, - huff_rollback: None, huff_weights: Default::default(), seen_content: Default::default(), fse_tables: FseTables::new(), diff --git a/zstd/src/encoding/match_generator/dict_prime.rs b/zstd/src/encoding/match_generator/dict_prime.rs index e09e5d9a..b8de6a38 100644 --- a/zstd/src/encoding/match_generator/dict_prime.rs +++ b/zstd/src/encoding/match_generator/dict_prime.rs @@ -184,11 +184,25 @@ impl MatchGeneratorDriver { | super::super::strategy::BackendTag::Row | super::super::strategy::BackendTag::HashChain => 4, }; - while start < dict_content.len() { - let end = (start + self.slice_size).min(dict_content.len()); - if end - start < min_primed_tail { - break; + // A dictionary too short to hash at all is left out whole: no match can + // reach it, so offsets stay right. A longer one goes in whole, since the + // decoder places every byte of it before the frame and a byte missing + // here shifts every offset into it. A last piece too short to hash is + // avoided by ending the piece before it that much earlier, or, when + // that would leave this piece too short instead, by running it to the + // end. + let hashable = dict_content.len() >= min_primed_tail; + while hashable && start < dict_content.len() { + let mut end = (start + self.slice_size).min(dict_content.len()); + if end != dict_content.len() && dict_content.len() - end < min_primed_tail { + let shortened = dict_content.len() - min_primed_tail; + end = if shortened - start >= min_primed_tail { + shortened + } else { + dict_content.len() + }; } + debug_assert!(end - start >= min_primed_tail); // The chunk goes straight into the history, as a block of input // does. let chunk = &dict_content[start..end]; diff --git a/zstd/src/encoding/match_generator/tests.rs b/zstd/src/encoding/match_generator/tests.rs index 60e3778d..9878c7b6 100644 --- a/zstd/src/encoding/match_generator/tests.rs +++ b/zstd/src/encoding/match_generator/tests.rs @@ -3396,19 +3396,21 @@ fn primed_snapshot_fast_attach_does_not_over_key_non_simple_backends() { ); } +/// A dictionary one byte past a whole slice is committed whole: the decoder +/// places every byte of it before the frame, so a byte left out would shift +/// every offset into it. The retention budget grows by all of it. #[test] -fn prime_with_dictionary_counts_only_committed_tail_budget() { +fn prime_with_dictionary_commits_a_short_last_piece() { let mut driver = MatchGeneratorDriver::new(8, 1); driver.reset(CompressionLevel::Fastest); let before = driver.simple_mut().max_window_size; - // One full slice plus a 1-byte tail that cannot be committed. driver.prime_with_dictionary(b"abcdefghi", [1, 4, 8]); assert_eq!( driver.simple_mut().max_window_size, - before + 8, - "retention budget must account only for dictionary bytes actually committed to history" + before + 9, + "every dictionary byte is in the history and in the budget" ); } @@ -3465,20 +3467,19 @@ fn row_prime_with_dictionary_preserves_history_for_first_full_block() { ); } +/// The Row backend takes a dictionary one byte past a whole slice whole too. #[test] -fn row_prime_with_dictionary_subtracts_uncommitted_tail_budget() { +fn row_prime_with_dictionary_commits_a_short_last_piece() { let mut driver = MatchGeneratorDriver::new(8, 1); driver.reset(CompressionLevel::Level(5)); let base_window = driver.row_matcher().max_window_size; - // Slice size is 8. The trailing byte cannot be committed (<4 tail), - // so it must be subtracted from retained budget. driver.prime_with_dictionary(b"abcdefghi", [1, 4, 8]); assert_eq!( driver.row_matcher().max_window_size, - base_window + 8, - "row retained window must exclude uncommitted 1-byte tail" + base_window + 9, + "every dictionary byte is in the history and in the budget" ); } diff --git a/zstd/src/encoding/row/mod.rs b/zstd/src/encoding/row/mod.rs index 4a34a716..12665f9b 100644 --- a/zstd/src/encoding/row/mod.rs +++ b/zstd/src/encoding/row/mod.rs @@ -131,8 +131,12 @@ enum LazyFinder { /// start, hoisted ONCE per block so the hot loops read no `Option` branch /// and no `Vec` header per access. Valid for the whole parse: the history /// buffer is never reallocated while a block is being scanned. +/// +/// `MLS` is the row key width (upstream's `mls` template argument), bound by +/// the per-tier parse monoliths so the hot loops hash with a constant width; +/// `0` reads it from `mls` instead, for the cold callers. #[derive(Clone, Copy)] -struct RowScan { +struct RowScan { base: *const u8, len: usize, hist_start: usize, @@ -150,9 +154,17 @@ struct RowScan { prefix_start: usize, dict_frame: bool, attached: bool, + /// The row key width, for an unbound (`MLS == 0`) view. + mls: u32, } -impl RowScan { +impl RowScan { + /// The row key width: the bound constant, or the stored one when unbound. + #[inline(always)] + fn key_mls(&self) -> u32 { + if MLS == 0 { self.mls } else { MLS } + } + /// Upstream `ZSTD_getLowestMatchIndex`: the whole history down to /// `lowLimit` while a dictionary is valid, else at most `maxDist` back. #[inline(always)] @@ -490,7 +502,9 @@ macro_rules! row_find_best_match { } else { (1u64 << row_entries) - 1 }; - let mut buf = [0u32; 64]; + // Filled up to `n` before it is read, so it is left uninitialised, as + // upstream's `matchBuffer` is: zeroing it cost 256 bytes per search. + let mut buf = [core::mem::MaybeUninit::::uninit(); 64]; let mut n = 0usize; { #[allow(unused_mut)] @@ -561,7 +575,7 @@ macro_rules! row_find_best_match { } } // SAFETY: `n < attempts_at_entry <= row_entries <= 64`. - unsafe { *buf.get_unchecked_mut(n) = raw }; + unsafe { buf.get_unchecked_mut(n).write(raw) }; n += 1; attempts -= 1; } @@ -570,7 +584,9 @@ macro_rules! row_find_best_match { // are gated at the match head (`MEM_read32(match) == MEM_read32(ip)`) // instead of at the current best length. let prefix_start = $ctx.prefix_start; - for &raw in &buf[..n] { + for slot in &buf[..n] { + // SAFETY: the first `n` entries were written above. + let raw = unsafe { slot.assume_init() }; let cand_idx = raw as usize - hist_start; // SAFETY: `cand_idx < cur_idx`; the gate reads 4 bytes at // `+ ml - 3` where `ml + 1 <= limit` (the count breaks out of the @@ -617,7 +633,7 @@ macro_rules! row_find_best_match { $ctx.base, $ctx.len, cur_idx, - $m.row_hash_mls, + $ctx.key_mls(), dict_row_hash_log + ROW_TAG_BITS, 0, ) @@ -839,7 +855,7 @@ macro_rules! hc_find_best_match { $ctx.base, $ctx.len, cur_idx, - $m.row_hash_mls, + $ctx.key_mls(), dict.hash_log, 0, ) @@ -1175,7 +1191,7 @@ macro_rules! dubt_find_best_match { $ctx.base, $ctx.len, cur_idx, - $m.row_hash_mls, + $ctx.key_mls(), dict.hash_log, 0, ) @@ -1297,7 +1313,7 @@ macro_rules! row_cache_insert_range { /// chosen match only, immediate-repcode loop after every stored sequence, /// `>> 8` miss acceleration with lazy skipping past step 8. macro_rules! lazy_parse_body { - ($m:expr, $handle:expr, $rl:expr, $finder:expr, $use_mask:literal, $maskmac:ident, $cpl:path) => {{ + ($m:expr, $handle:expr, $rl:expr, $finder:expr, $mls:ident, $use_mask:literal, $maskmac:ident, $cpl:path) => {{ #[allow(unused_labels)] 'parse: { debug_assert!($finder != FINDER_ROWS || $rl == $m.row_log); @@ -1317,7 +1333,7 @@ macro_rules! lazy_parse_body { } let block_end = current_abs_start + current_len; $m.enter_block(current_abs_start, block_end); - let scan = $m.scan_ctx(); + let scan = $m.scan_ctx_with::<$mls>(); let hist_start = scan.hist_start; let depth = $m.lazy_depth; // Upstream `ilimit`: `iend - 8 - ZSTD_ROW_HASH_CACHE_SIZE` for the @@ -1703,29 +1719,39 @@ macro_rules! gen_lazy_monolith { allow(dead_code) )] #[allow(unused_unsafe)] - unsafe fn $name( + unsafe fn $name( &mut self, mut handle_sequence: impl FnMut(Sequence), ) { - lazy_parse_body!(self, handle_sequence, ROW_LOG, FINDER, $use_mask, $maskmac, $cpl) + lazy_parse_body!(self, handle_sequence, ROW_LOG, FINDER, MLS, $use_mask, $maskmac, $cpl) } }; } -/// Bind the runtime finder and (for rows) `row_log` to the const parameters -/// of a lazy monolith: rows per `ROW_LOG` 4..=6, the chain and the tree -/// once each (they ignore `ROW_LOG`). Cold, once per block. +/// Bind the runtime finder, (for rows) `row_log` and the key width to the +/// const parameters of a lazy monolith: rows per `ROW_LOG` 4..=6, the chain +/// and the tree once each (they ignore `ROW_LOG`), each per `MLS` 4..=6, as +/// upstream instantiates its searches per `(rowLog, mls)`. Cold, once per +/// block. macro_rules! dispatch_lazy { ($self:ident . $m:ident :: <$k:ty> ( $($arg:expr),* )) => { + match $self.row_hash_mls { + 4 => dispatch_lazy!(@finder $self . $m :: <$k, 4> ( $($arg),* )), + 5 => dispatch_lazy!(@finder $self . $m :: <$k, 5> ( $($arg),* )), + 6 => dispatch_lazy!(@finder $self . $m :: <$k, 6> ( $($arg),* )), + _ => unreachable!("row_hash_mls is clamped to 4..=6 in configure()"), + } + }; + (@finder $self:ident . $m:ident :: <$k:ty, $mls:literal> ( $($arg:expr),* )) => { match $self.finder { LazyFinder::Rows => match $self.row_log { - 4 => $self.$m::<$k, 4, FINDER_ROWS>($($arg),*), - 5 => $self.$m::<$k, 5, FINDER_ROWS>($($arg),*), - 6 => $self.$m::<$k, 6, FINDER_ROWS>($($arg),*), + 4 => $self.$m::<$k, 4, FINDER_ROWS, $mls>($($arg),*), + 5 => $self.$m::<$k, 5, FINDER_ROWS, $mls>($($arg),*), + 6 => $self.$m::<$k, 6, FINDER_ROWS, $mls>($($arg),*), _ => unreachable!("row_log is clamped to 4..=6 in configure()"), }, - LazyFinder::Chain => $self.$m::<$k, 4, FINDER_CHAIN>($($arg),*), - LazyFinder::Tree => $self.$m::<$k, 4, FINDER_TREE>($($arg),*), + LazyFinder::Chain => $self.$m::<$k, 4, FINDER_CHAIN, $mls>($($arg),*), + LazyFinder::Tree => $self.$m::<$k, 4, FINDER_TREE, $mls>($($arg),*), } }; } @@ -3710,7 +3736,8 @@ impl RowMatchGenerator { /// Upstream zstd `ZSTD_hash4/5/6` of the `mls`-byte key at `base[idx..]` /// with `salt` XORed into the product, reduced to the top `bits`: /// `ZSTD_hashPtrSalted` for the row tables, `ZSTD_hashPtr` (salt 0) for - /// the hash chain. + /// the hash chain. The salt only spreads bucket indices between tables + /// and protects no secret, so a constant one is not a cryptographic value. /// /// # Safety /// `idx + ROW_HASH_KEY_LEN <= len` and `base` must point to `len` bytes. @@ -3745,17 +3772,26 @@ impl RowMatchGenerator { /// Hash-chain bucket of `abs_pos` (upstream `ZSTD_hashPtr(p, hashLog, mls)`). /// Callers pass positions at least `ROW_HASH_KEY_LEN` bytes before the end. #[inline(always)] - fn hc_hash_at(&self, ctx: RowScan, abs_pos: usize) -> usize { + fn hc_hash_at(&self, ctx: RowScan, abs_pos: usize) -> usize { let idx = abs_pos - ctx.hist_start; debug_assert!(idx + ROW_HASH_KEY_LEN <= ctx.len); let bits = self.row_hash_log + self.row_log; // SAFETY: caller contract, `ctx` views the live history. - unsafe { Self::key_hash_raw(ctx.base, ctx.len, idx, self.row_hash_mls, bits, 0) as usize } + unsafe { Self::key_hash_raw(ctx.base, ctx.len, idx, ctx.key_mls(), bits, 0) as usize } } - /// The live history as a [`RowScan`] for one parse. + /// The live history as a [`RowScan`] for one parse, its key width read at + /// run time. #[inline(always)] fn scan_ctx(&self) -> RowScan { + self.scan_ctx_with::<0>() + } + + /// The live history as a [`RowScan`] whose key width is the constant + /// `MLS` (`0`: read at run time), which must be this matcher's. + #[inline(always)] + fn scan_ctx_with(&self) -> RowScan { + debug_assert!(MLS == 0 || MLS == self.row_hash_mls); let history = self.live_history(); let dict_frame = self.loaded_dict_end != 0; RowScan { @@ -3768,13 +3804,18 @@ impl RowMatchGenerator { prefix_start: self.prefix_low, dict_frame, attached: dict_frame && self.dict_plan.is_some_and(|p| p.attach), + mls: self.row_hash_mls, } } /// `(row, tag)` of `abs_pos` through a hoisted [`RowScan`]; `None` when /// fewer than `ROW_HASH_KEY_LEN` bytes follow the position. #[inline(always)] - fn row_hash_at(&self, ctx: RowScan, abs_pos: usize) -> Option<(usize, u8)> { + fn row_hash_at( + &self, + ctx: RowScan, + abs_pos: usize, + ) -> Option<(usize, u8)> { let idx = abs_pos - ctx.hist_start; if idx + ROW_HASH_KEY_LEN > ctx.len { return None; @@ -3782,14 +3823,7 @@ impl RowMatchGenerator { let total_bits = self.row_hash_log + ROW_TAG_BITS; // SAFETY: `idx + ROW_HASH_KEY_LEN <= ctx.len`, `ctx` views the live history. let combined = unsafe { - Self::key_hash_raw( - ctx.base, - ctx.len, - idx, - self.row_hash_mls, - total_bits, - ctx.salt, - ) + Self::key_hash_raw(ctx.base, ctx.len, idx, ctx.key_mls(), total_bits, ctx.salt) }; let row_mask = (1usize << self.row_hash_log) - 1; let row = ((combined >> ROW_TAG_BITS) as usize) & row_mask; diff --git a/zstd/src/encoding/streaming_encoder.rs b/zstd/src/encoding/streaming_encoder.rs index 5958de96..cb4cbabe 100644 --- a/zstd/src/encoding/streaming_encoder.rs +++ b/zstd/src/encoding/streaming_encoder.rs @@ -396,7 +396,6 @@ impl CompressionContext { copy_kernel: crate::encoding::fastpath::select_kernel(), last_huff_table: None, huff_table_spare: None, - huff_rollback: None, huff_weights: Default::default(), seen_content: Default::default(), fse_tables: FseTables::new(), diff --git a/zstd/src/encoding/streaming_encoder/tests.rs b/zstd/src/encoding/streaming_encoder/tests.rs index a9e9b844..43d633be 100644 --- a/zstd/src/encoding/streaming_encoder/tests.rs +++ b/zstd/src/encoding/streaming_encoder/tests.rs @@ -410,26 +410,12 @@ fn the_reported_footprint_covers_what_compressing_retained() { assert_eq!(enc.heap_size(), after, "restoring must undo the removal"); // The entropy state kept between blocks is the other half of the figure: - // the FSE tables each axis holds in its two slots, and the Huffman tables - // the emit paths copy into before a block that may turn out raw. All of it - // survives the frame, so a caller sizing a context has to be told about it. - // Checked by removal rather than by comparison: the match-finder's share - // alone exceeds it, so any "total is at least the parts" assertion would - // pass with the term missing entirely. - let entropy = enc.context.state.fse_tables.heap_size() - + enc - .context - .state - .huff_rollback - .as_ref() - .map_or(0, |table| table.heap_size()) - + enc - .context - .state - .block_scratch - .huff_rollback - .as_ref() - .map_or(0, |table| table.heap_size()); + // the FSE tables each axis holds in its two slots. It survives the frame, + // so a caller sizing a context has to be told about it. Checked by removal + // rather than by comparison: the match-finder's share alone exceeds it, so + // any "total is at least the parts" assertion would pass with the term + // missing entirely. + let entropy = enc.context.state.fse_tables.heap_size(); assert!( entropy > 0, "compressing should have left entropy state retained", @@ -439,16 +425,12 @@ fn the_reported_footprint_covers_what_compressing_retained() { &mut enc.context.state.fse_tables, crate::encoding::frame_compressor::FseTables::new(), ); - let saved_rollback = enc.context.state.huff_rollback.take(); - let saved_scratch_rollback = enc.context.state.block_scratch.huff_rollback.take(); assert_eq!( with_entropy - enc.heap_size(), entropy, "the entropy state is retained but not counted in the reported footprint", ); enc.context.state.fse_tables = saved_tables; - enc.context.state.huff_rollback = saved_rollback; - enc.context.state.block_scratch.huff_rollback = saved_scratch_rollback; assert_eq!(enc.heap_size(), with_entropy, "restoring must undo removal"); } use crate::io::{Error, ErrorKind, Read, Write}; diff --git a/zstd/tests/cross_validation.rs b/zstd/tests/cross_validation.rs index 2dc0260f..9eeba550 100644 --- a/zstd/tests/cross_validation.rs +++ b/zstd/tests/cross_validation.rs @@ -653,6 +653,48 @@ fn cross_every_level_with_ldm_and_dictionary_roundtrips_both_ways() { ); } +/// A raw-content dictionary whose start the input repeats, so the input's +/// matches reach the dictionary's first byte, compressed with its size pledged +/// (as `zstd -D dict file` does). The pledge shrinks the window, so the +/// dictionary is taken in window-sized pieces, and a length one to three bytes +/// past a whole number of them leaves a last piece too short to hash: those +/// bytes still belong to the history, or every offset into the dictionary +/// lands that many bytes off. The frame has to decode through the C decoder at +/// every level and every tail length. +#[test] +fn cross_dictionary_with_a_short_last_piece_roundtrips_at_every_level() { + let mut failed = Vec::new(); + for tail in 0..=3usize { + let dict_raw = generate_huffman_friendly(0xD1C7, 65_536 + tail, 40); + let data = &dict_raw[..11_000]; + let dictionary = + structured_zstd::encoding::EncoderDictionary::from_serialized_or_raw_content(&dict_raw) + .expect("any bytes are a raw-content dictionary"); + for level in 1i32..=22 { + let mut context = structured_zstd::encoding::CompressionContext::new( + CompressionLevel::from_level(level), + ); + context.set_content_checksum(true).unwrap(); + context.set_encoder_dictionary(dictionary.clone()).unwrap(); + let mut encoder = + structured_zstd::encoding::StreamingEncoder::with_context(Vec::new(), &mut context); + encoder.set_pledged_content_size(data.len() as u64).unwrap(); + std::io::Write::write_all(&mut encoder, data).unwrap(); + let ours = encoder.finish().unwrap(); + + let mut c_dec = zstd::bulk::Decompressor::with_dictionary(&dict_raw).unwrap(); + match c_dec.decompress(ours.as_slice(), data.len()) { + Ok(decoded) if decoded == data => {} + other => failed.push((tail, level, other.err().map(|e| e.to_string()))), + } + } + } + assert!( + failed.is_empty(), + "(tail, level) whose frame does not decode: {failed:?}" + ); +} + /// A frame whose descriptor has the Reserved_bit (bit 3) set is refused by /// libzstd and by this decoder, and the same frame with only the Unused_bit /// (bit 4) set is accepted by both. libzstd's code is