Skip to main content

structured_zstd/encoding/
frame_compressor.rs

1//! Utilities and interfaces for encoding an entire frame. Allows reusing resources
2
3use alloc::vec::Vec;
4use core::convert::TryInto;
5#[cfg(feature = "hash")]
6use twox_hash::XxHash64;
7
8#[cfg(feature = "hash")]
9use core::hash::Hasher;
10
11use super::{
12    CompressionLevel, LiteralCompressionMode, Matcher, block_header::BlockHeader,
13    frame_header::FrameHeader, levels::*, match_generator::MatchGeneratorDriver,
14};
15use crate::common::MAX_BLOCK_SIZE;
16use crate::fse::fse_encoder::{FSETable, default_ll_table, default_ml_table, default_of_table};
17
18use crate::io::{Read, Write};
19
20/// A dictionary prepared for the ENCODER side, analogous to zstd's `CDict`
21/// (vs the decoder's [`Dictionary`](crate::decoding::Dictionary) / `DDict`).
22///
23/// It carries the entropy tables, content, and repeat-offset history the
24/// compressor needs, but is a distinct type with **no decode path**: there is
25/// no way to turn it into a [`DictionaryHandle`](crate::decoding::DictionaryHandle)
26/// or feed it to a [`FrameDecoder`](crate::decoding::FrameDecoder). That keeps
27/// the compress-only state (which may have been parsed without building the
28/// decode lookup tables, see
29/// [`set_dictionary_from_bytes`](FrameCompressor::set_dictionary_from_bytes))
30/// from ever reaching the decode side — the encoder/decoder dictionary split
31/// mirrors C zstd's `CDict` / `DDict`.
32/// Cloning one is a handle, not a copy: it is attached to a compressor by
33/// value, so a dictionary serving many frames would otherwise have its parsed
34/// tables and content duplicated for each of them — on exactly the path where
35/// one dictionary is prepared once precisely to be used again and again.
36///
37/// The encoder entropy tables a dictionary seeds into each frame's first block
38/// (upstream zstd `cdict->cBlockState`) are built once, here, and shared by
39/// every compressor the dictionary is attached to.
40#[derive(Clone)]
41pub struct EncoderDictionary {
42    pub(crate) inner: SharedEncoderDictionary,
43    /// Size of the serialized dictionary this was built from (header, entropy
44    /// tables, repeat offsets and content); the CDict cParams tier key
45    /// (upstream `ZSTD_createCDict(dictBuffer, dictSize, level)`). Falls back
46    /// to the content length when the wrapped [`Dictionary`] was handed over
47    /// already parsed ([`Self::from_dictionary`]) — exact for raw-content
48    /// dictionaries, a close lower bound otherwise.
49    serialized_len: usize,
50}
51
52/// What an [`EncoderDictionary`] shares between its clones: the parsed
53/// dictionary and the encoder entropy tables derived from it, in one
54/// allocation so a clone is one reference count, not two. Reads through to the
55/// [`Dictionary`](crate::decoding::Dictionary) for its content, offsets and id.
56pub(crate) struct EncoderDictionaryParts {
57    dictionary: crate::decoding::Dictionary,
58    pub(crate) entropy: CachedDictionaryEntropy,
59}
60
61impl core::ops::Deref for EncoderDictionaryParts {
62    type Target = crate::decoding::Dictionary;
63
64    fn deref(&self) -> &Self::Target {
65        &self.dictionary
66    }
67}
68
69/// Shared owner of [`EncoderDictionaryParts`]: `Arc` on atomic-pointer
70/// targets, `Rc` otherwise, as [`SharedFseTable`] is.
71#[cfg(target_has_atomic = "ptr")]
72pub(crate) type SharedEncoderDictionary = alloc::sync::Arc<EncoderDictionaryParts>;
73#[cfg(not(target_has_atomic = "ptr"))]
74pub(crate) type SharedEncoderDictionary = alloc::rc::Rc<EncoderDictionaryParts>;
75
76impl EncoderDictionary {
77    /// Share `dictionary` with the entropy tables it seeds, built once here.
78    fn prepared(dictionary: crate::decoding::Dictionary, serialized_len: usize) -> Self {
79        let entropy = CachedDictionaryEntropy::from_dictionary(&dictionary);
80        Self {
81            inner: SharedEncoderDictionary::new(EncoderDictionaryParts {
82                dictionary,
83                entropy,
84            }),
85            serialized_len,
86        }
87    }
88
89    /// Wrap an already-parsed [`Dictionary`](crate::decoding::Dictionary) for
90    /// encoder use. A fully-decoded dictionary is valid here; only the encoder
91    /// entropy tables, content, and offset history are read. The CDict cParams
92    /// tier is keyed by the content length here; prefer [`Self::from_bytes`]
93    /// when the serialized blob is at hand — it keys the tier by the exact
94    /// serialized size as upstream `ZSTD_createCDict` does.
95    pub fn from_dictionary(dictionary: crate::decoding::Dictionary) -> Self {
96        let serialized_len = dictionary.dict_content.len();
97        Self::prepared(dictionary, serialized_len)
98    }
99
100    /// Parse a serialized dictionary blob for encoder use, skipping the decode
101    /// lookup-table build the encoder never reads (see
102    /// `Dictionary::decode_dict_for_encoding`). The encoder entropy tables — and
103    /// thus the emitted frame — are identical to a full parse.
104    pub fn from_bytes(
105        raw_dictionary: &[u8],
106    ) -> Result<Self, crate::decoding::errors::DictionaryDecodeError> {
107        Ok(Self::prepared(
108            crate::decoding::Dictionary::decode_dict_for_encoding(raw_dictionary)?,
109            raw_dictionary.len(),
110        ))
111    }
112
113    /// Load whichever kind of dictionary `raw_dictionary` holds, the way
114    /// `zstd -D` does: a serialized blob is parsed, anything else is taken as
115    /// raw content (see
116    /// [`Dictionary::from_serialized_or_raw_content`](crate::decoding::Dictionary::from_serialized_or_raw_content)).
117    ///
118    /// Either way the blob's own length is what the compression-parameter tier
119    /// is chosen by, which is why this exists rather than parsing and calling
120    /// [`Self::from_dictionary`]: that keys the tier on the content length, and
121    /// for a serialized dictionary the entropy tables in between can put the
122    /// two on opposite sides of a boundary.
123    pub fn from_serialized_or_raw_content(
124        raw_dictionary: &[u8],
125    ) -> Result<Self, crate::decoding::errors::DictionaryDecodeError> {
126        // Parsed for the encoder, which reads the entropy probabilities, the
127        // content and the offsets and never the decode lookup tables: routing a
128        // serialized blob through the full parser builds those tables for
129        // nothing. The emitted frame is identical either way — only the wasted
130        // build is dropped (see `Dictionary::decode_dict_for_encoding`).
131        if raw_dictionary.starts_with(&crate::decoding::DICTIONARY_MAGIC) {
132            return Self::from_bytes(raw_dictionary);
133        }
134        Ok(Self::prepared(
135            crate::decoding::Dictionary::from_raw_content(0, raw_dictionary.to_vec())?,
136            raw_dictionary.len(),
137        ))
138    }
139
140    /// Heap bytes the prepared dictionary holds: the shared allocation itself,
141    /// the parsed dictionary's content and tables, and the entropy tables it
142    /// seeds. Clones share all of it, so this counts once however many
143    /// compressors the dictionary is attached to.
144    ///
145    /// # Examples
146    /// ```
147    /// use structured_zstd::encoding::EncoderDictionary;
148    ///
149    /// let dictionary =
150    ///     EncoderDictionary::from_serialized_or_raw_content(b"raw content, no header").unwrap();
151    /// assert!(dictionary.heap_size() >= b"raw content, no header".len());
152    /// ```
153    pub fn heap_size(&self) -> usize {
154        // `Arc` and `Rc` both place two word-sized counts ahead of the value,
155        // padded to its alignment.
156        let shared = core::alloc::Layout::new::<[usize; 2]>()
157            .extend(core::alloc::Layout::new::<EncoderDictionaryParts>())
158            .expect("a fixed-size layout extends")
159            .0
160            .pad_to_align()
161            .size();
162        shared + self.inner.dictionary.heap_bytes() + self.inner.entropy.heap_size()
163    }
164
165    /// Heap bytes that `handles` alone keep alive: each distinct dictionary
166    /// among them once, and only when no handle outside them shares it.
167    ///
168    /// For an owner holding several handles, such as a context whose
169    /// compressors each keep the dictionary they last ran with: what it reports
170    /// should count one dictionary once, and none that someone else also holds.
171    ///
172    /// # Examples
173    /// ```
174    /// use structured_zstd::encoding::EncoderDictionary;
175    ///
176    /// let dictionary = EncoderDictionary::from_serialized_or_raw_content(b"some shared history").unwrap();
177    /// let copy = dictionary.clone();
178    /// // Both handles are in the set: counted once.
179    /// assert_eq!(
180    ///     EncoderDictionary::exclusive_heap_size([&dictionary, &copy]),
181    ///     dictionary.heap_size()
182    /// );
183    /// // `copy` is held elsewhere: nothing is this set's alone.
184    /// assert_eq!(EncoderDictionary::exclusive_heap_size([&dictionary]), 0);
185    /// ```
186    pub fn exclusive_heap_size<'a, I>(handles: I) -> usize
187    where
188        I: IntoIterator<Item = &'a EncoderDictionary>,
189        I::IntoIter: Clone,
190    {
191        let handles = handles.into_iter();
192        let mut total = 0;
193        for (position, handle) in handles.clone().enumerate() {
194            let same = |other: &&EncoderDictionary| {
195                SharedEncoderDictionary::ptr_eq(&handle.inner, &other.inner)
196            };
197            // The first handle to an allocation speaks for all of them.
198            if handles.clone().take(position).any(|other| same(&other)) {
199                continue;
200            }
201            let held_here = handles.clone().filter(same).count();
202            if SharedEncoderDictionary::strong_count(&handle.inner) == held_here {
203                total += handle.heap_size();
204            }
205        }
206        total
207    }
208
209    /// The content and serialized sizes the encoder's matcher is hinted with.
210    pub(crate) fn sizes(&self) -> crate::encoding::DictionarySizes {
211        crate::encoding::DictionarySizes {
212            content: self.inner.dict_content.len(),
213            serialized: self.serialized_len,
214        }
215    }
216
217    /// Bytes of dictionary content, the part a frame can reference.
218    ///
219    /// This is the extent a caller weighs a source against when deciding
220    /// whether a frame is still about the dictionary or about its own content.
221    ///
222    /// # Examples
223    ///
224    /// ```
225    /// use structured_zstd::encoding::EncoderDictionary;
226    ///
227    /// let dictionary = EncoderDictionary::from_serialized_or_raw_content(&[7u8; 4096])
228    ///     .expect("raw content is always usable");
229    /// assert_eq!(dictionary.content_size(), 4096);
230    /// ```
231    pub fn content_size(&self) -> usize {
232        self.inner.dict_content.len()
233    }
234
235    /// The dictionary id.
236    ///
237    /// Zero is a raw-content dictionary, which has no header to carry an id.
238    /// Such a dictionary attaches like any other; what changes is the frame,
239    /// which omits the `Dictionary_ID` field rather than storing a zero, so a
240    /// decoder has to be handed the same bytes explicitly.
241    pub fn id(&self) -> u32 {
242        self.inner.id
243    }
244}
245
246/// An interface for compressing arbitrary data with the ZStandard compression algorithm.
247///
248/// `FrameCompressor` will generally be used by:
249/// 1. Initializing a compressor by providing a buffer of data using `FrameCompressor::new()`
250/// 2. Starting compression and writing that compression into a vec using `FrameCompressor::begin`
251///
252/// # Examples
253/// ```
254/// use structured_zstd::encoding::{FrameCompressor, CompressionLevel};
255/// let mock_data: &[_] = &[0x1, 0x2, 0x3, 0x4];
256/// let mut output = std::vec::Vec::new();
257/// // Initialize a compressor.
258/// let mut compressor = FrameCompressor::new(CompressionLevel::Uncompressed);
259/// compressor.set_source(mock_data);
260/// compressor.set_drain(&mut output);
261///
262/// // `compress` writes the compressed output into the provided buffer.
263/// compressor.compress();
264/// ```
265pub struct FrameCompressor<
266    R: Read = &'static [u8],
267    W: Write = Vec<u8>,
268    M: Matcher = MatchGeneratorDriver,
269> {
270    uncompressed_data: Option<R>,
271    compressed_data: Option<W>,
272    compression_level: CompressionLevel,
273    dictionary: Option<EncoderDictionary>,
274    source_size_hint: Option<u64>,
275    state: CompressState<M>,
276    /// When true, emitted frames omit the 4-byte magic number prefix
277    /// (`ZSTD_f_zstd1_magicless`). Default false. The caller is
278    /// responsible for ensuring the decoder is configured for the
279    /// matching format — wire-format only round-trips with a
280    /// magicless-aware decoder.
281    magicless: bool,
282    /// Whether to emit a trailing XXH64 content checksum and set the frame
283    /// header's `Content_Checksum_flag` (semantics of upstream
284    /// `ZSTD_c_checksumFlag`). Default `false`, matching the upstream
285    /// library default; combined with the `hash` feature at frame-build
286    /// time, so without `hash` no checksum is emitted regardless. Set via
287    /// [`Self::set_content_checksum`].
288    content_checksum: bool,
289    /// Diagnostic: skip the block pre-splitter and cut full blocks only
290    /// (upstream's block structure under `ZSTD_generateSequences`, whose
291    /// sequence-collecting mode never accrues the savings the splitter
292    /// requires). Set via [`Self::set_pre_split_disabled`]; default `false`.
293    pre_split_disabled: bool,
294    /// Whether to record `Frame_Content_Size` in the frame header when the
295    /// total size is known (semantics of upstream `ZSTD_c_contentSizeFlag`).
296    /// Default `true`, matching upstream. With the flag off the header
297    /// carries a window descriptor instead (single-segment requires an FCS,
298    /// so it is disabled too). Set via [`Self::set_content_size_flag`].
299    content_size_flag: bool,
300    /// Whether to record the dictionary ID in the frame header when a
301    /// dictionary is attached (semantics of upstream `ZSTD_c_dictIDFlag`).
302    /// Default `true`, matching upstream. Decoders can still decode the
303    /// frame by being handed the right dictionary explicitly. Set via
304    /// [`Self::set_dictionary_id_flag`].
305    dict_id_flag: bool,
306    /// Upper bound on emitted block sizes (semantics of upstream
307    /// `ZSTD_c_targetCBlockSize`): capping the RAW block length at the
308    /// target bounds every physical block's compressed payload at the
309    /// target too (a compressed block never exceeds its raw input — the
310    /// raw-block fallback fires otherwise), so blocks land at or under
311    /// `target + 3` header bytes on the wire. `None` = no target (full
312    /// 128 KiB blocks). Set via [`Self::set_target_block_size`].
313    target_block_size: Option<u32>,
314    #[cfg(feature = "hash")]
315    hasher: XxHash64,
316    /// Block-layout introspection populated at the end of every
317    /// successful `compress()`. `None` until the first call.
318    /// Behind the `lsm` feature gate.
319    #[cfg(feature = "lsm")]
320    frame_emit_info: Option<crate::encoding::frame_emit_info::FrameEmitInfo>,
321    /// When `true`, `compress()` XXH64-hashes each block's
322    /// uncompressed bytes and appends the low-32-bit digest to
323    /// `block_checksums`. Default `false` (zero cost). Gated on
324    /// `all(lsm, hash)` because XXH64 lives behind the `hash`
325    /// feature; an `lsm`-only build has no way to compute digests.
326    #[cfg(all(feature = "lsm", feature = "hash"))]
327    per_block_checksums_enabled: bool,
328    /// Per-block XXH64 (low 32 bits) digests captured during
329    /// `compress()` when `per_block_checksums_enabled` is set. Ordered
330    /// by block-emit order. `None` until the first call after enabling.
331    /// Gated on `all(lsm, hash)` (see `per_block_checksums_enabled`).
332    #[cfg(all(feature = "lsm", feature = "hash"))]
333    block_checksums: Option<alloc::vec::Vec<u32>>,
334    /// Per-physical-block decompressed (regenerated) sizes captured
335    /// during `compress()`, in block-emit order (1:1 with
336    /// `frame_emit_info.blocks`). Always captured under `lsm` (no
337    /// opt-in, unlike `block_checksums`) because `FrameEmitInfo` is
338    /// always built under `lsm` and `decompressed_byte_range` needs
339    /// the per-block sizes. Cleared and refilled per frame.
340    #[cfg(feature = "lsm")]
341    block_decompressed_sizes: alloc::vec::Vec<u32>,
342    tuning: FrameTuning,
343}
344
345/// What a compressor keeps of the fine-grained parameters (#27) beside the
346/// level: the parts the literal-compression gates and the block splitter read
347/// at frame start, where the matcher holds the rest. Shared by
348/// [`FrameCompressor`] and [`CompressionContext`](crate::encoding::CompressionContext).
349#[derive(Clone, Copy, Default)]
350pub(crate) struct FrameTuning {
351    /// A public-parameter [`Strategy`](crate::encoding::Strategy) override:
352    /// its tag and lazy depth (the collapsed `Lazy` tag needs the depth for
353    /// the pre-split tier). `Some` wins over the level-derived
354    /// `state.strategy_tag` so the gates and the dict-attach cutoff see the
355    /// strategy the matcher runs; `None` keeps the level's.
356    pub(crate) strategy: Option<(crate::encoding::strategy::StrategyTag, u8)>,
357    /// Public `target_length` override, kept so the raw-literals gate can be
358    /// recomputed per frame: a dictionary attached or cleared after
359    /// `set_parameters` can move the strategy the gate reads.
360    pub(crate) target_length: Option<u32>,
361    /// Public literal-compression mode (upstream
362    /// `ZSTD_c_literalCompressionMode`), kept for the same per-frame
363    /// recomputation of the raw-literals gate.
364    pub(crate) literal_compression: LiteralCompressionMode,
365}
366
367impl FrameTuning {
368    /// The tuning `overrides` asks for.
369    pub(crate) fn from_overrides(overrides: &crate::encoding::parameters::ParamOverrides) -> Self {
370        Self {
371            strategy: overrides.strategy.map(|s| (s.tag(), s.lazy_depth())),
372            target_length: overrides.target_length,
373            literal_compression: overrides.literal_compression,
374        }
375    }
376}
377
378#[derive(Clone, Default)]
379pub(crate) struct CachedDictionaryEntropy {
380    pub(crate) huff: Option<crate::huff0::huff0_encoder::HuffmanTable>,
381    pub(crate) ll_previous: Option<PreviousFseTable>,
382    pub(crate) ml_previous: Option<PreviousFseTable>,
383    pub(crate) of_previous: Option<PreviousFseTable>,
384}
385
386impl CachedDictionaryEntropy {
387    /// Heap bytes the cached dictionary entropy holds: the literals Huffman
388    /// table plus any `Custom` LL/ML/OF FSE tables (the `Arc`-boxed `FSETable`
389    /// payload and its flat state array). `Default` / `Rle` variants own no heap.
390    pub(crate) fn heap_size(&self) -> usize {
391        let mut total = self.huff.as_ref().map_or(0, |h| h.heap_size());
392        for prev in [&self.ll_previous, &self.ml_previous, &self.of_previous] {
393            if let Some(PreviousFseTable::Custom(table)) = prev {
394                total +=
395                    core::mem::size_of::<crate::fse::fse_encoder::FSETable>() + table.heap_size();
396            }
397        }
398        total
399    }
400
401    /// Derive the encoder-side entropy tables a dictionary seeds for the first
402    /// block of each frame (the upstream zstd `cdict->cBlockState`): the literals
403    /// Huffman table plus the literal-length / match-length / offset FSE
404    /// "previous" tables. Shared by [`FrameCompressor`] and
405    /// [`crate::encoding::StreamingEncoder`] so both seed identically.
406    pub(crate) fn from_dictionary(dictionary: &crate::decoding::Dictionary) -> Self {
407        Self {
408            huff: dictionary.huf.table.to_encoder_table(),
409            ll_previous: dictionary
410                .fse
411                .literal_lengths
412                .to_encoder_table()
413                .map(|table| PreviousFseTable::Custom(SharedFseTable::new(table))),
414            ml_previous: dictionary
415                .fse
416                .match_lengths
417                .to_encoder_table()
418                .map(|table| PreviousFseTable::Custom(SharedFseTable::new(table))),
419            of_previous: dictionary
420                .fse
421                .offsets
422                .to_encoder_table()
423                .map(|table| PreviousFseTable::Custom(SharedFseTable::new(table))),
424        }
425    }
426}
427
428/// Shared owner for a custom "previous" FSE encoder table. `Arc` on
429/// atomic-pointer targets, `Rc` otherwise (keeps `no_std` no-atomics
430/// builds compiling, single-thread there anyway), mirroring
431/// `decoding::dictionary::SharedDictionary`. Cloning the cached
432/// dictionary entropy into the per-frame state is then a refcount bump,
433/// not a full `FSETable` copy — the upstream zstd references `cdict->cBlockState`
434/// instead of rebuilding it per frame.
435#[cfg(target_has_atomic = "ptr")]
436pub(crate) type SharedFseTable = alloc::sync::Arc<FSETable>;
437#[cfg(not(target_has_atomic = "ptr"))]
438pub(crate) type SharedFseTable = alloc::rc::Rc<FSETable>;
439
440/// Bytes a [`SharedFseTable`] allocation carries in FRONT of the table: the two
441/// reference counts, then whatever padding the table's alignment adds. One
442/// allocation holds both, so a caller sizing a context is told about the whole
443/// of it rather than the payload alone.
444const fn shared_table_overhead() -> usize {
445    let counts = 2 * core::mem::size_of::<usize>();
446    let align = core::mem::align_of::<FSETable>();
447    counts.div_ceil(align) * align
448}
449
450#[derive(Clone)]
451pub(crate) enum PreviousFseTable {
452    // Default tables are immutable and already stored alongside the state, so
453    // repeating them only needs a lightweight marker instead of cloning FSETable.
454    Default,
455    // Shared handle: cloning (per-frame dictionary entropy seed) is a refcount
456    // bump. The table is only ever read or REPLACED wholesale (a block that
457    // builds a new table swaps in a fresh `SharedFseTable`), never mutated in
458    // place, so sharing is sound.
459    Custom(SharedFseTable),
460    Rle(u8),
461}
462
463impl PreviousFseTable {
464    pub(crate) fn as_table<'a>(&'a self, default: &'a FSETable) -> Option<&'a FSETable> {
465        match self {
466            Self::Default => Some(default),
467            Self::Custom(table) => Some(table),
468            Self::Rle(_) => None,
469        }
470    }
471}
472
473pub(crate) struct FseTables {
474    /// The three predefined LL/ML/OF tables are functions of
475    /// compile-time-constant distributions. The
476    /// [`fse_encoder::FseDefaultTable`] type alias resolves to
477    /// `&'static FSETable` when a process-wide cache is available
478    /// (atomic-pointer targets, or no-atomic targets with the
479    /// `critical-section` feature) and to `Box<FSETable>` on the
480    /// cache-less no-atomic path (one per-frame allocation, dropped
481    /// with the compressor — no `Box::leak`, no unbounded growth).
482    /// Both arms `Deref` to `FSETable`, so consumers in
483    /// `encoding/blocks/compressed.rs` borrow through `&` uniformly
484    /// without seeing the per-target divergence.
485    pub(crate) ll_default: crate::fse::fse_encoder::FseDefaultTable,
486    pub(crate) ll_previous: Option<PreviousFseTable>,
487    pub(crate) ml_default: crate::fse::fse_encoder::FseDefaultTable,
488    pub(crate) ml_previous: Option<PreviousFseTable>,
489    pub(crate) of_default: crate::fse::fse_encoder::FseDefaultTable,
490    pub(crate) of_previous: Option<PreviousFseTable>,
491    /// Where a block builds the table it is about to emit, before that table
492    /// becomes the axis's `*_previous`.
493    ///
494    /// Upstream's `ZSTD_blockState_t` keeps `prevCBlock` and `nextCBlock` for
495    /// exactly this: the entropy build reads the previous tables and writes the
496    /// next ones, so the two never alias and committing a block is a pointer
497    /// swap (`ZSTD_blockState_confirmRepcodesAndEntropyTables`). A block that
498    /// ends up raw simply does not swap. Holding the slot here is what lets a
499    /// table be built in place instead of on the stack.
500    /// `None` until an axis first builds a table. Lazy because a state that
501    /// never emits a custom table must not pay for a slot: the block splitter
502    /// makes one of these per probe, and eagerly giving each three tables
503    /// faulted in pages for buffers most probes never wrote to.
504    pub(crate) ll_next: Option<SharedFseTable>,
505    pub(crate) ml_next: Option<SharedFseTable>,
506    pub(crate) of_next: Option<SharedFseTable>,
507}
508
509impl FseTables {
510    /// Undo a block's table confirmation, keeping the tables it built as the
511    /// next block's buffers.
512    ///
513    /// A block that loses to a raw one has already had its tables confirmed
514    /// into the previous slots, and the caller holds a clone of what they
515    /// were. Restoring only that clone leaves the axis holding the SAME table
516    /// in both slots — the restored one in `previous` and the one confirmation
517    /// displaced into `next` — so the next block finds its buffer shared and
518    /// allocates eleven kilobytes instead of writing into it, on every block
519    /// of a run that keeps falling back.
520    ///
521    /// Handing the discarded table to `next` fixes both halves at once: the
522    /// restored table is unique again, and the buffer the block just filled is
523    /// exactly what the next one wants to build into.
524    pub(crate) fn roll_back_confirmation(&mut self, saved: [Option<PreviousFseTable>; 3]) {
525        let [saved_ll, saved_ml, saved_of] = saved;
526        for (previous, next, restored) in [
527            (&mut self.ll_previous, &mut self.ll_next, saved_ll),
528            (&mut self.ml_previous, &mut self.ml_next, saved_ml),
529            (&mut self.of_previous, &mut self.of_next, saved_of),
530        ] {
531            let discarded = core::mem::replace(previous, restored);
532            // The table the block built, and it is not the one just restored:
533            // whatever `next` held was the outgoing previous, which is what was
534            // restored, so it is a duplicate and this one is free and unique.
535            //
536            // The exception is `RepeatLast`, where confirmation never replaced
537            // `previous` at all — the discarded handle IS the restored one.
538            // Parking it in `next` then leaves the two sharing a handle, and the
539            // next custom build cannot write into a table `previous` still
540            // holds, so it allocates another one.
541            let built = match discarded {
542                Some(PreviousFseTable::Custom(built))
543                    if !matches!(previous.as_ref(), Some(PreviousFseTable::Custom(back))
544                        if SharedFseTable::ptr_eq(&built, back)) =>
545                {
546                    Some(built)
547                }
548                _ => None,
549            };
550            match built {
551                Some(built) => *next = Some(built),
552                // Nothing new to park. The block settled on a predefined or RLE
553                // table, or repeated the last one, so confirmation may have
554                // parked the outgoing custom table here — and if that is the
555                // table now back in `previous`, holding it here too is what
556                // keeps the next build from writing into it.
557                None => {
558                    if let (Some(spare), Some(PreviousFseTable::Custom(back))) =
559                        (next.as_ref(), previous.as_ref())
560                        && SharedFseTable::ptr_eq(spare, back)
561                    {
562                        *next = None;
563                    }
564                }
565            }
566        }
567    }
568
569    /// Keep the frame's tables as buffers before the next frame overwrites
570    /// them.
571    ///
572    /// A frame start replaces every previous slot — with the dictionary's
573    /// seed, or with nothing — and dropping the table that was there costs a
574    /// reused compressor an eleven-kilobyte allocation per axis on the next
575    /// frame that builds one. The slot the axis builds into is exactly where
576    /// it belongs. Only an unshared table is worth keeping: one still held by
577    /// the dictionary entropy cache cannot be built into anyway.
578    pub(crate) fn park_previous_before_frame(&mut self) {
579        for (previous, next) in [
580            (&mut self.ll_previous, &mut self.ll_next),
581            (&mut self.ml_previous, &mut self.ml_next),
582            (&mut self.of_previous, &mut self.of_next),
583        ] {
584            // Decided before taking: a handle the cache still holds stays
585            // where it is rather than being pulled out and dropped.
586            let worth_keeping = matches!(
587                previous.as_ref(),
588                Some(PreviousFseTable::Custom(handle))
589                    if SharedFseTable::strong_count(handle) == 1
590            );
591            // An occupied slot is not the same as a usable one. A frame that
592            // settled on a predefined table parks the dictionary cache's handle
593            // here, and a shared handle cannot be built into — so leaving it
594            // there while the uniquely owned table is dropped costs the
595            // allocation this exists to avoid, on the first custom build of
596            // every frame.
597            let next_is_a_buffer = next
598                .as_ref()
599                .is_some_and(|handle| SharedFseTable::strong_count(handle) == 1);
600            if !next_is_a_buffer
601                && worth_keeping
602                && let Some(PreviousFseTable::Custom(handle)) = previous.take()
603            {
604                *next = Some(handle);
605            }
606        }
607    }
608
609    /// Heap bytes the retained encoder tables hold.
610    ///
611    /// Both slots of an axis count: the previous table is what the next block
612    /// reads, the next slot is the buffer it builds into, and each lives as
613    /// long as the compressor. A handle shared with the dictionary entropy
614    /// cache is left out, because that cache reports its own tables and the
615    /// caller would otherwise be told about the same allocation twice.
616    pub(crate) fn heap_size(&self) -> usize {
617        // The table AND the reference counts in front of it: a shared handle is
618        // one allocation holding both, and reporting only the payload
619        // understates every retained table by the control block and whatever
620        // padding its alignment adds.
621        let per_table = core::mem::size_of::<FSETable>() + shared_table_overhead();
622        let mut total = 0;
623        for previous in [&self.ll_previous, &self.ml_previous, &self.of_previous] {
624            if let Some(PreviousFseTable::Custom(handle)) = previous
625                && SharedFseTable::strong_count(handle) == 1
626            {
627                total += per_table;
628            }
629        }
630        for next in [&self.ll_next, &self.ml_next, &self.of_next] {
631            // Same ownership test as the previous slots, and for the same
632            // reason: a dictionary-seeded axis that settles on a predefined
633            // table parks the CACHE's handle here, and that cache reports the
634            // table itself.
635            if let Some(handle) = next
636                && SharedFseTable::strong_count(handle) == 1
637            {
638                total += per_table;
639            }
640        }
641        // Where there is neither a pointer atomic nor `critical-section` to
642        // guard a process-wide cache, each default table is an owned box built
643        // per compressor, so it is this struct's allocation to report. With the
644        // cache it is a `&'static` shared by every compressor and counts as
645        // nothing.
646        #[cfg(not(any(target_has_atomic = "ptr", feature = "critical-section")))]
647        {
648            total += 3 * core::mem::size_of::<FSETable>();
649        }
650        total
651    }
652
653    pub fn new() -> Self {
654        Self {
655            ll_next: None,
656            ml_next: None,
657            of_next: None,
658            ll_default: default_ll_table(),
659            ll_previous: None,
660            ml_default: default_ml_table(),
661            ml_previous: None,
662            of_default: default_of_table(),
663            of_previous: None,
664        }
665    }
666
667    /// Borrow the LL default table as `&FSETable`.
668    ///
669    /// Test-only now: the encoder and the estimator both destructure
670    /// `FseTables` so they can hold a `*_next` slot mutably, and a method
671    /// borrowing the whole struct cannot coexist with that. Tests keep it
672    /// because they touch one table at a time.
673    ///
674    /// Abstracts the cfg
675    /// split in [`crate::fse::fse_encoder::FseDefaultTable`] —
676    /// `&'static FSETable` (atomic / `critical-section`) auto-derefs
677    /// directly; `Box<FSETable>` (cache-less no-atomic) derefs
678    /// through `Box`. Both arms yield `&FSETable` uniformly so
679    /// downstream consumers can stay cfg-agnostic.
680    #[inline]
681    #[allow(clippy::borrow_deref_ref)]
682    #[cfg(test)]
683    pub(crate) fn ll_default_ref(&self) -> &FSETable {
684        &*self.ll_default
685    }
686
687    /// Borrow the ML default table as `&FSETable`. See [`Self::ll_default_ref`].
688    #[inline]
689    #[allow(clippy::borrow_deref_ref)]
690    #[cfg(test)]
691    pub(crate) fn ml_default_ref(&self) -> &FSETable {
692        &*self.ml_default
693    }
694
695    /// Borrow the OF default table as `&FSETable`. See [`Self::ll_default_ref`].
696    #[inline]
697    #[allow(clippy::borrow_deref_ref)]
698    #[cfg(test)]
699    pub(crate) fn of_default_ref(&self) -> &FSETable {
700        &*self.of_default
701    }
702}
703
704const PRESPLIT_BLOCK_MIN: usize = 3500;
705const PRESPLIT_THRESHOLD_PENALTY_RATE: u64 = 16;
706const PRESPLIT_THRESHOLD_BASE: u64 = PRESPLIT_THRESHOLD_PENALTY_RATE - 2;
707const PRESPLIT_THRESHOLD_PENALTY: i32 = 3;
708const PRESPLIT_CHUNK_SIZE: usize = 8 << 10;
709const PRESPLIT_HASH_LOG_MAX: usize = 10;
710const PRESPLIT_HASH_TABLE_SIZE: usize = 1 << PRESPLIT_HASH_LOG_MAX;
711const PRESPLIT_KNUTH: u32 = 0x9E37_79B9;
712/// Upstream zstd `SEGMENT_SIZE` in `ZSTD_splitBlock_fromBorders` (`zstd_preSplit.c:201`).
713/// Two `SEGMENT_SIZE`-byte fingerprints — one from the start, one from the end —
714/// drive the cheap border heuristic; a third one from the middle disambiguates
715/// where in the block the transition sits.
716const PRESPLIT_BORDERS_SEGMENT: usize = 512;
717
718#[derive(Clone)]
719struct PreSplitFingerprint {
720    events: [u32; PRESPLIT_HASH_TABLE_SIZE],
721    nb_events: usize,
722}
723
724impl Default for PreSplitFingerprint {
725    fn default() -> Self {
726        Self {
727            events: [0; PRESPLIT_HASH_TABLE_SIZE],
728            nb_events: 0,
729        }
730    }
731}
732
733/// Grow `out` ahead of the next block so block emission never lands on an
734/// amortized-doubling reallocation mid-frame (whose transient old+new copy
735/// spikes peak memory to ~3x the output), sizing the reservation from the
736/// compression ratio observed so far instead of the whole-input worst case.
737///
738/// `blocks_start` is where this frame's blocks begin in `out`, `consumed`
739/// the input bytes already emitted as blocks, `remaining` the input
740/// bytes still to compress (an estimate is fine: a low one only means one
741/// more re-estimate later), and `block_capacity` the active block-size cap
742/// (`FrameCompressor::block_capacity`) so a small `targetCBlockSize` does
743/// not keep a 128 KiB floor in the buffer or undercount header density.
744/// Incompressible input re-estimates to ~the full `compress_bound` after
745/// the first block — the old up-front policy's worst case — while
746/// compressible input stays at output scale.
747fn reserve_for_next_block(
748    out: &mut Vec<u8>,
749    blocks_start: usize,
750    consumed: u64,
751    remaining: usize,
752    block_capacity: usize,
753) {
754    // Worst-case single-block output: 3-byte header + raw payload, plus
755    // slack for the 4-byte frame checksum trailer and a few extra sub-block
756    // headers from the post-split emitters, so neither can reallocate.
757    let block_bound = remaining.min(block_capacity) + 3 + 16;
758    if out.capacity() - out.len() >= block_bound {
759        return;
760    }
761    let produced = (out.len() - blocks_start) as u64;
762    let estimate = if consumed == 0 {
763        // No ratio signal yet (capacity exhausted before the first block —
764        // only reachable with a caller-shrunk `out`): one block's bound.
765        block_bound
766    } else {
767        // remaining * observed ratio + per-block headers + 1/16 slack so a
768        // slightly-worsening tail doesn't force a reallocation per block.
769        // u128 keeps the product exact for multi-GiB frames.
770        let scaled = ((remaining as u128 * produced as u128) / consumed as u128) as u64;
771        let headers = (remaining as u64 / block_capacity.max(1) as u64 + 1) * 3;
772        usize::try_from(scaled + scaled / 16 + headers + 64).unwrap_or(usize::MAX)
773    };
774    // `reserve_exact`: the estimate already carries its own slack, and the
775    // whole-buffer doubling policy is exactly what this function exists to
776    // avoid. The `produced`-sized floor keeps growth geometric when the
777    // ratio estimate lands BELOW one block's bound (highly compressible
778    // input): without it every block would trigger a block-sized
779    // reallocation — O(blocks) buffer copies — while with it the buffer at
780    // least doubles its produced span per reallocation (O(log) copies) and
781    // the peak stays at output scale.
782    out.reserve_exact(estimate.max(block_bound + produced as usize));
783}
784
785/// The rate and the width stay ARGUMENTS here, though upstream generates one
786/// function per tier and says the speed of the pass relies on compile-time
787/// constant propagation (`zstd_preSplit.c:32`, `ZSTD_GEN_RECORD_FINGERPRINT` at
788/// `:87`). Monomorphising the walk on both was tried and measured on 8 MiB of
789/// repeated log lines at level 6, three interleaved readings a side: cycles
790/// overlapping (966-984 M against 973-997 M) and instructions UP, 1,252 M to
791/// 1,268 M. Rust inlines this into a walk the caller already picked a tier for,
792/// so the propagation is there without the four copies.
793fn presplit_hash2(bytes: &[u8], hash_log: usize) -> usize {
794    debug_assert!(hash_log >= 8);
795    if hash_log == 8 {
796        return bytes[0] as usize;
797    }
798    debug_assert!(hash_log <= PRESPLIT_HASH_LOG_MAX);
799    let value = u16::from_le_bytes([bytes[0], bytes[1]]) as u32;
800    (value.wrapping_mul(PRESPLIT_KNUTH) >> (32 - hash_log)) as usize
801}
802
803fn presplit_record_fingerprint(
804    fp: &mut PreSplitFingerprint,
805    src: &[u8],
806    sampling_rate: usize,
807    hash_log: usize,
808) {
809    // Only the slots this hash can reach. The table is sized for the widest
810    // hash the splitter uses; a tier on a narrower one was clearing four times
811    // the memory it then touched, sixteen times a block.
812    fp.events[..1usize << hash_log].fill(0);
813    fp.nb_events = 0;
814    if src.len() < 2 {
815        return;
816    }
817    let limit = src.len() - 1;
818    let mut n = 0usize;
819    while n < limit {
820        fp.events[presplit_hash2(&src[n..], hash_log)] += 1;
821        n += sampling_rate;
822    }
823    // Upstream zstd parity: zstd_preSplit.c records the integer division, not the
824    // rounded-up number of sampled events from the loop above.
825    fp.nb_events += limit / sampling_rate;
826}
827
828/// Single-byte histogram pass — matches upstream zstd `HIST_add` over a small
829/// segment with `hashLog == 8` (the `hash2` shortcut at
830/// `zstd_preSplit.c:36` returns the raw byte). The byChunks path uses
831/// 2-byte hashing for `hashLog >= 9`; this helper exists so the borders
832/// heuristic doesn't pay for that wider hash on its 512-byte windows.
833fn presplit_record_byte_histogram(fp: &mut PreSplitFingerprint, src: &[u8]) {
834    fp.events.fill(0);
835    for &b in src {
836        fp.events[b as usize] += 1;
837    }
838    // Upstream zstd `HIST_add` returns the maximum symbol; the caller then sets
839    // `nbEvents = SEGMENT_SIZE` explicitly (see `zstd_preSplit.c:213`).
840    fp.nb_events = src.len();
841}
842
843fn presplit_distance(lhs: &PreSplitFingerprint, rhs: &PreSplitFingerprint, hash_log: usize) -> u64 {
844    let slots = 1usize << hash_log;
845    // 64-bit, as upstream's `fpDistance` is (`zstd_preSplit.c`). Both factors
846    // are per-block sample counts bounded by the block size, so each product is
847    // under 2^34 and the sum over at most 2^10 slots is under 2^44 — the 128-bit
848    // arithmetic this used to do could not overflow either, it just cost several
849    // instructions a slot on a loop the splitter runs sixteen times a block.
850    let rn = rhs.nb_events as u64;
851    let ln = lhs.nb_events as u64;
852    // Four accumulators over slices the optimiser can see the length of, which
853    // is what lets this go wide: the splitter runs this loop once per chunk and
854    // sixteen times a block, so it counts more slots than the sampling reads
855    // bytes.
856    let (l, r) = (&lhs.events[..slots], &rhs.events[..slots]);
857    let mut acc = [0u64; 4];
858    let mut idx = 0;
859    while idx + 4 <= slots {
860        for lane in 0..4 {
861            let left = u64::from(l[idx + lane]) * rn;
862            let right = u64::from(r[idx + lane]) * ln;
863            acc[lane] += left.abs_diff(right);
864        }
865        idx += 4;
866    }
867    let mut distance = acc[0] + acc[1] + acc[2] + acc[3];
868    while idx < slots {
869        distance += (u64::from(l[idx]) * rn).abs_diff(u64::from(r[idx]) * ln);
870        idx += 1;
871    }
872    distance
873}
874
875fn presplit_fingerprints_differ(
876    reference: &PreSplitFingerprint,
877    new_fp: &PreSplitFingerprint,
878    penalty: i32,
879    hash_log: usize,
880) -> bool {
881    debug_assert!(reference.nb_events > 0);
882    debug_assert!(new_fp.nb_events > 0);
883    let p50 = reference.nb_events as u64 * new_fp.nb_events as u64;
884    let deviation = presplit_distance(reference, new_fp, hash_log);
885    // Plain `*`: p50 <= (block-sample-count)^2 and the (base+penalty) factor is
886    // a small constant, so the product stays well under u64::MAX.
887    let threshold =
888        p50 * (PRESPLIT_THRESHOLD_BASE + penalty as u64) / PRESPLIT_THRESHOLD_PENALTY_RATE;
889    deviation >= threshold
890}
891
892fn presplit_merge_events(
893    acc: &mut PreSplitFingerprint,
894    new_fp: &PreSplitFingerprint,
895    hash_log: usize,
896) {
897    // Plain `+`: `acc` accumulates only the chunks of a single block (caller
898    // loops within one block, <= MAX_BLOCK_SIZE), so the merged sample counts
899    // stay far under u32 / usize bounds — no overflow.
900    //
901    // Only the slots this tier's hash reaches: the rest are never recorded into
902    // and never read, so merging them is a quarter of the table for nothing.
903    // Paired slices rather than indices, so the bound is one check instead of
904    // one per slot and the adds can go wide.
905    let slots = 1usize << hash_log;
906    for (into, from) in acc.events[..slots].iter_mut().zip(&new_fp.events[..slots]) {
907        *into += *from;
908    }
909    acc.nb_events += new_fp.nb_events;
910}
911
912fn split_block_by_chunks(block: &[u8], level: usize) -> usize {
913    debug_assert_eq!(block.len(), MAX_BLOCK_SIZE as usize);
914    debug_assert!((1..=4).contains(&level));
915    let (sampling_rate, hash_log) = match level - 1 {
916        0 => (43, 8),
917        1 => (11, 9),
918        2 => (5, 10),
919        _ => (1, 10),
920    };
921
922    let mut past = PreSplitFingerprint::default();
923    let mut new_events = PreSplitFingerprint::default();
924    let mut penalty = PRESPLIT_THRESHOLD_PENALTY;
925    presplit_record_fingerprint(
926        &mut past,
927        &block[..PRESPLIT_CHUNK_SIZE],
928        sampling_rate,
929        hash_log,
930    );
931    // No pre-check on the ends before the walk. It reads as free — two
932    // fingerprints against the sixty-three the walk takes — but a sample of the
933    // ends cannot stand in for the walk: an A-B-A block has matching ends and a
934    // boundary in the middle, and four megabytes repeated at nearly the window
935    // distance lost 13% to exactly that. Upstream has no such check either; the
936    // cheap gate that keeps the walk off hopeless input is `savings`, which the
937    // caller already applies.
938    let mut pos = PRESPLIT_CHUNK_SIZE;
939    while pos <= block.len() - PRESPLIT_CHUNK_SIZE {
940        presplit_record_fingerprint(
941            &mut new_events,
942            &block[pos..pos + PRESPLIT_CHUNK_SIZE],
943            sampling_rate,
944            hash_log,
945        );
946        if presplit_fingerprints_differ(&past, &new_events, penalty, hash_log) {
947            return pos;
948        }
949        presplit_merge_events(&mut past, &new_events, hash_log);
950        if penalty > 0 {
951            penalty -= 1;
952        }
953        pos += PRESPLIT_CHUNK_SIZE;
954    }
955    block.len()
956}
957
958/// Upstream zstd port of `ZSTD_splitBlock_fromBorders` (`zstd_preSplit.c:198`).
959/// Records two 512-byte byte-histograms — one from each end of a 128 KB
960/// block — and a third from the middle as a tie-breaker; returns either
961/// a quantised split point (32 KB / 64 KB / 96 KB) or the full block
962/// size when the two ends look indistinguishable. Cheaper than the
963/// chunk-based path because it touches at most 1.5 KB of input
964/// regardless of block size.
965fn split_block_from_borders(block: &[u8]) -> usize {
966    debug_assert_eq!(block.len(), MAX_BLOCK_SIZE as usize);
967    let block_size = block.len();
968    let mut past = PreSplitFingerprint::default();
969    let mut new_fp = PreSplitFingerprint::default();
970    presplit_record_byte_histogram(&mut past, &block[..PRESPLIT_BORDERS_SEGMENT]);
971    presplit_record_byte_histogram(&mut new_fp, &block[block_size - PRESPLIT_BORDERS_SEGMENT..]);
972    // Upstream zstd uses `penalty = 0, hash_log = 8` — i.e. raw byte histogram
973    // distance with no threshold padding (`zstd_preSplit.c:214`).
974    if !presplit_fingerprints_differ(&past, &new_fp, 0, 8) {
975        return block_size;
976    }
977
978    let mut middle = PreSplitFingerprint::default();
979    let mid_start = block_size / 2 - PRESPLIT_BORDERS_SEGMENT / 2;
980    presplit_record_byte_histogram(
981        &mut middle,
982        &block[mid_start..mid_start + PRESPLIT_BORDERS_SEGMENT],
983    );
984
985    let dist_from_begin = presplit_distance(&past, &middle, 8);
986    let dist_from_end = presplit_distance(&new_fp, &middle, 8);
987    // Upstream zstd `SEGMENT_SIZE * SEGMENT_SIZE / 3` (`zstd_preSplit.c:221`):
988    // if the middle is roughly equidistant from both ends, the change
989    // sits near the centre — split at the midpoint.
990    let min_distance = (PRESPLIT_BORDERS_SEGMENT as u64) * (PRESPLIT_BORDERS_SEGMENT as u64) / 3;
991    if dist_from_begin.abs_diff(dist_from_end) < min_distance {
992        return 64 * 1024;
993    }
994    // Larger `dist_from_begin` (i.e. `middle` farther from the head
995    // fingerprint, equivalently closer to the tail) means the new
996    // statistics already dominate the centre — the transition
997    // happened EARLY → emit a small 32 KB head and let the 96 KB
998    // tail absorb the rest. Inverse case: `dist_from_end` larger
999    // (middle still resembles the head) means the transition is
1000    // LATE → emit a 96 KB head so the trailing 32 KB carries the
1001    // new statistics alone.
1002    if dist_from_begin > dist_from_end {
1003        32 * 1024
1004    } else {
1005        96 * 1024
1006    }
1007}
1008
1009/// XXH64 (low 32 bits, seed 0) over `data`. Shared helper for the
1010/// per-physical-block checksum sidecar so encoder and decoder hash
1011/// the exact same byte ranges with the exact same parameters. Gated
1012/// at `all(lsm, hash)` because the only consumer is the lsm-side
1013/// `block_checksums` sidecar; non-lsm builds carry no reference to
1014/// this helper at all.
1015#[cfg(all(feature = "lsm", feature = "hash"))]
1016#[inline]
1017pub(crate) fn xxh64_block_low32(data: &[u8]) -> u32 {
1018    let mut h = XxHash64::with_seed(0);
1019    h.write(data);
1020    h.finish() as u32
1021}
1022
1023/// Bench-only entry point for the upstream zstd-parity comparator test in
1024/// `tests/block_splitter_parity.rs`. Dispatches to the same
1025/// `_from_borders` (split_level == 0) / `_by_chunks` (split_level ∈
1026/// 1..=4) ports that `optimal_block_size` itself routes
1027/// through. Caller is responsible for passing exactly
1028/// `MAX_BLOCK_SIZE` bytes (per upstream zstd `ZSTD_splitBlock` contract —
1029/// "@blockSize must be == 128 KB" in `zstd_preSplit.h`).
1030#[cfg(feature = "bench-internals")]
1031pub(crate) fn block_splitter_decision_for_bench(block: &[u8], split_level: usize) -> usize {
1032    assert_eq!(
1033        block.len(),
1034        MAX_BLOCK_SIZE as usize,
1035        "block_splitter_decision_for_bench expects exactly MAX_BLOCK_SIZE bytes"
1036    );
1037    assert!(
1038        split_level <= 4,
1039        "block_splitter_decision_for_bench: split_level must be in 0..=4, got {split_level}"
1040    );
1041    if split_level == 0 {
1042        split_block_from_borders(block)
1043    } else {
1044        split_block_by_chunks(block, split_level)
1045    }
1046}
1047
1048/// Pull a pre-split window into cache with one bandwidth-bound sequential
1049/// pass before the strided fingerprint histogram + match scan read it.
1050///
1051/// The borrowed (no-copy) over-window path matches in place on the caller's
1052/// input, so the pre-split fingerprint is the FIRST touch of that 128 KiB
1053/// region — a cache-cold read. `presplit_record_fingerprint` reads it with a
1054/// `sampling_rate` stride and interleaved random writes into the 1 KiB events
1055/// table, a latency-bound pattern that pays full DRAM miss latency per line
1056/// (measured ~3x the cost of an ERMS streaming read of the same bytes). The
1057/// owned path never hits this because its history-mirror copy already warmed
1058/// the bytes; this restores that warmth without the copy's write half. One
1059/// dependent load per 64-byte line (the i9 line size) streams under the
1060/// hardware prefetcher, so the cold read is paid once at memory bandwidth and
1061/// every subsequent strided sample lands in L1/L2. `black_box` keeps the loop
1062/// from being optimized away as a dead read.
1063#[inline]
1064fn warm_presplit_window(window: &[u8]) {
1065    let mut acc = 0u8;
1066    let mut i = 0usize;
1067    while i < window.len() {
1068        acc ^= window[i];
1069        i += 64;
1070    }
1071    core::hint::black_box(acc);
1072}
1073
1074/// [`optimal_block_size_with`] at a level's default pre-split tier; the
1075/// frame loop resolves the tier itself, so only tests and the
1076/// `bench-internals` block-boundary probe read this form.
1077#[cfg(any(test, feature = "bench-internals"))]
1078pub(crate) fn optimal_block_size(
1079    level: CompressionLevel,
1080    block: &[u8],
1081    remaining_src_size: usize,
1082    block_size_max: usize,
1083    savings: i64,
1084) -> usize {
1085    optimal_block_size_with(
1086        crate::encoding::levels::config::level_pre_split(level),
1087        block,
1088        remaining_src_size,
1089        block_size_max,
1090        savings,
1091    )
1092}
1093
1094/// [`optimal_block_size`] with the pre-split level already resolved
1095/// (`None` = never split, only full blocks).
1096///
1097/// Out of line on purpose: inlined into the per-block frame loop it grew the
1098/// loop body and shifted the code layout around `run_fast_kernel_block`,
1099/// costing the Fast levels 17-24 % on 1 MiB+ inputs on x86 (measured on the
1100/// i9; the kernel's own instructions were byte-identical). One call per block
1101/// is noise; the compact caller is not.
1102#[inline(never)]
1103pub(crate) fn optimal_block_size_with(
1104    pre_split: Option<usize>,
1105    block: &[u8],
1106    remaining_src_size: usize,
1107    block_size_max: usize,
1108    savings: i64,
1109) -> usize {
1110    let Some(split_level) = pre_split else {
1111        return remaining_src_size.min(block_size_max);
1112    };
1113    if remaining_src_size < MAX_BLOCK_SIZE as usize || block_size_max < MAX_BLOCK_SIZE as usize {
1114        return remaining_src_size.min(block_size_max);
1115    }
1116    if savings < 3 {
1117        return MAX_BLOCK_SIZE as usize;
1118    }
1119    if block.len() < MAX_BLOCK_SIZE as usize {
1120        return remaining_src_size.min(block_size_max);
1121    }
1122    // Upstream zstd `ZSTD_splitBlock` dispatch (`zstd_preSplit.c:234`):
1123    // `split_level == 0` → cheap borders heuristic;
1124    // `split_level == 1..=4` → byChunks with internal sampling level
1125    // `split_level - 1`.
1126    // NOTE: gating the sampling tier behind the cheap borders tier (run it
1127    // only when the block's ends already disagree) was measured and REJECTED:
1128    // it costs 2.8-3.4% of the decode corpus across levels 3-15, because the
1129    // boundaries the sampling finds are mostly inside blocks whose ends look
1130    // alike. The cost of the sampling on uniform input has to come off some
1131    // other way.
1132    let raw_split = if split_level == 0 {
1133        split_block_from_borders(&block[..MAX_BLOCK_SIZE as usize])
1134    } else {
1135        split_block_by_chunks(&block[..MAX_BLOCK_SIZE as usize], split_level)
1136    };
1137    raw_split
1138        .max(PRESPLIT_BLOCK_MIN)
1139        .min(MAX_BLOCK_SIZE as usize)
1140}
1141
1142/// Load `dict` into the matcher for the frame about to start, `size_hint`
1143/// bytes long when known, once `state.strategy_tag` holds the strategy the
1144/// frame runs. Shared by the frame compressor and the streaming context so a
1145/// reused one of either loads its dictionary the same way.
1146///
1147/// A dictionary a reused matcher kept resident across the reset (its bytes
1148/// and index still in place) only gets its offset history back: priming it
1149/// again would commit it into the window a second time. Otherwise it is
1150/// primed, or restored from the matcher's copy-mode snapshot when the source
1151/// is past the strategy's attach cutoff.
1152pub(crate) fn load_frame_dictionary<M: Matcher>(
1153    state: &mut CompressState<M>,
1154    level: CompressionLevel,
1155    dict: &EncoderDictionary,
1156    size_hint: Option<u64>,
1157) {
1158    // Upstream zstd `ZSTD_shouldAttachDict` (`zstd_compress.c`): a
1159    // precomputed-dictionary table is COPIED into the working context
1160    // only when the source is larger than a per-strategy cutoff; at or
1161    // below it (and for unknown size) the upstream zstd ATTACHES the dictionary
1162    // tables by reference (no per-frame table touch at all). We don't
1163    // have an attach-by-reference path yet, so:
1164    //   - large source (> cutoff): reuse the captured prime snapshot
1165    //     (a table copy) instead of re-hashing the dictionary — the
1166    //     upstream zstd COPY regime, where the copy is cheaper than re-priming;
1167    //   - small / unknown source: re-prime (the snapshot copy of the
1168    //     whole table would cost MORE than the sparse re-prime here,
1169    //     which is exactly why the upstream zstd attaches by reference instead).
1170    // `attachDictSizeCutoffs` per strategy: fast 8K, dfast 16K,
1171    // greedy/lazy/btopt 32K, btultra/btultra2 8K. Expressed as the
1172    // ceil-log bucket (8K = 2^13, 16K = 2^14, 32K = 2^15) so the
1173    // decision uses the SAME bucketed representation as the driver's
1174    // attach/copy gate (`reset_size_log`) — comparing
1175    // `source_size_ceil_log(hint)` on the full u64 avoids the `as usize`
1176    // truncation that could diverge from the driver on 32-bit targets.
1177    // For a power-of-two cutoff `2^k`, `ceil_log2(hint) > k` is exactly
1178    // `hint > 2^k`, so this is identical to the raw `hint > cutoff` on
1179    // 64-bit.
1180    let cutoff_log = match state.strategy_tag {
1181        // Keep the copy-snapshot gate in sync with the matcher's own
1182        // attach cutoff, so Fast never captures or restores a snapshot
1183        // for a mode it did not resolve.
1184        crate::encoding::strategy::StrategyTag::Fast => {
1185            crate::encoding::levels::config::FAST_ATTACH_DICT_CUTOFF_LOG
1186        }
1187        crate::encoding::strategy::StrategyTag::BtUltra
1188        | crate::encoding::strategy::StrategyTag::BtUltra2 => 13,
1189        crate::encoding::strategy::StrategyTag::Dfast => 14,
1190        crate::encoding::strategy::StrategyTag::Greedy
1191        | crate::encoding::strategy::StrategyTag::Lazy
1192        | crate::encoding::strategy::StrategyTag::Btlazy2
1193        | crate::encoding::strategy::StrategyTag::BtOpt => 15,
1194    };
1195    if state.matcher.dictionary_is_resident() {
1196        // Re-borrow fast path: the previous frame's reset kept this
1197        // dict's bytes + cached index resident, so skip the re-commit /
1198        // re-index and only reapply the offset history.
1199        state
1200            .matcher
1201            .reapply_resident_dictionary(dict.inner.offset_hist);
1202        return;
1203    }
1204    let prefer_copy_snapshot = size_hint
1205        .is_some_and(|s| crate::encoding::levels::config::source_size_ceil_log(s) > cutoff_log);
1206    let restored = prefer_copy_snapshot && state.matcher.restore_primed_dictionary(level);
1207    if !restored {
1208        state
1209            .matcher
1210            .prime_with_dictionary(dict.inner.dict_content.as_slice(), dict.inner.offset_hist);
1211        if prefer_copy_snapshot {
1212            state.matcher.capture_primed_dictionary(level);
1213        }
1214    }
1215}
1216
1217/// Record in `state` the strategy the matcher runs for the next frame (the
1218/// honoured public override, else the size- and dictionary-adaptive
1219/// resolution in `params`) and its pre-split tier; the literal gates and the
1220/// block splitter read these, and upstream indexes `splitLevels` by the
1221/// effective strategy too. Shared by the frame compressor and the streaming
1222/// encoder so both entry points cut and gate blocks identically.
1223pub(crate) fn sync_effective_strategy<M: Matcher>(
1224    state: &mut CompressState<M>,
1225    level: CompressionLevel,
1226    params: &crate::encoding::levels::config::LevelParams,
1227    strategy_override: Option<(crate::encoding::strategy::StrategyTag, u8)>,
1228) {
1229    match strategy_override {
1230        Some((tag, lazy_depth)) => {
1231            state.strategy_tag = tag;
1232            state.pre_split = Some(crate::encoding::levels::config::pre_split_for(
1233                tag, lazy_depth,
1234            ));
1235        }
1236        None => {
1237            state.strategy_tag = params.strategy_tag;
1238            state.pre_split = if matches!(level, CompressionLevel::Uncompressed) {
1239                None
1240            } else {
1241                params.pre_split()
1242            };
1243        }
1244    }
1245}
1246
1247/// Upstream `ZSTD_literalsCompressionIsDisabled`: an explicit
1248/// [`LiteralCompressionMode`] decides outright; under `Auto`, raw literals iff
1249/// the EFFECTIVE cParams are the fast strategy with `targetLength > 0`.
1250/// The effective strategy tag gates this (a strategy override can move a
1251/// negative level off fast). For the fast strategy the level table sets
1252/// `targetLength > 0` exactly on the negative (acceleration) rows (a CDict's
1253/// too), so where [`gate_target_length`] has no value `level < 0` is that
1254/// test.
1255pub(crate) fn literal_compression_disabled(
1256    strategy_tag: crate::encoding::strategy::StrategyTag,
1257    level: CompressionLevel,
1258    target_length_override: Option<u32>,
1259    mode: LiteralCompressionMode,
1260) -> bool {
1261    match mode {
1262        LiteralCompressionMode::Disable => true,
1263        LiteralCompressionMode::Enable => false,
1264        LiteralCompressionMode::Auto => {
1265            strategy_tag == crate::encoding::strategy::StrategyTag::Fast
1266                && target_length_override.map_or_else(
1267                    || matches!(level, CompressionLevel::Level(n) if n < 0),
1268                    |tl| tl > 0,
1269                )
1270        }
1271    }
1272}
1273
1274/// The `target_length` the raw-literals gate reads for a frame: the caller's,
1275/// else, for a dictionary frame a strategy knob moves onto the fast strategy,
1276/// the CDict row's, which the fast matcher runs as its step (upstream
1277/// `ZSTD_overrideCParams` replaces only the strategy, and
1278/// `ZSTD_literalsCompressionIsDisabled` reads the resulting cParams). Every
1279/// other frame's fast targetLength is what the level says, which
1280/// [`literal_compression_disabled`] derives itself.
1281pub(crate) fn gate_target_length(
1282    level: CompressionLevel,
1283    tuning: &FrameTuning,
1284    dictionary: Option<&EncoderDictionary>,
1285) -> Option<u32> {
1286    match (tuning.target_length, tuning.strategy, dictionary) {
1287        (Some(target_length), _, _) => Some(target_length),
1288        (None, Some((crate::encoding::strategy::StrategyTag::Fast, _)), Some(dict))
1289            if !dict.inner.dict_content.is_empty() =>
1290        {
1291            Some(
1292                crate::encoding::cparams::get_cdict_cparams(
1293                    crate::encoding::levels::config::numeric_level(level),
1294                    dict.sizes().serialized,
1295                    &crate::encoding::parameters::ParamOverrides::default(),
1296                )
1297                .target_length,
1298            )
1299        }
1300        _ => None,
1301    }
1302}
1303
1304/// The level params the matcher's reset resolves for a frame: through the
1305/// dictionary's CDict tier when a dictionary is in play, else by source size.
1306/// Resolved without the caller's knobs: the frame gates read only the
1307/// strategy from these, and a strategy override replaces it
1308/// ([`sync_effective_strategy`]) on a dictionary frame as on any other.
1309pub(crate) fn resolve_frame_params(
1310    level: CompressionLevel,
1311    hint: Option<u64>,
1312    dictionary: Option<&EncoderDictionary>,
1313) -> crate::encoding::levels::config::LevelParams {
1314    match dictionary {
1315        Some(dict) if !dict.inner.dict_content.is_empty() => {
1316            crate::encoding::levels::config::resolve_level_params_with_dict(
1317                level,
1318                hint,
1319                dict.sizes(),
1320                &crate::encoding::parameters::ParamOverrides::default(),
1321            )
1322            .0
1323        }
1324        _ => crate::encoding::levels::config::resolve_level_params(level, hint),
1325    }
1326}
1327
1328pub(crate) struct CompressState<M: Matcher> {
1329    pub(crate) matcher: M,
1330    /// Grid fingerprints of the frame's blocks, which is what tells the
1331    /// raw-skip that a block duplicating an earlier one is not the noise it
1332    /// looks like. See [`SeenContentGrid`](crate::encoding::incompressible::SeenContentGrid).
1333    pub(crate) seen_content: crate::encoding::incompressible::SeenContentGrid,
1334    /// Widest literal-copy kernel this CPU can run, resolved once when the
1335    /// compressor is built. The emit path reads it; it never re-probes.
1336    pub(crate) copy_tier: crate::decoding::simd_copy::ExactCopyTier,
1337    pub(crate) last_huff_table: Option<crate::huff0::huff0_encoder::HuffmanTable>,
1338    /// Recycled `HuffmanTable` buffers: when a block clears or replaces
1339    /// `last_huff_table`, the old table parks here instead of dropping, so
1340    /// the next frame's dictionary entropy seed `clone_from`s into existing
1341    /// allocations. Without this, every dict-seeded frame whose last block
1342    /// ended raw/RLE paid a fresh two-Vec table clone per frame.
1343    pub(crate) huff_table_spare: Option<crate::huff0::huff0_encoder::HuffmanTable>,
1344    /// Where a block copies `last_huff_table` before encoding, so it can be put
1345    /// back if the compressed form loses to a raw block.
1346    ///
1347    /// A slot rather than a local because the copy is per block: cloning into a
1348    /// fresh `Option` took two `Vec`s every time, while cloning into one that
1349    /// already holds a table reuses them. Restoring is a swap, which also hands
1350    /// the discarded table back here as the next block's buffer.
1351    pub(crate) huff_rollback: Option<crate::huff0::huff0_encoder::HuffmanTable>,
1352    /// The Huffman weight builder's three buffers, kept across blocks and
1353    /// frames. The cheap build path takes a tree and two weight buffers per
1354    /// call, and it runs once per block plus once per split candidate wherever
1355    /// the block splitter probes, so taking them fresh each time was the single
1356    /// largest source of per-frame allocations: 1,880 of a frame's 4,000 at
1357    /// level 3. Lives here rather than in the per-block scratch, which the
1358    /// block emitter takes out of this state while a block is in flight.
1359    pub(crate) huff_weights: crate::huff0::huff0_encoder::WeightScratch,
1360    pub(crate) fse_tables: FseTables,
1361    pub(crate) block_scratch: crate::encoding::blocks::CompressedBlockScratch,
1362    /// Offset history for repeat offset encoding: [rep0, rep1, rep2].
1363    /// Initialized to [1, 4, 8] per RFC 8878 §3.1.2.5.
1364    pub(crate) offset_hist: [u32; 3],
1365    /// Strategy tag resolved from the current `CompressionLevel` at every
1366    /// `matcher.reset()` call. Used by the literal-compression gates
1367    /// (`min_literals_to_compress`, `min_gain`) in
1368    /// `encoding::blocks::compressed` to mirror upstream zstd's strategy-aware
1369    /// thresholds (`zstd_compress_literals.c:114-127, 187-188`).
1370    ///
1371    /// **Invariant (required of every construction site):** must be
1372    /// initialized from the active `CompressionLevel` via
1373    /// `StrategyTag::for_compression_level`, and re-synced from the
1374    /// active level alongside every `matcher.reset()` call so the
1375    /// level-aware gates stay correct after a level change. The two
1376    /// reset sites that own this sync are `FrameCompressor::compress`
1377    /// and `StreamingEncoder::ensure_frame_started`. There is no
1378    /// `Default` impl — production constructors
1379    /// (`FrameCompressor::new`, `new_with_matcher`, the streaming
1380    /// encoder constructor) plumb this explicitly. Tests that build
1381    /// `CompressState` by hand must also supply a value.
1382    pub(crate) strategy_tag: crate::encoding::strategy::StrategyTag,
1383    /// Pre-split tier of the effective strategy (upstream `splitLevels`),
1384    /// synced with `strategy_tag`; `None` never pre-splits (raw frames).
1385    pub(crate) pre_split: Option<u8>,
1386    /// Whether the HUF literal table build runs the #167 table-log search
1387    /// (`true`) or the cheap single-build (`false`). The search is a clean
1388    /// ratio win over upstream zstd but costs ~1.5 us per literal section —
1389    /// negligible on large inputs, ~20% on small ones. The Fast and DoubleFast
1390    /// matchers are byte-faithful to upstream zstd, so the cheap path ties them;
1391    /// the search is therefore gated ON only for large (> 128 KiB) Fast and
1392    /// DoubleFast frames. Higher strategies always keep it (their matchers
1393    /// diverge, making the search load-bearing for ratio). Set per frame
1394    /// alongside `strategy_tag` via [`huf_search_enabled`].
1395    pub(crate) huf_optimal_search: bool,
1396    /// Mirror of upstream zstd's `ZSTD_literalsCompressionIsDisabled`
1397    /// (zstd_compress_internal.h): in the default (`auto`) literal-compression
1398    /// mode the literals section is emitted RAW (no Huffman) when
1399    /// `strategy == ZSTD_fast && targetLength > 0`. For the levels we resolve,
1400    /// that is exactly the negative levels (Fast strategy with `targetLength =
1401    /// -level > 0`; L1/L2 are Fast with `targetLength == 0`). C trades the
1402    /// literal-Huffman pass for speed there, so matching it keeps both the frame
1403    /// size and the encode cost in parity on the negative band. Set per frame
1404    /// alongside `strategy_tag`.
1405    pub(crate) literal_compression_disabled: bool,
1406}
1407
1408/// Whether the HUF literal build should run the #167 table-log search for a
1409/// frame of `source_size` bytes (see [`CompressState::huf_optimal_search`]).
1410/// Upstream gates the optimal-depth tableLog probe to
1411/// `HUF_OPTIMAL_DEPTH_THRESHOLD = ZSTD_btultra` (huf.h:117): only btultra /
1412/// btultra2 search the tableLog, every lower strategy (fast .. btopt) takes the
1413/// single-shot fast path (`HUF_optimalTableLog`, huf_compress.c:1284-1287).
1414/// Mirror that so our literal tableLog choice tracks upstream's instead of
1415/// spending the search to beat it on ratio at a speed cost.
1416pub(crate) fn huf_search_enabled(
1417    strategy: crate::encoding::strategy::StrategyTag,
1418    _source_size: Option<u64>,
1419) -> bool {
1420    use crate::encoding::strategy::StrategyTag;
1421    matches!(strategy, StrategyTag::BtUltra | StrategyTag::BtUltra2)
1422}
1423
1424impl<M: Matcher> CompressState<M> {
1425    /// Clears `last_huff_table`, parking the table's buffers in
1426    /// `huff_table_spare` for reuse instead of dropping them.
1427    #[inline]
1428    /// Heap bytes the compressor keeps between blocks and frames beyond the
1429    /// match finder: the FSE tables both slots of each axis hold, the rollback
1430    /// slot the emit paths copy a Huffman table into before a block that may not
1431    /// be kept, and the block scratch with everything it holds — its literal and
1432    /// sequence buffers, the splitter's workspace, and the nested estimator
1433    /// scratch.
1434    ///
1435    /// All of it survives a frame, so a caller sizing a context has to see it.
1436    pub(crate) fn retained_scratch_heap_size(&self) -> usize {
1437        self.fse_tables.heap_size()
1438            + self
1439                .huff_rollback
1440                .as_ref()
1441                .map_or(0, |table| table.heap_size())
1442            + self.block_scratch.retained_heap_size()
1443    }
1444
1445    pub(crate) fn clear_huff_table(&mut self) {
1446        if let Some(table) = self.last_huff_table.take() {
1447            self.park_huff_table(table);
1448        }
1449    }
1450
1451    /// Replaces `last_huff_table` with `table`, parking any displaced table
1452    /// in `huff_table_spare` for reuse.
1453    #[inline]
1454    pub(crate) fn replace_huff_table(&mut self, table: crate::huff0::huff0_encoder::HuffmanTable) {
1455        if let Some(old) = self.last_huff_table.replace(table) {
1456            self.park_huff_table(old);
1457        }
1458    }
1459
1460    /// Keeps a table's buffers rather than dropping them. The dictionary seed
1461    /// wants one spare to `clone_from` into, once per frame; every further
1462    /// table a block displaces goes to the weight builder instead, which takes
1463    /// one per block and per split candidate. Overwriting the single spare
1464    /// dropped the previous table on every block, so the builds that followed
1465    /// allocated their buffers again.
1466    #[inline]
1467    fn park_huff_table(&mut self, table: crate::huff0::huff0_encoder::HuffmanTable) {
1468        if self.huff_table_spare.is_none() {
1469            self.huff_table_spare = Some(table);
1470        } else {
1471            self.huff_weights.recycle(table);
1472        }
1473    }
1474}
1475
1476/// Per-frame setup resolved once by [`FrameCompressor::prepare_frame`] and
1477/// consumed by the block loop + [`FrameCompressor::finish_frame`]. Lets the
1478/// owned `compress()` and the borrowed one-shot path share identical
1479/// reset / dict-prime / entropy-seed setup and frame-tail emission.
1480struct FramePrep {
1481    window_size: u64,
1482    use_dictionary_state: bool,
1483    source_size_hint_known: bool,
1484    initial_size_hint: Option<u64>,
1485}
1486
1487/// Initial capacity for the `all_blocks` accumulator, by source-size hint.
1488/// The frame header is written only after all input is read (so
1489/// Frame_Content_Size is known), so compressed blocks accumulate in memory
1490/// first. Seed-size tiers (mirrors upstream zstd `ZSTD_CStreamOutSize` naming):
1491/// - tiny (`<= 4 KiB` hint): payload-bound seed, `>=` anything a tiny input's
1492///   compressed output could need.
1493/// - small (`<= 64 KiB` hint): absorbs one or two `Vec::extend` doublings
1494///   without over-allocating.
1495/// - default (one upstream zstd block, `130 KiB`): the value the rest of the encoder
1496///   is sized around; larger inputs amortise the first doublings cheaply and
1497///   the residue is dominated by internal `compress_block_encoded` buffers.
1498///
1499/// Shared by the owned (`run_owned_block_loop`) and borrowed
1500/// (`run_borrowed_block_loop`) paths so the tier table can't drift between them.
1501///
1502/// `block_capacity` (the active `targetCBlockSize` cap, or the 128 KiB
1503/// format ceiling) bounds every tier: with a small target the first
1504/// allocation tracks one capped block + header/checksum slack instead of
1505/// keeping the upstream zstd-sized floor that only later growth respects.
1506fn initial_all_blocks_cap(initial_size_hint: Option<u64>, block_capacity: usize) -> usize {
1507    const TINY_THRESHOLD: u64 = 4 * 1024;
1508    const SMALL_THRESHOLD: u64 = 64 * 1024;
1509    const TINY_CAP: usize = 4 * 1024;
1510    const SMALL_CAP: usize = 16 * 1024;
1511    const DEFAULT_CAP: usize = 130 * 1024;
1512    let first_block_cap = block_capacity + 3 + 16;
1513    match initial_size_hint {
1514        Some(h) if h <= TINY_THRESHOLD => TINY_CAP.min(first_block_cap),
1515        Some(h) if h <= SMALL_THRESHOLD => SMALL_CAP.min(first_block_cap),
1516        _ => DEFAULT_CAP.min(first_block_cap),
1517    }
1518}
1519
1520/// Per-block feeder for `run_owned_block_loop`.
1521///
1522/// `fill_block` appends source bytes to `buf` (which already holds any
1523/// carried pre-split suffix) until `buf.len() == block_capacity` or the
1524/// source is exhausted, returning `(bytes_appended, reached_eof)`.
1525/// `reached_eof` is true when no more input follows this block: either the
1526/// block could not be filled to `block_capacity`, or it filled exactly and the
1527/// source is confirmed exhausted (the slice knows its length; the reader probes
1528/// one byte ahead). An input that is an exact multiple of the block size
1529/// therefore marks its final full block `last_block` rather than emitting a
1530/// spurious trailing empty block.
1531///
1532/// The slice impl exists so the slice entry points
1533/// (`compress_independent_frame_into`, `compress_oneshot_*` fallbacks)
1534/// append with one `extend_from_slice` — the generic reader impl must
1535/// `resize` an initialized target region before `Read::read` can fill it,
1536/// which costs a zero-fill memset of the whole block on every frame.
1537pub(crate) trait OwnedBlockSource {
1538    fn fill_block(
1539        &mut self,
1540        buf: &mut Vec<u8>,
1541        block_capacity: usize,
1542        size_hint_remaining: Option<u64>,
1543    ) -> (usize, bool);
1544}
1545
1546impl OwnedBlockSource for &[u8] {
1547    fn fill_block(
1548        &mut self,
1549        buf: &mut Vec<u8>,
1550        block_capacity: usize,
1551        _size_hint_remaining: Option<u64>,
1552    ) -> (usize, bool) {
1553        let want = block_capacity - buf.len();
1554        let take = want.min(self.len());
1555        buf.extend_from_slice(&self[..take]);
1556        *self = &self[take..];
1557        // EOF when this fill could not top the block to `block_capacity`
1558        // (`take < want`) OR it exactly consumed the last input bytes
1559        // (`self` now empty). The slice knows its own length, so a block that
1560        // exactly fills capacity at end-of-input is reported as the final
1561        // block here — the loop marks it `last_block` instead of emitting a
1562        // spurious trailing empty block on the next iteration. Mirrors the C
1563        // encoder, which marks the last real block last on `ZSTD_e_end`.
1564        (take, take < want || self.is_empty())
1565    }
1566}
1567
1568/// Adapter routing a generic [`Read`] source through [`OwnedBlockSource`]:
1569/// preserves the historical sizing behaviour — an initialized target region
1570/// bounded by the source-size hint, grown (doubling, capped) only when the
1571/// hint under-counted.
1572/// `peeked` holds a single look-ahead byte: when a block fills exactly to
1573/// `block_capacity`, `fill_block` reads one more byte to learn whether the
1574/// stream ended on that boundary. A `None` from that probe sets EOF (so the
1575/// just-filled block is marked last, mirroring the C encoder on `ZSTD_e_end`);
1576/// a byte is stashed here and prepended to the next block instead of leaking a
1577/// spurious trailing empty block when the input is an exact multiple of the
1578/// block size.
1579pub(crate) struct ReaderBlockSource<Rd> {
1580    pub(crate) reader: Rd,
1581    peeked: Option<u8>,
1582}
1583
1584impl<Rd> ReaderBlockSource<Rd> {
1585    pub(crate) fn new(reader: Rd) -> Self {
1586        Self {
1587            reader,
1588            peeked: None,
1589        }
1590    }
1591}
1592
1593impl<Rd: Read> OwnedBlockSource for ReaderBlockSource<Rd> {
1594    fn fill_block(
1595        &mut self,
1596        buf: &mut Vec<u8>,
1597        block_capacity: usize,
1598        size_hint_remaining: Option<u64>,
1599    ) -> (usize, bool) {
1600        let start = buf.len();
1601        let mut filled = start;
1602        let mut reached_eof = false;
1603        // Prepend the look-ahead byte read past the previous full block. In
1604        // stream order it follows any carried pre-split suffix already in
1605        // `buf`, so it is appended after that suffix and counted as part of
1606        // this block's appended bytes.
1607        if let Some(b) = self.peeked.take() {
1608            buf.push(b);
1609            filled += 1;
1610        }
1611        // Size the read buffer to the bytes this block actually expects
1612        // rather than always zero-filling a full MAX_BLOCK_SIZE: a small
1613        // frame otherwise pays a 128 KiB `resize(_, 0)` memset per block
1614        // just to read a few KiB (the zero-fill past `filled` is then
1615        // truncated away).
1616        //
1617        // Overflow-free by construction (no `saturating_*` masking):
1618        // `filled <= block_capacity` always (the read only ever targets
1619        // `[filled..len]` with `len <= block_capacity`, and a carried-over
1620        // pre-split suffix is a `split_off` below `block_capacity`), so
1621        // `block_capacity - filled` never underflows; pinning `remaining`
1622        // to `block_capacity` before the `usize` cast keeps the cast and
1623        // the final add within `usize` on every target.
1624        let initial_target = match size_hint_remaining {
1625            Some(remaining) => {
1626                let remaining = remaining.min(block_capacity as u64) as usize;
1627                filled + remaining.min(block_capacity - filled)
1628            }
1629            // Unknown hint, or an inexact hint already met by prior blocks:
1630            // read against the full block window.
1631            None => block_capacity,
1632        };
1633        if buf.len() < initial_target {
1634            buf.resize(initial_target, 0);
1635        }
1636        loop {
1637            if reached_eof || filled == block_capacity {
1638                break;
1639            }
1640            if filled == buf.len() {
1641                // Hint under-counted the block; grow toward block_capacity
1642                // (doubling, capped) so reading continues without paying a
1643                // full-buffer zero up front. `len <= block_capacity` so the
1644                // double stays well within `usize`; `filled < block_capacity`
1645                // here (the `== block_capacity` break fired otherwise), so
1646                // `filled + 1 <= block_capacity`.
1647                let grow_to = (buf.len() * 2).clamp(filled + 1, block_capacity);
1648                buf.resize(grow_to, 0);
1649            }
1650            let read_end = buf.len();
1651            let new_bytes = self.reader.read(&mut buf[filled..read_end]).unwrap();
1652            if new_bytes == 0 {
1653                reached_eof = true;
1654                break;
1655            }
1656            filled += new_bytes;
1657        }
1658        // Look ahead one byte when the block filled exactly to capacity: a
1659        // 0-byte read means the stream ended on the block boundary, so this
1660        // block is the last one (the loop marks it `last_block`); otherwise
1661        // stash the byte for the next block. Without this, an input that is an
1662        // exact multiple of the block size would emit a spurious trailing
1663        // empty block (the next iteration reads 0 and serializes an empty
1664        // last Raw block). A blocking reader's probe read is consistent with
1665        // the existing pull model — the next `fill_block` would block on the
1666        // same byte anyway.
1667        if !reached_eof && filled == block_capacity {
1668            let mut probe = [0u8; 1];
1669            if self.reader.read(&mut probe).unwrap() == 0 {
1670                reached_eof = true;
1671            } else {
1672                self.peeked = Some(probe[0]);
1673            }
1674        }
1675        buf.truncate(filled);
1676        (filled - start, reached_eof)
1677    }
1678}
1679
1680impl<R: Read, W: Write> FrameCompressor<R, W, MatchGeneratorDriver> {
1681    /// Create a new `FrameCompressor`
1682    pub fn new(compression_level: CompressionLevel) -> Self {
1683        Self {
1684            uncompressed_data: None,
1685            compressed_data: None,
1686            compression_level,
1687            dictionary: None,
1688            source_size_hint: None,
1689            state: CompressState {
1690                matcher: MatchGeneratorDriver::new(1024 * 128, 1),
1691                copy_tier: crate::decoding::simd_copy::ExactCopyTier::resolve(),
1692                last_huff_table: None,
1693                huff_table_spare: None,
1694                huff_rollback: None,
1695                huff_weights: Default::default(),
1696                seen_content: Default::default(),
1697                fse_tables: FseTables::new(),
1698                block_scratch: crate::encoding::blocks::CompressedBlockScratch::new(),
1699                offset_hist: [1, 4, 8],
1700                strategy_tag: crate::encoding::strategy::StrategyTag::for_compression_level(
1701                    compression_level,
1702                ),
1703                pre_split: crate::encoding::levels::config::level_pre_split(compression_level)
1704                    .map(|tier| tier as u8),
1705                huf_optimal_search: true,
1706                literal_compression_disabled: matches!(
1707                    compression_level,
1708                    crate::encoding::CompressionLevel::Level(n) if n < 0
1709                ),
1710            },
1711            magicless: false,
1712            content_checksum: false,
1713            pre_split_disabled: false,
1714            content_size_flag: true,
1715            dict_id_flag: true,
1716            target_block_size: None,
1717            #[cfg(feature = "hash")]
1718            hasher: XxHash64::with_seed(0),
1719            #[cfg(feature = "lsm")]
1720            frame_emit_info: None,
1721            #[cfg(all(feature = "lsm", feature = "hash"))]
1722            per_block_checksums_enabled: false,
1723            #[cfg(all(feature = "lsm", feature = "hash"))]
1724            block_checksums: None,
1725            #[cfg(feature = "lsm")]
1726            block_decompressed_sizes: alloc::vec::Vec::new(),
1727            tuning: FrameTuning::default(),
1728        }
1729    }
1730
1731    /// Configure fine-grained compression parameters (#27).
1732    ///
1733    /// Resets the base [`CompressionLevel`](crate::encoding::CompressionLevel)
1734    /// to the parameters' level and installs the per-knob overrides
1735    /// (window/hash/chain/search logs, strategy, LDM) applied at the next
1736    /// frame. Pass `None`-equivalent (a builder that overrides nothing)
1737    /// to fall back to plain level-based compression.
1738    ///
1739    /// ```rust
1740    /// use structured_zstd::encoding::{
1741    ///     CompressionLevel, CompressionParameters, FrameCompressor, Strategy,
1742    /// };
1743    /// let params = CompressionParameters::builder(CompressionLevel::Level(19))
1744    ///     .strategy(Strategy::Btultra2)
1745    ///     .enable_long_distance_matching(true)
1746    ///     .build()
1747    ///     .unwrap();
1748    /// let mut compressor: FrameCompressor = FrameCompressor::new(CompressionLevel::Default);
1749    /// compressor.set_parameters(&params);
1750    /// let compressed = compressor.compress_independent_frame(b"some data to compress");
1751    /// assert!(!compressed.is_empty());
1752    /// ```
1753    pub fn set_parameters(&mut self, params: &crate::encoding::CompressionParameters) {
1754        self.compression_level = params.level();
1755        let overrides = params.overrides();
1756        self.tuning = FrameTuning::from_overrides(&overrides);
1757        // Keep `state.strategy_tag` consistent immediately so the borrowed
1758        // one-shot eligibility gate (`borrowed_eligible`) and literal gates
1759        // are correct even before the next `compress()` re-sync. Resolve it
1760        // size-adaptively (same `resolve_level_params` path `prepare_frame`
1761        // uses) so a hint already set here yields the same strategy the matcher
1762        // will run, not the bare level-only mapping.
1763        // The dictionary counts only when the frame will prime it (same gate
1764        // as `prepare_frame`'s `use_dictionary_state`): uncompressed mode
1765        // ignores an attached dictionary and has no CDict tier to resolve.
1766        let with_dictionary = !matches!(self.compression_level, CompressionLevel::Uncompressed)
1767            && self.state.matcher.supports_dictionary_priming();
1768        let params = self.resolve_frame_params(self.source_size_hint, with_dictionary);
1769        self.sync_effective_strategy(&params);
1770        self.state.huf_optimal_search =
1771            huf_search_enabled(self.state.strategy_tag, self.source_size_hint);
1772        self.state.literal_compression_disabled = literal_compression_disabled(
1773            self.state.strategy_tag,
1774            self.compression_level,
1775            gate_target_length(
1776                self.compression_level,
1777                &self.tuning,
1778                self.dictionary.as_ref().filter(|_| with_dictionary),
1779            ),
1780            self.tuning.literal_compression,
1781        );
1782        self.state.matcher.set_param_overrides(Some(overrides));
1783    }
1784
1785    /// Whether the borrowed (no per-block history copy) one-shot loop is
1786    /// valid for an `input_len`-byte slice under the resolved `prep`.
1787    ///
1788    /// `Uncompressed` resolves to `StrategyTag::Fast` but must emit stored
1789    /// Raw blocks, which the borrowed loop's
1790    /// `compress_block_encoded_borrowed` (RLE/raw-fast/compressed) does NOT
1791    /// do, so exclude it; it then takes the owned path's dedicated
1792    /// Uncompressed arm.
1793    ///
1794    /// No window-size gate: over-window inputs are handled too. The owned
1795    /// path bounds matches to the last `advertised_window` bytes via
1796    /// `window_low` and evicts its history, sliding the table's stored
1797    /// positions down with it; the borrowed path
1798    /// computes the identical `window_low = block_end - advertised_window`
1799    /// and the kernel rejects any hash candidate below it, while the
1800    /// per-position `put` during the scan keeps in-window slots current,
1801    /// so it produces byte-identical output to the owned (evicting) path
1802    /// without ever copying the input into `history`, even when the input
1803    /// far exceeds the window.
1804    ///
1805    /// BUT gate on `input_len <= u32::MAX`: the Fast kernel stores ABSOLUTE
1806    /// positions in a `u32` hash table, and the borrowed scan walks
1807    /// absolute input offsets up to `block_end == input.len()`. Past 4 GiB
1808    /// those offsets truncate / overflow the `u32` position math
1809    /// (`base_off + ip0 as u32`, `window_low`), panicking or corrupting.
1810    /// The owned/evicting path keeps the scanned window bounded (positions
1811    /// stay small), so >4 GiB inputs fall back to it.
1812    fn borrowed_eligible(&self, input_len: usize, prep: &FramePrep) -> bool {
1813        if matches!(self.compression_level, CompressionLevel::Uncompressed)
1814            || input_len > u32::MAX as usize
1815        {
1816            return false;
1817        }
1818        if prep.use_dictionary_state {
1819            // The borrowed dict scan runs in VIRTUAL `[dict][input]` coordinates,
1820            // so the position space is `dict_content.len() + input_len`, not just
1821            // `input_len`. A large attached dictionary plus an otherwise-allowed
1822            // input can exceed the `u32` floor the kernel asserts — fall back to
1823            // the owned (copy) path in that case.
1824            let fits_u32 = self
1825                .dictionary
1826                .as_ref()
1827                .and_then(|dict| dict.inner.dict_content.len().checked_add(input_len))
1828                .is_some_and(|virtual_len| virtual_len <= u32::MAX as usize);
1829            if !fits_u32 {
1830                return false;
1831            }
1832            // Dictionary frames: only the Simple (Fast) backend in attach mode
1833            // has a borrowed (no input copy) dict scan. Copy-mode dict frames
1834            // and the other backends still take the owned path.
1835            return self.state.matcher.borrowed_dict_supported();
1836        }
1837        // The borrowed (no-copy, in-place over-window) scan exists for the
1838        // Simple (Fast), Dfast, and Row backends, and for the HashChain
1839        // backend's lazy CHAIN parser; BT/optimal (BinaryTree search) stay on
1840        // the owned path. Every borrowed scan applies the per-position
1841        // `window_low = abs_ip - advertised_window` offset cap so over-window
1842        // inputs are matched in place (no input->history copy), matching C's
1843        // continuous-index + windowLow one-shot behaviour.
1844        self.state.matcher.borrowed_supported()
1845    }
1846
1847    /// Compress `input` as one frame's worth of blocks into `out` (appended
1848    /// from its current end): the borrowed in-place loop when
1849    /// [`Self::borrowed_eligible`], else the owned (history-copying) loop fed
1850    /// an in-place `&[u8]` cursor. Returns `total_uncompressed`; the caller
1851    /// emits the frame header (before this call, when the content size is
1852    /// known) or the drain tail.
1853    fn run_one_frame(&mut self, input: &[u8], prep: &FramePrep, out: &mut Vec<u8>) -> u64 {
1854        if self.borrowed_eligible(input.len(), prep) {
1855            // A reused Dfast matcher's tables hold earlier frames, which its
1856            // borrowed kernel (numbering this input from zero) would read as
1857            // its own: emptied, the frame is the one a fresh matcher writes
1858            // (upstream: first and later uses of a context compress the
1859            // same). Measured cheaper than the owned path, which retires them
1860            // by moving the floor but copies the input and runs slower: level 3
1861            // on a reused ZSTD_compress2 context, the owned path costs 3.5% at
1862            // 1 KiB frames, 4.1% at 4 KiB and 10.8% at 1 MiB, and gains at most
1863            // 1% below 256 bytes, within noise. The emptied tables are sized
1864            // from the frame's own source, so a tiny frame clears a tiny table.
1865            if !self.state.matcher.borrowed_frame_is_independent() {
1866                self.state.matcher.forget_earlier_frames();
1867            }
1868            self.run_borrowed_block_loop(input, out)
1869        } else {
1870            let mut cursor: &[u8] = input;
1871            self.run_owned_block_loop(&mut cursor, prep.initial_size_hint, true, out)
1872        }
1873    }
1874
1875    /// Compress one contiguous `&[u8]` as a single independent Zstd frame,
1876    /// writing the frame bytes into `out` (its previous contents are
1877    /// replaced and its allocation reused), reusing this compressor's heavy
1878    /// state across calls.
1879    ///
1880    /// This is the reusable-compression-context (CCtx-equivalent) entry
1881    /// point, mirroring C `ZSTD_compress2` over a reused `ZSTD_CCtx`:
1882    /// construct ONE `FrameCompressor` and call this in a loop to emit N
1883    /// independent, self-describing frames (each carrying its own header,
1884    /// blocks, and checksum, decodable in isolation, with no cross-frame
1885    /// match history). Every call resets the per-frame state via
1886    /// [`Self::prepare_frame`]: only the allocations are kept, so the
1887    /// dominant per-frame setup cost (table allocation + dictionary prime)
1888    /// is paid once instead of N times. Passing the same `out` buffer each
1889    /// call additionally reuses the output allocation, matching C's
1890    /// caller-owned `dst` buffer (no per-frame output allocation).
1891    ///
1892    /// Reusing the context + `out` across many small frames (the typical
1893    /// per-block-frame workload) is far cheaper than a fresh
1894    /// [`compress_slice_to_vec`](crate::encoding::compress_slice_to_vec)
1895    /// per block, which allocates and primes from scratch each time.
1896    ///
1897    /// The input is read in place: no [`Self::set_source`] /
1898    /// [`Self::set_drain`] setup is required, and the input lifetime is not
1899    /// baked into the compressor type, so successive calls may pass slices
1900    /// with unrelated lifetimes. When the Fast (Simple) backend is active
1901    /// and no dictionary is set, the matcher references the input directly
1902    /// (no per-block history copy); other backends / dictionary use copy
1903    /// each block into history exactly as the streaming
1904    /// [`compress`](Self::compress) path does. The source-size hint is
1905    /// derived from the input length on every call, so per-frame table
1906    /// sizing tracks each frame's actual size regardless of any earlier
1907    /// hint.
1908    ///
1909    /// A sticky dictionary set via
1910    /// [`set_dictionary`](Self::set_dictionary) (or its variants) is primed
1911    /// into every frame, mirroring `ZSTD_CCtx_loadDictionary` /
1912    /// `ZSTD_CCtx_refCDict`.
1913    ///
1914    /// # Panics
1915    ///
1916    /// Panics on encoder error, matching [`Self::compress`] and
1917    /// [`compress_slice_to_vec`](crate::encoding::compress_slice_to_vec).
1918    pub fn compress_independent_frame_into(&mut self, input: &[u8], out: &mut Vec<u8>) {
1919        // Size the next frame from the actual payload, not a stale hint a
1920        // previous call may have left behind (a wrong hint would change the
1921        // resolved window/header and could flip borrowed eligibility).
1922        self.source_size_hint = Some(input.len() as u64);
1923        let prep = self.prepare_frame();
1924        // Content size is known up front (one-shot), so write the frame
1925        // header FIRST and emit blocks STRAIGHT into `out` — no separate
1926        // `all_blocks` accumulator and no header+blocks copy (which was the
1927        // dominant per-frame memmove + the only un-amortized per-frame alloc
1928        // even when the compressor is reused).
1929        let total_uncompressed = input.len() as u64;
1930        let emit_checksum = cfg!(feature = "hash") && self.content_checksum;
1931        let checksum_len = if emit_checksum { 4 } else { 0 };
1932        out.clear();
1933        // Reserve the header plus ONE block's worst case up front; the block
1934        // loops then grow `out` from the compression ratio observed so far
1935        // (`reserve_for_next_block`). Reserving `compress_bound(input_len)`
1936        // here held a whole-input-sized allocation for the entire frame —
1937        // ~100 MiB peak on a 100 MiB stream whose compressed output is a few
1938        // MiB, where the reference implementation's context peaks at
1939        // window-sized state. Small frames (<= one block) still get their
1940        // full bound in one shot, so the reused-`out` steady state is
1941        // unchanged. 18 = max frame header (magic 4 + descriptor 1 + window
1942        // 1 + dict id 4 + FCS 8).
1943        let first_block_bound = input.len().min(self.block_capacity()) + 3;
1944        out.reserve(18 + first_block_bound + checksum_len);
1945        self.append_frame_header(total_uncompressed, &prep, out);
1946        let header_len = out.len();
1947        let _ = self.run_one_frame(input, &prep, out);
1948        #[cfg(feature = "hash")]
1949        if self.content_checksum {
1950            out.extend_from_slice(&(self.hasher.finish() as u32).to_le_bytes());
1951        }
1952        #[cfg(feature = "lsm")]
1953        {
1954            let blocks_end = out.len() - checksum_len;
1955            self.populate_frame_emit_info(header_len, &out[header_len..blocks_end], emit_checksum);
1956        }
1957        #[cfg(not(feature = "lsm"))]
1958        let _ = header_len;
1959    }
1960
1961    /// Convenience wrapper over [`Self::compress_independent_frame_into`]
1962    /// that allocates and returns a fresh `Vec` per call. Prefer the
1963    /// `_into` form in tight per-block-frame loops to reuse one output
1964    /// buffer across frames (the CCtx-equivalent zero-per-call-alloc
1965    /// output, matching C's caller-owned `dst`).
1966    ///
1967    /// ```rust
1968    /// use structured_zstd::encoding::{FrameCompressor, CompressionLevel};
1969    /// let mut cctx: FrameCompressor = FrameCompressor::new(CompressionLevel::Default);
1970    /// let frame_a = cctx.compress_independent_frame(b"first block payload");
1971    /// let frame_b = cctx.compress_independent_frame(b"second block payload");
1972    /// assert!(!frame_a.is_empty() && !frame_b.is_empty());
1973    /// ```
1974    pub fn compress_independent_frame(&mut self, input: &[u8]) -> Vec<u8> {
1975        let mut out = Vec::new();
1976        self.compress_independent_frame_into(input, &mut out);
1977        out
1978    }
1979
1980    /// Borrowed one-shot block loop: walks `input` in `MAX_BLOCK_SIZE`
1981    /// strides (the Fast backend never pre-splits, so boundaries match the
1982    /// owned loop), scanning each block range in place against the
1983    /// borrowed window via `compress_block_encoded_borrowed` — no
1984    /// per-block `commit_space` copy. Returns `(all_blocks,
1985    /// total_uncompressed)`. Caller guarantees Fast backend + no
1986    /// dictionary; over-window inputs are fine (matches are bounded by
1987    /// `window_low` exactly as the owned evicting path).
1988    fn run_borrowed_block_loop(&mut self, input: &[u8], out: &mut Vec<u8>) -> u64 {
1989        // Blocks are appended to `out` starting here. `out` may already hold
1990        // the frame header (the one-shot compress-into-Vec path writes it
1991        // first, since the content size is known up front, and the loop
1992        // emits blocks straight after it — no separate `all_blocks` Vec and
1993        // no header+blocks copy). Output-size reads below are taken RELATIVE
1994        // to `blocks_start` so a header prefix never skews the upstream zstd split
1995        // `savings` gate (which would change block boundaries / wire output).
1996        let blocks_start = out.len();
1997        let total_uncompressed = input.len() as u64;
1998        // Empty input: emit a single empty last Raw block (mirrors the
1999        // owned loop's empty-file special case).
2000        if input.is_empty() {
2001            let header = BlockHeader {
2002                last_block: true,
2003                block_type: crate::blocks::block::BlockType::Raw,
2004                block_size: 0,
2005            };
2006            header.serialize(out);
2007            #[cfg(feature = "lsm")]
2008            self.block_decompressed_sizes.push(0);
2009            #[cfg(all(feature = "lsm", feature = "hash"))]
2010            if let Some(checksums) = self.block_checksums.as_mut() {
2011                checksums.push(xxh64_block_low32(&[]));
2012            }
2013            return total_uncompressed;
2014        }
2015        // SAFETY: `input` outlives this call (held by the caller across
2016        // the call) and is not mutated. Only the Simple backend is active
2017        // (gated by `compress_oneshot_borrowed`).
2018        unsafe {
2019            self.state.matcher.set_borrowed_window(input);
2020        }
2021        // Panic-safety: clear the borrowed `(ptr, len)` on EVERY exit,
2022        // including an unwind from an `assert!` inside the block loop, so
2023        // a caught-and-reused compressor never retains a dangling window.
2024        // (The next frame's `reset()` also clears it before any read, but
2025        // this guard makes the invariant local and unwind-proof.)
2026        struct ClearBorrowedOnDrop(*mut MatchGeneratorDriver);
2027        impl Drop for ClearBorrowedOnDrop {
2028            fn drop(&mut self) {
2029                // SAFETY: at drop (normal return or unwind) the loop's
2030                // borrows of the matcher have ended, so this is the only
2031                // access. `addr_of_mut!` produced this pointer without an
2032                // intermediate `&mut`, so the interleaved `&mut` uses in
2033                // the loop did not invalidate it.
2034                unsafe { (*self.0).clear_borrowed_window() };
2035            }
2036        }
2037        let _clear_guard = ClearBorrowedOnDrop(core::ptr::addr_of_mut!(self.state.matcher));
2038        let block_capacity = self.block_capacity();
2039        let mut start = 0usize;
2040        while start < input.len() {
2041            reserve_for_next_block(
2042                out,
2043                blocks_start,
2044                start as u64,
2045                input.len() - start,
2046                block_capacity,
2047            );
2048            // Upstream zstd `ZSTD_compress_frameChunk`: size each block via the cheap
2049            // fingerprint pre-splitter so a full 128 KiB block is cut at a
2050            // statistical boundary when it pays. `savings = consumed -
2051            // produced` mirrors the upstream zstd gate (the first block and
2052            // incompressible input keep the full 128 KiB). The borrowed window
2053            // already spans the whole input, so a smaller block is just a
2054            // narrower `(block_start, block_end)` range into it.
2055            let savings = start as i64 - (out.len() - blocks_start) as i64;
2056            // Borrowed path only: warm the pre-split window before the
2057            // cache-cold strided fingerprint read. Gated to exactly the
2058            // conditions under which `optimal_block_size` reads `block`
2059            // (a pre-split level, a full 128 KiB block remaining, the
2060            // block-size cap admits a full block, and `savings >= 3` so the
2061            // splitter actually runs) — so non-pre-split levels, the first
2062            // block, and the trailing partial block pay nothing. See
2063            // `warm_presplit_window`.
2064            let pre_split = self.pre_split_level();
2065            if savings >= 3
2066                && input.len() - start >= MAX_BLOCK_SIZE as usize
2067                && block_capacity >= MAX_BLOCK_SIZE as usize
2068                && pre_split.is_some()
2069            {
2070                warm_presplit_window(&input[start..start + MAX_BLOCK_SIZE as usize]);
2071            }
2072            let block_len = optimal_block_size_with(
2073                pre_split,
2074                &input[start..],
2075                input.len() - start,
2076                block_capacity,
2077                savings,
2078            );
2079            let end = (start + block_len).min(input.len());
2080            let block = &input[start..end];
2081            let last_block = end == input.len();
2082            #[cfg(feature = "hash")]
2083            if self.content_checksum {
2084                self.hasher.write(block);
2085            }
2086            let dict_active =
2087                self.dictionary.is_some() && self.state.matcher.supports_dictionary_priming();
2088            crate::encoding::levels::compress_block_encoded_borrowed(
2089                &mut self.state,
2090                self.compression_level,
2091                last_block,
2092                block,
2093                start,
2094                end,
2095                out,
2096                dict_active,
2097                #[cfg(feature = "lsm")]
2098                Some(&mut self.block_decompressed_sizes),
2099                #[cfg(all(feature = "lsm", feature = "hash"))]
2100                self.block_checksums.as_mut(),
2101            );
2102            start = end;
2103        }
2104        // `_clear_guard` drops here, clearing the borrowed window.
2105        total_uncompressed
2106    }
2107}
2108
2109impl<R: Read, W: Write, M: Matcher> FrameCompressor<R, W, M> {
2110    /// Create a new `FrameCompressor` with a custom matching algorithm implementation
2111    pub fn new_with_matcher(matcher: M, compression_level: CompressionLevel) -> Self {
2112        Self {
2113            uncompressed_data: None,
2114            compressed_data: None,
2115            dictionary: None,
2116            source_size_hint: None,
2117            state: CompressState {
2118                matcher,
2119                copy_tier: crate::decoding::simd_copy::ExactCopyTier::resolve(),
2120                last_huff_table: None,
2121                huff_table_spare: None,
2122                huff_rollback: None,
2123                huff_weights: Default::default(),
2124                seen_content: Default::default(),
2125                fse_tables: FseTables::new(),
2126                block_scratch: crate::encoding::blocks::CompressedBlockScratch::new(),
2127                offset_hist: [1, 4, 8],
2128                strategy_tag: crate::encoding::strategy::StrategyTag::for_compression_level(
2129                    compression_level,
2130                ),
2131                pre_split: crate::encoding::levels::config::level_pre_split(compression_level)
2132                    .map(|tier| tier as u8),
2133                huf_optimal_search: true,
2134                literal_compression_disabled: matches!(
2135                    compression_level,
2136                    crate::encoding::CompressionLevel::Level(n) if n < 0
2137                ),
2138            },
2139            compression_level,
2140            magicless: false,
2141            content_checksum: false,
2142            pre_split_disabled: false,
2143            content_size_flag: true,
2144            dict_id_flag: true,
2145            target_block_size: None,
2146            #[cfg(feature = "hash")]
2147            hasher: XxHash64::with_seed(0),
2148            #[cfg(feature = "lsm")]
2149            frame_emit_info: None,
2150            #[cfg(all(feature = "lsm", feature = "hash"))]
2151            per_block_checksums_enabled: false,
2152            #[cfg(all(feature = "lsm", feature = "hash"))]
2153            block_checksums: None,
2154            #[cfg(feature = "lsm")]
2155            block_decompressed_sizes: alloc::vec::Vec::new(),
2156            tuning: FrameTuning::default(),
2157        }
2158    }
2159
2160    /// Enable or disable magicless frame format (`ZSTD_f_zstd1_magicless`).
2161    ///
2162    /// When set to `true`, emitted frames omit the 4-byte magic number
2163    /// prefix. The matching decoder must be configured to expect a
2164    /// magicless stream — wire-format only round-trips with a
2165    /// magicless-aware decoder.
2166    pub fn set_magicless(&mut self, magicless: bool) {
2167        self.magicless = magicless;
2168    }
2169
2170    /// Enable or disable the trailing XXH64 content checksum
2171    /// (semantics of upstream `ZSTD_c_checksumFlag`). Default `false`,
2172    /// matching the upstream library default (`ZSTD_c_checksumFlag = 0`)
2173    /// so out-of-the-box frames carry the same layout and pay the same
2174    /// costs as the reference implementation.
2175    ///
2176    /// When `false`, emitted frames set `Content_Checksum_flag = 0` and carry
2177    /// no trailing digest; such frames are valid (RFC 8878) and decode
2178    /// correctly in any [`ContentChecksum`](crate::decoding::ContentChecksum)
2179    /// mode. Without the `hash` feature no checksum is emitted regardless of
2180    /// this setting.
2181    pub fn set_content_checksum(&mut self, emit: bool) {
2182        self.content_checksum = emit;
2183    }
2184
2185    /// Enable or disable recording `Frame_Content_Size` in the frame header
2186    /// when the total size is known (semantics of upstream
2187    /// `ZSTD_c_contentSizeFlag`). Default `true`, matching upstream. With
2188    /// the flag off the header carries a window descriptor instead (and the
2189    /// single-segment layout, which requires an FCS, is disabled).
2190    pub fn set_content_size_flag(&mut self, emit: bool) {
2191        self.content_size_flag = emit;
2192    }
2193
2194    /// Enable or disable recording the dictionary ID in the frame header
2195    /// when a dictionary is attached (semantics of upstream
2196    /// `ZSTD_c_dictIDFlag`). Default `true`, matching upstream. Frames
2197    /// emitted with the flag off still decode when the decoder is handed
2198    /// the dictionary explicitly.
2199    pub fn set_dictionary_id_flag(&mut self, emit: bool) {
2200        self.dict_id_flag = emit;
2201    }
2202
2203    /// Set an upper bound on emitted block sizes (semantics of upstream
2204    /// `ZSTD_c_targetCBlockSize`): every physical block's payload is capped
2205    /// at `target` bytes (+3-byte block header on the wire), trading some
2206    /// ratio for bounded per-block latency. The value is clamped to
2207    /// `[MIN_TARGET_BLOCK_SIZE, MAX_BLOCK_SIZE]` (the upstream bounds).
2208    /// `None` removes the target.
2209    pub fn set_target_block_size(&mut self, target: Option<u32>) {
2210        self.target_block_size = target.map(|t| {
2211            t.clamp(
2212                crate::common::MIN_TARGET_BLOCK_SIZE,
2213                crate::common::MAX_BLOCK_SIZE,
2214            )
2215        });
2216    }
2217
2218    /// The active block-size cap: the configured target, or the format's
2219    /// 128 KiB block ceiling.
2220    fn block_capacity(&self) -> usize {
2221        let requested = self
2222            .target_block_size
2223            .map_or(crate::common::MAX_BLOCK_SIZE as usize, |t| t as usize);
2224        // Upstream zstd sizes a block as `MIN(maxBlockSize, windowSize)`
2225        // (`ZSTD_compress.c`). A block wider than the window can never be
2226        // held by the matcher, which asserts on it, so a small `window_log`
2227        // must shrink the block rather than overrun the window.
2228        let window = self.state.matcher.window_size() as usize;
2229        if window == 0 {
2230            requested
2231        } else {
2232            requested.min(window)
2233        }
2234    }
2235
2236    /// Before calling [FrameCompressor::compress] you need to set the source.
2237    ///
2238    /// This is the data that is compressed and written into the drain.
2239    pub fn set_source(&mut self, uncompressed_data: R) -> Option<R> {
2240        self.uncompressed_data.replace(uncompressed_data)
2241    }
2242
2243    /// Before calling [FrameCompressor::compress] you need to set the drain.
2244    ///
2245    /// As the compressor compresses data, the drain serves as a place for the output to be writte.
2246    pub fn set_drain(&mut self, compressed_data: W) -> Option<W> {
2247        self.compressed_data.replace(compressed_data)
2248    }
2249
2250    /// Diagnostic switch: cut full blocks only, never pre-split. Used by
2251    /// the sequence capture so its block structure matches upstream's
2252    /// `ZSTD_generateSequences` (see `pre_split_disabled`).
2253    #[cfg(feature = "bench-internals")]
2254    pub fn set_pre_split_disabled(&mut self, disabled: bool) {
2255        self.pre_split_disabled = disabled;
2256    }
2257
2258    /// The pre-split level the block loops apply: the effective strategy's
2259    /// tier (`state.pre_split`) unless the diagnostic switch is on.
2260    fn pre_split_level(&self) -> Option<usize> {
2261        if self.pre_split_disabled {
2262            None
2263        } else {
2264            self.state.pre_split.map(usize::from)
2265        }
2266    }
2267
2268    /// The level params the matcher's reset resolves for the next frame
2269    /// (dictionary-aware when a dictionary will be used).
2270    fn resolve_frame_params(
2271        &self,
2272        hint: Option<u64>,
2273        with_dictionary: bool,
2274    ) -> crate::encoding::levels::config::LevelParams {
2275        resolve_frame_params(
2276            self.compression_level,
2277            hint,
2278            self.dictionary.as_ref().filter(|_| with_dictionary),
2279        )
2280    }
2281
2282    /// Record the strategy the matcher actually runs for the next frame (a
2283    /// public-parameter override, else the size- and dictionary-adaptive
2284    /// resolution in `params`) and its pre-split tier: the literal gates and
2285    /// the block splitter read these, and upstream indexes `splitLevels` by
2286    /// the effective strategy too.
2287    fn sync_effective_strategy(&mut self, params: &crate::encoding::levels::config::LevelParams) {
2288        sync_effective_strategy(
2289            &mut self.state,
2290            self.compression_level,
2291            params,
2292            self.tuning.strategy,
2293        );
2294    }
2295
2296    /// Provide a hint about the total uncompressed size for the next frame.
2297    ///
2298    /// When set, the encoder selects smaller hash tables and windows for
2299    /// small inputs, matching the C zstd source-size-class behavior.
2300    ///
2301    /// This hint applies only to frame payload bytes (`size`). Dictionary
2302    /// history is primed separately and does not inflate the hinted size or
2303    /// advertised frame window.
2304    /// Must be called before [`compress`](Self::compress).
2305    pub fn set_source_size_hint(&mut self, size: u64) {
2306        self.source_size_hint = Some(size);
2307    }
2308
2309    /// Total heap bytes this compressor's allocations hold, excluding the
2310    /// inline struct: the match-finder tables / history / recycled buffers and
2311    /// the primed-dictionary snapshot (via the matcher), the retained
2312    /// Huffman tables (active + recycled spare), the retained dictionary
2313    /// content, the cached dictionary entropy tables (literals Huffman +
2314    /// LL/ML/OF FSE), and the per-block sidecar buffers. Lets a context
2315    /// report its true footprint through `ZSTD_sizeof_CCtx`.
2316    pub fn heap_size(&self) -> usize {
2317        let mut total = self.state.matcher.heap_size();
2318        total += self
2319            .state
2320            .last_huff_table
2321            .as_ref()
2322            .map_or(0, |table| table.heap_size());
2323        total += self
2324            .state
2325            .huff_table_spare
2326            .as_ref()
2327            .map_or(0, |table| table.heap_size());
2328        // The weight builder's buffers are kept between blocks and frames, so
2329        // a reused compressor holds them for as long as it lives.
2330        total += self.state.huff_weights.heap_size();
2331        total += self.state.retained_scratch_heap_size();
2332        total += self.state.seen_content.heap_size();
2333        total += self
2334            .dictionary
2335            .as_ref()
2336            .map_or(0, EncoderDictionary::heap_size);
2337        #[cfg(all(feature = "lsm", feature = "hash"))]
2338        {
2339            total += self
2340                .block_checksums
2341                .as_ref()
2342                .map_or(0, |v| v.capacity() * core::mem::size_of::<u32>());
2343        }
2344        #[cfg(feature = "lsm")]
2345        {
2346            total += self.block_decompressed_sizes.capacity() * core::mem::size_of::<u32>();
2347        }
2348        total
2349    }
2350
2351    /// Compress the uncompressed data from the provided source as one Zstd frame and write it to the provided drain
2352    ///
2353    /// This will repeatedly call [Read::read] on the source to fill up blocks until the source returns 0 on the read call.
2354    /// All compressed blocks are buffered in memory so that the frame header can include the
2355    /// `Frame_Content_Size` field (which requires knowing the total uncompressed size). The
2356    /// entire frame — header, blocks, and optional checksum — is then written to the drain
2357    /// at the end. This means peak memory usage is O(compressed_size).
2358    ///
2359    /// To avoid endlessly encoding from a potentially endless source (like a network socket) you can use the
2360    /// [Read::take] function
2361    /// Per-frame setup values resolved by [`Self::prepare_frame`] and
2362    /// consumed by the block loop + [`Self::finish_frame`]. Lets the
2363    /// owned `compress()` and the borrowed one-shot path share the exact
2364    /// same reset / dict-prime / entropy-seed setup and frame tail.
2365    pub fn compress(&mut self) {
2366        let prep = self.prepare_frame();
2367        // Take the reader out so `run_owned_block_loop` can borrow it
2368        // mutably alongside `&mut self` (the rest of the loop touches
2369        // `self.state` / `self.hasher`, disjoint from the reader). Restored
2370        // before the frame tail so a reused compressor keeps its source.
2371        //
2372        // Deliberately NOT restored on unwind: if the block loop panics the
2373        // source has been partially consumed, so handing it back would let a
2374        // `catch_unwind` caller "successfully" compress the remaining tail
2375        // from an arbitrary midpoint — silent data corruption. Leaving the
2376        // slot empty makes any post-panic reuse fail loudly at the `expect`
2377        // below (matcher/entropy state is equally unre-usable after an
2378        // unwind; the reference implementation likewise requires a context
2379        // reset after an error).
2380        let mut source = self
2381            .uncompressed_data
2382            .take()
2383            .expect("source must be set via set_source before compress()");
2384        // Streaming drain: the content size is only known at EOF, so the
2385        // frame header can't precede the blocks — accumulate them in a local
2386        // buffer and let `finish_frame` write header + blocks to the drain.
2387        let mut all_blocks: Vec<u8> = Vec::with_capacity(initial_all_blocks_cap(
2388            prep.initial_size_hint,
2389            self.block_capacity(),
2390        ));
2391        let mut block_source = ReaderBlockSource::new(&mut source);
2392        let total_uncompressed = self.run_owned_block_loop(
2393            &mut block_source,
2394            prep.initial_size_hint,
2395            false,
2396            &mut all_blocks,
2397        );
2398        self.uncompressed_data = Some(source);
2399        self.finish_frame(all_blocks, total_uncompressed, &prep);
2400    }
2401
2402    fn prepare_frame(&mut self) -> FramePrep {
2403        // The raw-skip's memory of what this frame has already emitted. Frames
2404        // are independent, so carrying it over would let one frame's content
2405        // hold the skip off for the next; the allocation is kept.
2406        self.state.seen_content.reset_for_frame();
2407        // Reset per-frame introspection state so a re-used compressor
2408        // doesn't carry over the previous frame's layout/checksums.
2409        #[cfg(feature = "lsm")]
2410        {
2411            self.frame_emit_info = None;
2412            // Always captured under lsm (drives `decompressed_byte_range`);
2413            // clear, keep the allocation for a reused compressor.
2414            self.block_decompressed_sizes.clear();
2415        }
2416        #[cfg(all(feature = "lsm", feature = "hash"))]
2417        {
2418            if self.per_block_checksums_enabled {
2419                self.block_checksums = Some(alloc::vec::Vec::new());
2420            } else {
2421                self.block_checksums = None;
2422            }
2423        }
2424        let initial_size_hint = self.source_size_hint;
2425        let source_size_hint_known = initial_size_hint.is_some();
2426        let use_dictionary_state =
2427            !matches!(self.compression_level, CompressionLevel::Uncompressed)
2428                && self.state.matcher.supports_dictionary_priming()
2429                && self.dictionary.is_some();
2430        if let Some(size_hint) = self.source_size_hint.take() {
2431            // Keep source-size hint scoped to payload bytes; dictionary priming
2432            // is applied separately and should not force larger matcher sizing.
2433            self.state.matcher.set_source_size_hint(size_hint);
2434        }
2435        // Hand the matcher the dictionary's sizes so the frame runs the CDict's
2436        // cParams tier and sizes its dictionary tables from the content. Set
2437        // before `reset` (which consumes it) and only when a dictionary will
2438        // actually be primed.
2439        if use_dictionary_state && let Some(dict) = self.dictionary.as_ref() {
2440            self.state.matcher.set_dictionary_size_hint(dict.sizes());
2441        }
2442        // Clearing buffers to allow re-using of the compressor
2443        self.state.matcher.reset(self.compression_level);
2444        self.state.offset_hist = [1, 4, 8];
2445        // Sync `state.strategy_tag` to the level resolved at this reset so
2446        // the literal-compression gates (`min_literals_to_compress` /
2447        // `min_gain` in `encoding::blocks::compressed`) see the correct
2448        // strategy for the next frame. Frame-by-frame level changes go
2449        // through this same `compress()` entry point, so re-syncing here
2450        // covers level switches without touching the matcher dispatch.
2451        // A public-parameter strategy override (#27) wins over the level's
2452        // derived tag so the literal-compression gates and dict-attach cutoff
2453        // below see the strategy the matcher actually runs. Otherwise resolve
2454        // the strategy SIZE-ADAPTIVELY through the same path the matcher's reset
2455        // used (`resolve_level_params` -> `get_cparams`, the port of upstream
2456        // `ZSTD_getCParams`): a small frame promotes a level to a higher
2457        // strategy (e.g. L13 over a <=16 KiB frame becomes btultra). Re-deriving
2458        // from the bare level would make the literal-compression / HUF-search
2459        // gates disagree with the matcher's actual parse on small frames (the
2460        // gate would think btlazy2 and skip the HUF table-log search the btultra
2461        // frame runs, costing a few bytes on small literal sections).
2462        // A dictionary frame runs the CDict's strategy (upstream
2463        // `ZSTD_resetCCtx_usingCDict`), so resolve through the same
2464        // dictionary-aware path the matcher's reset took.
2465        let params = self.resolve_frame_params(initial_size_hint, use_dictionary_state);
2466        self.sync_effective_strategy(&params);
2467        // `initial_size_hint` (captured before the `.take()` above) — by here
2468        // `self.source_size_hint` is None.
2469        self.state.huf_optimal_search =
2470            huf_search_enabled(self.state.strategy_tag, initial_size_hint);
2471        // Recomputed per frame from the persisted override: attaching or
2472        // clearing a dictionary after `set_parameters` can move the strategy
2473        // this gate reads.
2474        self.state.literal_compression_disabled = literal_compression_disabled(
2475            self.state.strategy_tag,
2476            self.compression_level,
2477            gate_target_length(
2478                self.compression_level,
2479                &self.tuning,
2480                self.dictionary.as_ref().filter(|_| use_dictionary_state),
2481            ),
2482            self.tuning.literal_compression,
2483        );
2484        let cached_entropy = if use_dictionary_state {
2485            self.dictionary.as_ref().map(|dict| &dict.inner.entropy)
2486        } else {
2487            None
2488        };
2489        if use_dictionary_state && let Some(dict) = self.dictionary.as_ref() {
2490            // This state drives sequence encoding, while matcher priming below updates
2491            // the match generator's internal repeat-offset history for match finding.
2492            self.state.offset_hist = dict.inner.offset_hist;
2493            load_frame_dictionary(
2494                &mut self.state,
2495                self.compression_level,
2496                dict,
2497                initial_size_hint,
2498            );
2499        }
2500        if let Some(cache) = cached_entropy {
2501            // Refill an empty slot from the recycled spare before
2502            // `clone_from`: `Option::clone_from(None ← Some)` falls back to
2503            // a fresh clone (two Vec allocations), while `Some ← Some`
2504            // delegates to the table's buffer-reusing `clone_from`. Frames
2505            // whose last block cleared the table would otherwise re-clone
2506            // the dict seed every frame.
2507            match &cache.huff {
2508                Some(src) => {
2509                    if self.state.last_huff_table.is_none() {
2510                        self.state.last_huff_table = self.state.huff_table_spare.take();
2511                    }
2512                    match &mut self.state.last_huff_table {
2513                        Some(dst) => dst.clone_from(src),
2514                        slot => *slot = Some(src.clone()),
2515                    }
2516                }
2517                None => self.state.clear_huff_table(),
2518            }
2519        } else {
2520            self.state.clear_huff_table();
2521        }
2522        // Whatever the last frame ended on is about to be replaced. Keep it as
2523        // this frame's build buffer rather than dropping it, or a reused
2524        // compressor that emits one custom-table block per frame allocates a
2525        // table per axis every frame.
2526        self.state.fse_tables.park_previous_before_frame();
2527        // `clone_from` keeps frame-to-frame seeding cheap for reused compressors by
2528        // reusing existing allocations where possible instead of reallocating every frame.
2529        if let Some(cache) = cached_entropy {
2530            self.state
2531                .fse_tables
2532                .ll_previous
2533                .clone_from(&cache.ll_previous);
2534            self.state
2535                .fse_tables
2536                .ml_previous
2537                .clone_from(&cache.ml_previous);
2538            self.state
2539                .fse_tables
2540                .of_previous
2541                .clone_from(&cache.of_previous);
2542        } else {
2543            self.state.fse_tables.ll_previous = None;
2544            self.state.fse_tables.ml_previous = None;
2545            self.state.fse_tables.of_previous = None;
2546        }
2547        let ll_entropy = cached_entropy.and_then(|cache| match cache.ll_previous.as_ref() {
2548            Some(PreviousFseTable::Custom(table)) => Some(table.as_ref()),
2549            _ => None,
2550        });
2551        let ml_entropy = cached_entropy.and_then(|cache| match cache.ml_previous.as_ref() {
2552            Some(PreviousFseTable::Custom(table)) => Some(table.as_ref()),
2553            _ => None,
2554        });
2555        let of_entropy = cached_entropy.and_then(|cache| match cache.of_previous.as_ref() {
2556            Some(PreviousFseTable::Custom(table)) => Some(table.as_ref()),
2557            _ => None,
2558        });
2559        self.state.matcher.seed_dictionary_entropy(
2560            self.state.last_huff_table.as_ref(),
2561            ll_entropy,
2562            ml_entropy,
2563            of_entropy,
2564        );
2565        #[cfg(feature = "hash")]
2566        {
2567            self.hasher = XxHash64::with_seed(0);
2568        }
2569        let window_size = self.state.matcher.window_size();
2570        assert!(
2571            window_size != 0,
2572            "matcher reported window_size == 0, which is invalid"
2573        );
2574        FramePrep {
2575            window_size,
2576            use_dictionary_state,
2577            source_size_hint_known,
2578            initial_size_hint,
2579        }
2580    }
2581
2582    /// Owned streaming block loop: reads blocks from the caller-provided
2583    /// `source` reader, optionally pre-splits, hashes for the content
2584    /// checksum, and emits each block via `compress_block_encoded`,
2585    /// accumulating the block bytes. Returns `(all_blocks,
2586    /// total_uncompressed)`. The source is passed in (rather than read
2587    /// from `self.uncompressed_data`) so the streaming `compress` path can
2588    /// feed the configured reader while the slice paths
2589    /// (`compress_oneshot_borrowed`, `compress_independent_frame`) feed an
2590    /// in-place `&[u8]` cursor without baking its lifetime into the
2591    /// compressor type.
2592    fn run_owned_block_loop<S: OwnedBlockSource>(
2593        &mut self,
2594        source: &mut S,
2595        initial_size_hint: Option<u64>,
2596        // Whether `initial_size_hint` is the input's exact length (the
2597        // one-shot slice paths) or a caller-provided estimate (the streaming
2598        // `Read` path, where `set_source_size_hint` is advisory). An exact
2599        // hint drives the one-shot ratio reservation; an estimate is only
2600        // trusted up to a small lookahead past the bytes actually read.
2601        hint_is_exact: bool,
2602        out: &mut Vec<u8>,
2603    ) -> u64 {
2604        // Compressed blocks are appended to `out` from its current end. The
2605        // streaming drain path passes a fresh buffer (the frame header is
2606        // written to the drain afterward, since Frame_Content_Size is only
2607        // known once the reader hits EOF); the one-shot compress-into-Vec
2608        // path passes `out` already holding the header. The upstream zstd split
2609        // `savings` gate below accumulates block-relative (`before_len`)
2610        // output deltas, so a header prefix never skews it.
2611        let blocks_start = out.len();
2612        let mut total_uncompressed: u64 = 0;
2613        let mut pending_input: Vec<u8> = Vec::new();
2614        let mut reached_eof = false;
2615        let mut savings = 0i64;
2616        // One allocation for the whole frame's ingest buffer, instead of a
2617        // doubling chain of reallocations as the blocks arrive. A fresh
2618        // compressor starts with an empty buffer, so without this every frame
2619        // climbs the ladder again and hands the pages back at the end of it:
2620        // measured at level 3 over a 1 MB frame, three growth steps per frame
2621        // and about 2.4 MB of pages faulted back in each time, against none for
2622        // a reference that sizes its workspace once.
2623        //
2624        // An inexact hint is sized on too. The worry it would otherwise raise —
2625        // that a wild overestimate reserves memory the reader never fills — is
2626        // already answered twice over: the same hint has by this point sized
2627        // the window and the match-finder tables (it reaches the matcher
2628        // through `set_source_size_hint`, and the level parameters cap the
2629        // window by it), and `reserve_for_frame` clamps to the eviction ceiling
2630        // the buffer would reach anyway. So the reservation is proportionate to
2631        // allocations the hint has already caused, not a new class of waste.
2632        // The slack is one block, so the final top-up (which asks for a whole
2633        // block even when only a tail remains) does not reallocate; sized off
2634        // the ACTIVE block capacity, since a small window shrinks the block
2635        // below the format maximum.
2636        // Raw frames are excluded: they emit straight from the staged buffer
2637        // and never consult the match finder (the `in_place` gate below keeps
2638        // them off it whatever the backend supports), so sizing its history for
2639        // them holds a window's worth of memory the frame has no use for.
2640        if let Some(hint) = initial_size_hint
2641            && !matches!(self.compression_level, CompressionLevel::Uncompressed)
2642        {
2643            // `saturating_add`: a caller may pledge `u64::MAX`, and clamping a
2644            // reservation request at the address-space limit is the meaningful
2645            // answer — the matcher caps it at its eviction ceiling anyway.
2646            let mut target =
2647                (hint.min(usize::MAX as u64) as usize).saturating_add(self.block_capacity());
2648            if !hint_is_exact {
2649                // An advisory number is a claim about data that has not arrived,
2650                // so it is trusted only as far as the frame's own configuration
2651                // makes plausible: the window this LEVEL would choose, never an
2652                // overridden one. Overriding the window is itself a claim about
2653                // the data — one only the data can confirm — and taking it here
2654                // let a caller who promised gibibytes and delivered ten bytes
2655                // reserve two of them. Beyond this bound the buffer grows as it
2656                // did before, which costs a few reallocations on frames already
2657                // large enough for that to be noise.
2658                let level_window = crate::encoding::levels::config::resolve_level_params(
2659                    self.compression_level,
2660                    initial_size_hint,
2661                )
2662                .window_log;
2663                let plausible = (1usize << level_window).saturating_add(self.block_capacity());
2664                target = target.min(plausible);
2665            }
2666            self.state.matcher.reserve_for_frame(target);
2667        }
2668        // Compress block by block
2669        loop {
2670            // Read up to one upstream zstd block. When the pre-block splitter keeps a
2671            // suffix, top it back up before compressing the next block, matching
2672            // ZSTD_compress_frameChunk() over a contiguous input buffer.
2673            let block_capacity = self.block_capacity();
2674            // Always draw the block buffer from the matcher's recycled pool
2675            // (its capacity already covers the block size, so the resize below
2676            // stays in-place). Any carried pre-split suffix is copied in, and
2677            // `pending_input` is retained as a reusable carry buffer. The prior
2678            // approach `split_off`'d a fresh suffix Vec per pre-split and
2679            // `reserve_exact`-grew it to `block_capacity` every block; on a
2680            // heavily pre-split frame that churned one block-sized allocation
2681            // per split (~12 MB over ~90 splits on a 1 MiB corpus input).
2682            // Remaining-bytes expectation for the reader source's sizing
2683            // (`None` = unknown, or an inexact hint already met by prior
2684            // blocks). The slice source appends directly and ignores it.
2685            let size_hint_remaining = match initial_size_hint {
2686                Some(hint) if hint > total_uncompressed => Some(hint - total_uncompressed),
2687                _ => None,
2688            };
2689            // Preferred shape: read straight into the matcher's history, so
2690            // neither this block nor a pre-split remainder is ever copied. The
2691            // leftover from the previous iteration is already sitting there as
2692            // uncommitted bytes, which is why there is no `pending_input`
2693            // top-up on this path.
2694            // `Uncompressed` emits Raw blocks straight from the staged buffer,
2695            // so it stays on the staged path whatever the matcher supports. The
2696            // gate is on the LEVEL, not the backend: `fill_in_place` dispatches
2697            // on the matcher, and an external `M: Matcher` that implements it
2698            // would otherwise leave the payload sitting uncommitted while an
2699            // empty Raw block goes out.
2700            let in_place = if matches!(self.compression_level, CompressionLevel::Uncompressed) {
2701                None
2702            } else if reached_eof {
2703                // Nothing left to read; the carried remainder is already in the
2704                // matcher, so just re-inspect it.
2705                self.state
2706                    .matcher
2707                    .fill_in_place(0, &mut |_buf| (0, true))
2708                    .map(|_| 0usize)
2709            } else {
2710                let carried = self.state.matcher.uncommitted_input().len();
2711                let want = block_capacity.saturating_sub(carried);
2712                self.state
2713                    .matcher
2714                    .fill_in_place(want, &mut |buf| {
2715                        source.fill_block(buf, buf.len() + want, size_hint_remaining)
2716                    })
2717                    .map(|(appended, eof)| {
2718                        total_uncompressed += appended as u64;
2719                        reached_eof = eof;
2720                        appended
2721                    })
2722            };
2723
2724            let mut uncompressed_data;
2725            if in_place.is_some() {
2726                // Bytes live in the matcher; nothing staged here.
2727                uncompressed_data = Vec::new();
2728            } else {
2729                uncompressed_data = self.state.matcher.get_next_space();
2730                uncompressed_data.clear();
2731                uncompressed_data.extend_from_slice(&pending_input);
2732                pending_input.clear();
2733                if !reached_eof {
2734                    let (appended, eof) = source.fill_block(
2735                        &mut uncompressed_data,
2736                        block_capacity,
2737                        size_hint_remaining,
2738                    );
2739                    total_uncompressed += appended as u64;
2740                    reached_eof = eof;
2741                }
2742            }
2743            // Unified view of this iteration's candidate bytes, whichever path
2744            // produced them. Length only — the bytes themselves are read back
2745            // through the matcher on the in-place path.
2746            let available = if in_place.is_some() {
2747                self.state.matcher.uncommitted_input().len()
2748            } else {
2749                uncompressed_data.len()
2750            };
2751            let mut last_block = reached_eof;
2752            let remaining_for_split = if reached_eof {
2753                available
2754            } else {
2755                block_capacity
2756            };
2757            // Length this block will actually claim. The pre-split pass may
2758            // shorten it; on the in-place path the remainder simply stays
2759            // uncommitted in the matcher and heads the next block, so there is
2760            // no suffix copy at all.
2761            let mut block_len = available;
2762            if !matches!(self.compression_level, CompressionLevel::Uncompressed)
2763                && available == block_capacity
2764            {
2765                let split_at = {
2766                    let bytes: &[u8] = if in_place.is_some() {
2767                        self.state.matcher.uncommitted_input()
2768                    } else {
2769                        &uncompressed_data
2770                    };
2771                    optimal_block_size_with(
2772                        self.pre_split_level(),
2773                        bytes,
2774                        remaining_for_split,
2775                        block_capacity,
2776                        savings,
2777                    )
2778                };
2779                if split_at < available {
2780                    block_len = split_at;
2781                    last_block = false;
2782                    if in_place.is_none() {
2783                        // Staged path keeps its carry buffer: copy the kept
2784                        // suffix out and truncate the block being compressed.
2785                        pending_input.clear();
2786                        pending_input.extend_from_slice(&uncompressed_data[block_len..]);
2787                        uncompressed_data.truncate(block_len);
2788                    }
2789                }
2790            }
2791            // As we read, hash that data too (skipped when the content
2792            // checksum is disabled).
2793            #[cfg(feature = "hash")]
2794            if self.content_checksum {
2795                if in_place.is_some() {
2796                    let bytes = &self.state.matcher.uncommitted_input()[..block_len];
2797                    self.hasher.write(bytes);
2798                } else {
2799                    self.hasher.write(&uncompressed_data);
2800                }
2801            }
2802            // Per-physical-block XXH64 (low 32 bits) for the optional
2803            // per-block checksum sidecar. Hashing happens INSIDE the
2804            // block emitters (RLE / Raw fast-path / Compressed /
2805            // post-split partitions), so the digests vector has
2806            // exactly one entry per physical Block_Header written to
2807            // `all_blocks` — 1:1 with `FrameEmitInfo.blocks`. See
2808            // `enable_per_block_checksums` rustdoc.
2809            // Size the output ahead of this block's emission from the ratio
2810            // observed so far (see `reserve_for_next_block`); with no usable
2811            // size hint, ensure one block's worst case and let the doubling
2812            // growth policy amortize across blocks.
2813            // Bytes already emitted as blocks: everything read so far minus what
2814            // this block will claim and minus whatever stays buffered for the
2815            // next one (the staged carry, or the in-place uncommitted tail).
2816            let buffered_after = if in_place.is_some() {
2817                (available - block_len) as u64
2818            } else {
2819                pending_input.len() as u64
2820            };
2821            let emitted = total_uncompressed - block_len as u64 - buffered_after;
2822            match initial_size_hint {
2823                Some(hint) if hint >= total_uncompressed => {
2824                    // An advisory hint (streaming path) is only trusted up to
2825                    // a small lookahead past the bytes actually read: a hint
2826                    // far above the real input would otherwise reserve the
2827                    // whole phantom remainder up front.
2828                    let hint_remaining = hint - emitted;
2829                    let remaining = if hint_is_exact {
2830                        hint_remaining
2831                    } else {
2832                        let buffered = total_uncompressed - emitted;
2833                        const HINT_LOOKAHEAD: u64 = 64 * 1024;
2834                        hint_remaining.min(buffered + HINT_LOOKAHEAD)
2835                    };
2836                    reserve_for_next_block(
2837                        out,
2838                        blocks_start,
2839                        emitted,
2840                        remaining as usize,
2841                        self.block_capacity(),
2842                    );
2843                }
2844                _ => {
2845                    out.reserve(block_len + 3 + 16);
2846                }
2847            }
2848            // Special handling is needed for compression of a totally empty file
2849            if block_len == 0 {
2850                let header = BlockHeader {
2851                    last_block: true,
2852                    block_type: crate::blocks::block::BlockType::Raw,
2853                    block_size: 0,
2854                };
2855                header.serialize(out);
2856                #[cfg(feature = "lsm")]
2857                self.block_decompressed_sizes.push(0);
2858                #[cfg(all(feature = "lsm", feature = "hash"))]
2859                if let Some(checksums) = self.block_checksums.as_mut() {
2860                    checksums.push(xxh64_block_low32(&[]));
2861                }
2862                break;
2863            }
2864
2865            match self.compression_level {
2866                CompressionLevel::Uncompressed => {
2867                    // Always the staged buffer here — the ingest above refuses
2868                    // the in-place path for this level.
2869                    let header = BlockHeader {
2870                        last_block,
2871                        block_type: crate::blocks::block::BlockType::Raw,
2872                        block_size: uncompressed_data.len().try_into().unwrap(),
2873                    };
2874                    header.serialize(out);
2875                    #[cfg(feature = "lsm")]
2876                    self.block_decompressed_sizes
2877                        .push(uncompressed_data.len() as u32);
2878                    #[cfg(all(feature = "lsm", feature = "hash"))]
2879                    if let Some(checksums) = self.block_checksums.as_mut() {
2880                        checksums.push(xxh64_block_low32(&uncompressed_data));
2881                    }
2882                    out.extend_from_slice(&uncompressed_data);
2883                    savings +=
2884                        uncompressed_data.len() as i64 - (3 + uncompressed_data.len()) as i64;
2885                }
2886                CompressionLevel::Fastest
2887                | CompressionLevel::Default
2888                | CompressionLevel::Better
2889                | CompressionLevel::Best
2890                | CompressionLevel::Level(_) => {
2891                    let before_len = out.len();
2892                    // A primed dictionary makes "incompressible-looking"
2893                    // blocks matchable against the dict, so the raw-fast-
2894                    // path inside must be bypassed (it skips matching).
2895                    // Mirror prepare_frame's `use_dictionary_state`: a dict
2896                    // is only PRIMED (and thus matchable) when the matcher
2897                    // supports priming — a non-priming matcher ignores an
2898                    // attached dictionary, so the raw-fast-path must stay
2899                    // enabled for it. (This arm is already non-Uncompressed.)
2900                    let block_input = if in_place.is_some() {
2901                        crate::encoding::levels::BlockInput::InPlace(block_len)
2902                    } else {
2903                        crate::encoding::levels::BlockInput::Staged(uncompressed_data)
2904                    };
2905                    let dict_active = self.dictionary.is_some()
2906                        && self.state.matcher.supports_dictionary_priming();
2907                    compress_block_encoded(
2908                        &mut self.state,
2909                        self.compression_level,
2910                        last_block,
2911                        block_input,
2912                        out,
2913                        dict_active,
2914                        #[cfg(feature = "lsm")]
2915                        Some(&mut self.block_decompressed_sizes),
2916                        #[cfg(all(feature = "lsm", feature = "hash"))]
2917                        self.block_checksums.as_mut(),
2918                    );
2919                    savings += block_len as i64 - (out.len() - before_len) as i64;
2920                }
2921            }
2922            // The in-place path carries its remainder as uncommitted bytes in
2923            // the matcher rather than in `pending_input`, so the staged
2924            // emptiness test alone would exit while a split leftover still
2925            // needs a block.
2926            let carry_left = if in_place.is_some() {
2927                available - block_len
2928            } else {
2929                pending_input.len()
2930            };
2931            if last_block && carry_left == 0 {
2932                break;
2933            }
2934        }
2935        total_uncompressed
2936    }
2937
2938    /// Append the frame header bytes onto `out` once the total payload size
2939    /// is known (so `Frame_Content_Size` / `single_segment` can be set).
2940    /// Appends rather than returns so the one-shot path serializes straight
2941    /// into the reused output buffer with no per-frame header `Vec`.
2942    fn append_frame_header(&self, total_uncompressed: u64, prep: &FramePrep, out: &mut Vec<u8>) {
2943        // Match the upstream zstd framing policy (`ZSTD_writeFrameHeader`):
2944        // single-segment whenever the content size is known and the whole
2945        // source fits the active window (`contentSizeFlag && windowSize >=
2946        // srcSize`). A single-segment frame REQUIRES an FCS field, so
2947        // suppressing the content size (`content_size_flag` off) forces the
2948        // windowed layout. There is no lower size bound: small payloads
2949        // benefit most, since a windowed frame cannot encode a content size
2950        // below 256 in fewer than 4 FCS bytes (the 1-byte FCS class is
2951        // single-segment-only, see `find_fcs_field_size`), whereas a
2952        // single-segment frame stores it in one byte and omits the window
2953        // descriptor. The single-segment window equals the FCS, so a block
2954        // must never reference past the content: the post-hoc raw fallback in
2955        // the block emitters guarantees any non-shrinking block is stored raw,
2956        // and genuine matches stay within the already-emitted output.
2957        // Dictionary frames qualify too (the dictionary is decoder setup
2958        // state, not part of the regenerated segment), keeping the decoder's
2959        // single-allocation path (our decoder caps reservation to
2960        // min(window, FCS) either way).
2961        let single_segment = self.content_size_flag
2962            && prep.source_size_hint_known
2963            && total_uncompressed <= prep.window_size;
2964        let header = FrameHeader {
2965            frame_content_size: self.content_size_flag.then_some(total_uncompressed),
2966            single_segment,
2967            content_checksum: cfg!(feature = "hash") && self.content_checksum,
2968            dictionary_id: if prep.use_dictionary_state && self.dict_id_flag {
2969                // Id 0 is a raw-content dictionary: RFC 8878 spells "no
2970                // dictionary ID" as an absent field, not as a stored zero.
2971                self.dictionary
2972                    .as_ref()
2973                    .map(|dict| dict.inner.id)
2974                    .filter(|id| *id != 0)
2975                    .map(u64::from)
2976            } else {
2977                None
2978            },
2979            window_size: if single_segment {
2980                None
2981            } else {
2982                Some(prep.window_size)
2983            },
2984            magicless: self.magicless,
2985        };
2986        header.serialize(out);
2987    }
2988
2989    /// Write the frame header, accumulated block bytes, and optional
2990    /// trailing content checksum to the configured drain; populate
2991    /// `frame_emit_info` (lsm). Header and blocks are written separately to
2992    /// avoid shifting `all_blocks` to prepend the header. Used by
2993    /// `compress` and `compress_oneshot_borrowed`.
2994    fn finish_frame(&mut self, all_blocks: Vec<u8>, total_uncompressed: u64, prep: &FramePrep) {
2995        let mut header_buf: Vec<u8> = Vec::with_capacity(18);
2996        self.append_frame_header(total_uncompressed, prep, &mut header_buf);
2997        // Snapshot the checksum before borrowing the drain field so the
2998        // `self.hasher` read and the `self.compressed_data` write don't
2999        // both need `&mut self` simultaneously.
3000        #[cfg(feature = "hash")]
3001        let checksum_bytes = self
3002            .content_checksum
3003            .then(|| (self.hasher.finish() as u32).to_le_bytes());
3004        let drain = self.compressed_data.as_mut().unwrap();
3005        drain.write_all(&header_buf).unwrap();
3006        drain.write_all(&all_blocks).unwrap();
3007        // With the `hash` feature AND the content checksum enabled, the header
3008        // set `Content_Checksum_flag` and the 32-bit digest is written at the
3009        // end of the frame. Disabled => no trailing bytes, flag stays 0.
3010        #[cfg(feature = "hash")]
3011        if let Some(checksum_bytes) = checksum_bytes {
3012            drain.write_all(&checksum_bytes).unwrap();
3013        }
3014        #[cfg(feature = "lsm")]
3015        {
3016            let emit_checksum = cfg!(feature = "hash") && self.content_checksum;
3017            self.populate_frame_emit_info(header_buf.len(), &all_blocks, emit_checksum);
3018        }
3019    }
3020
3021    /// Assemble the frame (header + blocks + optional checksum) into the
3022    /// caller-provided `out` buffer, replacing its contents, and populate
3023    /// `frame_emit_info` (lsm). `out` is cleared first (its allocation is
3024    /// reused, the CCtx-equivalent zero-per-call-alloc output path) then
3025    /// grown once to the exact frame size. Used by
3026    /// `compress_independent_frame_into`. The single `all_blocks` copy into
3027    /// `out` is the same one copy `finish_frame` performs writing
3028    /// `all_blocks` into a `Vec` drain, no extra buffering vs the drain
3029    /// path.
3030    /// Walk `all_blocks` to recover per-block layout and store it in
3031    /// `frame_emit_info`. Each Block_Header is 3 bytes LE packing
3032    /// `(block_size << 3) | (block_type << 1) | last_block`. Physical body
3033    /// size differs by type: RLE bodies are always 1 byte (the repeated
3034    /// byte), Raw/Compressed bodies span `block_size`. `header_len` is the
3035    /// serialized frame-header length (frame offset of the first block).
3036    #[cfg(feature = "lsm")]
3037    fn populate_frame_emit_info(
3038        &mut self,
3039        header_len: usize,
3040        all_blocks: &[u8],
3041        emit_checksum: bool,
3042    ) {
3043        use crate::blocks::block::BlockType as BT;
3044        use crate::encoding::frame_emit_info::{FrameBlock, FrameEmitInfo};
3045        // All frame-offset arithmetic below is bounded by u32 on the wire
3046        // (Block_Size is a 21-bit field, frames bounded by MAX_BLOCK_SIZE *
3047        // #blocks). A pathologically large frame whose total emitted size
3048        // exceeds u32::MAX would overflow the cast; bail out by leaving
3049        // `frame_emit_info` at `None` rather than handing the caller a
3050        // silently-truncated layout. The overflow path is statically
3051        // unreachable on every realistic frame so the predictor amortises
3052        // the branch to zero cost.
3053        let frame_header_len: u32 = match u32::try_from(header_len) {
3054            Ok(v) => v,
3055            Err(_) => return,
3056        };
3057        let all_blocks_len_u32: u32 = match u32::try_from(all_blocks.len()) {
3058            Ok(v) => v,
3059            Err(_) => return,
3060        };
3061        let mut blocks: Vec<FrameBlock> = Vec::new();
3062        let mut cursor: usize = 0;
3063        while cursor + 3 <= all_blocks.len() {
3064            let mut header_u32 = [0u8; 4];
3065            header_u32[..3].copy_from_slice(&all_blocks[cursor..cursor + 3]);
3066            let raw = u32::from_le_bytes(header_u32);
3067            let last_block = (raw & 1) != 0;
3068            let block_type = match (raw >> 1) & 0b11 {
3069                0 => BT::Raw,
3070                1 => BT::RLE,
3071                2 => BT::Compressed,
3072                _ => BT::Reserved,
3073            };
3074            let block_size_field = raw >> 3;
3075            // RLE bodies are always 1 byte physical on the wire (the single
3076            // repeated byte); the spec's Block_Size field carries the
3077            // logical repeat count. Raw and Compressed bodies physically
3078            // span block_size_field bytes. Store the physical length in
3079            // body_size so the 'offset + header + body_size' arithmetic
3080            // always lands on the next block boundary, and surface the raw
3081            // spec field separately as block_size_field.
3082            let physical_body: u32 = match block_type {
3083                BT::RLE => 1,
3084                _ => block_size_field,
3085            };
3086            let cursor_u32: u32 = match u32::try_from(cursor) {
3087                Ok(v) => v,
3088                Err(_) => return,
3089            };
3090            let offset_in_frame = match frame_header_len.checked_add(cursor_u32) {
3091                Some(v) => v,
3092                None => return,
3093            };
3094            // Decompressed (regenerated) size, captured per physical block
3095            // during emit (1:1 with the wire blocks scanned here). Raw/RLE are
3096            // wire-derivable (`block_size_field`), so a short sidecar still
3097            // yields the correct value for them. A Compressed block's size is
3098            // NOT on the wire: if the sidecar is missing its entry, fabricating
3099            // 0 would publish a silently-wrong `decompressed_byte_range`. Since
3100            // this metadata is the authoritative mapping for a successful
3101            // encode, bail out (leave `frame_emit_info` at `None`) rather than
3102            // hand back a corrupt layout; the 1:1 push invariant makes this
3103            // unreachable in practice (debug_assert catches a regression).
3104            let decompressed_size = match self.block_decompressed_sizes.get(blocks.len()).copied() {
3105                Some(size) => size,
3106                None if matches!(block_type, BT::Raw | BT::RLE) => block_size_field,
3107                None => {
3108                    debug_assert!(
3109                        false,
3110                        "missing decompressed-size sidecar entry for compressed block {}",
3111                        blocks.len()
3112                    );
3113                    return;
3114                }
3115            };
3116            blocks.push(FrameBlock {
3117                offset_in_frame,
3118                header_size: 3,
3119                body_size: physical_body,
3120                block_size_field,
3121                block_type,
3122                last_block,
3123                decompressed_size,
3124            });
3125            cursor += 3 + physical_body as usize;
3126            if last_block {
3127                break;
3128            }
3129        }
3130        // Fail closed on a structurally incomplete scan: the loop must have
3131        // consumed the whole block section AND ended on a parsed last block.
3132        // A premature `last_block` (bytes left over) or a run-off without any
3133        // last block would otherwise publish an invalid public `FrameEmitInfo`.
3134        // Unreachable for a well-formed self-produced frame (debug_assert
3135        // catches a regression); on release we bail, leaving `frame_emit_info`
3136        // at `None` rather than handing back a corrupt layout.
3137        if cursor != all_blocks.len() || !blocks.last().is_some_and(|b| b.last_block) {
3138            debug_assert!(
3139                false,
3140                "incomplete block scan in populate_frame_emit_info: cursor={} len={} last_block={:?}",
3141                cursor,
3142                all_blocks.len(),
3143                blocks.last().map(|b| b.last_block)
3144            );
3145            return;
3146        }
3147        let checksum_range = if emit_checksum {
3148            let cs_start = match frame_header_len.checked_add(all_blocks_len_u32) {
3149                Some(v) => v,
3150                None => return,
3151            };
3152            let cs_end = match cs_start.checked_add(4) {
3153                Some(v) => v,
3154                None => return,
3155            };
3156            Some(cs_start..cs_end)
3157        } else {
3158            None
3159        };
3160        let body_total = match frame_header_len.checked_add(all_blocks_len_u32) {
3161            Some(v) => v,
3162            None => return,
3163        };
3164        let total_size = if checksum_range.is_some() {
3165            match body_total.checked_add(4) {
3166                Some(v) => v,
3167                None => return,
3168            }
3169        } else {
3170            body_total
3171        };
3172        self.frame_emit_info = Some(FrameEmitInfo {
3173            frame_header_range: 0..frame_header_len,
3174            blocks,
3175            checksum_range,
3176            total_size,
3177        });
3178    }
3179
3180    /// Layout of the most recently emitted frame.
3181    ///
3182    /// Returns `None` if [`compress`](Self::compress) has not been
3183    /// called yet on this compressor. After a successful `compress()`
3184    /// the returned `FrameEmitInfo` describes the frame header range,
3185    /// every emitted block's offset / size / type, and the optional
3186    /// trailing content-checksum range — all in frame-absolute byte
3187    /// offsets matching the bytes written to the drain.
3188    ///
3189    /// Behind the `lsm` Cargo feature.
3190    #[cfg(feature = "lsm")]
3191    pub fn last_frame_emit_info(&self) -> Option<&crate::encoding::frame_emit_info::FrameEmitInfo> {
3192        self.frame_emit_info.as_ref()
3193    }
3194
3195    /// Opt in to per-block XXH64 checksum computation during
3196    /// [`compress`](Self::compress). Default off; zero cost when
3197    /// disabled. The captured digests are accessible via
3198    /// [`last_frame_block_checksums`](Self::last_frame_block_checksums).
3199    ///
3200    /// One checksum is emitted per physical FrameBlock written to
3201    /// the drain: 1:1 cardinality with
3202    /// [`last_frame_emit_info`](Self::last_frame_emit_info)'s
3203    /// `blocks` vector. On the post-split optimization path
3204    /// (Level 16-22 with large window) the per-partition decompressed
3205    /// range is hashed inside the partition loop so the digest count
3206    /// still matches the emitted block count. The decoder collects
3207    /// per-physical-block digests on the same granularity, so
3208    /// element-wise equality holds round-trip.
3209    ///
3210    /// Behind `all(feature = "lsm", feature = "hash")` — the XXH64
3211    /// primitive lives behind the `hash` feature, so this method only
3212    /// compiles when both are enabled.
3213    #[cfg(all(feature = "lsm", feature = "hash"))]
3214    pub fn enable_per_block_checksums(&mut self) {
3215        self.per_block_checksums_enabled = true;
3216    }
3217
3218    /// Per-block XXH64 (low 32 bits) digests captured during the most
3219    /// recent `compress()` call. `None` unless
3220    /// [`enable_per_block_checksums`](Self::enable_per_block_checksums)
3221    /// was called before `compress()`.
3222    ///
3223    /// Behind `all(feature = "lsm", feature = "hash")`.
3224    #[cfg(all(feature = "lsm", feature = "hash"))]
3225    pub fn last_frame_block_checksums(&self) -> Option<&[u32]> {
3226        self.block_checksums.as_deref()
3227    }
3228
3229    /// Get a mutable reference to the source
3230    pub fn source_mut(&mut self) -> Option<&mut R> {
3231        self.uncompressed_data.as_mut()
3232    }
3233
3234    /// Get a mutable reference to the drain
3235    pub fn drain_mut(&mut self) -> Option<&mut W> {
3236        self.compressed_data.as_mut()
3237    }
3238
3239    /// Get a reference to the source
3240    pub fn source(&self) -> Option<&R> {
3241        self.uncompressed_data.as_ref()
3242    }
3243
3244    /// Get a reference to the drain
3245    pub fn drain(&self) -> Option<&W> {
3246        self.compressed_data.as_ref()
3247    }
3248
3249    /// Retrieve the source
3250    pub fn take_source(&mut self) -> Option<R> {
3251        self.uncompressed_data.take()
3252    }
3253
3254    /// Retrieve the drain
3255    pub fn take_drain(&mut self) -> Option<W> {
3256        self.compressed_data.take()
3257    }
3258
3259    /// Before calling [FrameCompressor::compress] you can replace the matcher
3260    pub fn replace_matcher(&mut self, mut match_generator: M) -> M {
3261        core::mem::swap(&mut match_generator, &mut self.state.matcher);
3262        match_generator
3263    }
3264
3265    /// Before calling [FrameCompressor::compress] you can replace the compression level.
3266    ///
3267    /// This also clears any fine-grained parameter overrides installed via
3268    /// [`set_parameters`](Self::set_parameters): reverting to a bare level
3269    /// means plain level-based tuning, not the previous frame's customized
3270    /// strategy / LDM / log / literal-mode overrides. To keep overriding, call
3271    /// [`set_parameters`](Self::set_parameters) again with the new base level.
3272    pub fn set_compression_level(
3273        &mut self,
3274        compression_level: CompressionLevel,
3275    ) -> CompressionLevel {
3276        let old = self.compression_level;
3277        self.compression_level = compression_level;
3278        // Resync the raw-literals gate: negative levels disable literal (Huffman)
3279        // compression (C `ZSTD_literalsCompressionIsDisabled`). `prepare_frame`
3280        // never recomputes this, so it must be refreshed on the level switch the
3281        // same way the constructors and `set_parameters` do.
3282        self.state.literal_compression_disabled = matches!(
3283            compression_level,
3284            CompressionLevel::Level(n) if n < 0
3285        );
3286        // Drop sticky overrides so the level switch yields plain geometry,
3287        // the literal mode included: the gate above is the bare level's rule.
3288        self.tuning = FrameTuning::default();
3289        self.state.matcher.clear_param_overrides();
3290        old
3291    }
3292
3293    /// Get the current compression level
3294    pub fn compression_level(&self) -> CompressionLevel {
3295        self.compression_level
3296    }
3297
3298    /// Attach a pre-parsed dictionary to be used for subsequent compressions.
3299    ///
3300    /// In compressed modes, the dictionary id is written only when the active
3301    /// matcher supports dictionary priming.
3302    /// Uncompressed mode and non-priming matchers ignore the attached dictionary
3303    /// at encode time.
3304    pub fn set_dictionary(
3305        &mut self,
3306        dictionary: crate::decoding::Dictionary,
3307    ) -> Result<Option<EncoderDictionary>, crate::decoding::errors::DictionaryDecodeError> {
3308        self.attach_dictionary(EncoderDictionary::from_dictionary(dictionary))
3309    }
3310
3311    /// Parse and attach a dictionary blob, in either of the two forms upstream
3312    /// `ZSTD_CCtx_loadDictionary` takes (`ZSTD_dct_auto`): a blob prefixed with
3313    /// [`DICTIONARY_MAGIC`](crate::decoding::DICTIONARY_MAGIC) is a serialized
3314    /// dictionary, and anything else is raw content, which is why any file can
3315    /// be handed to `zstd -D`. Raw content has no id, so the frame records none
3316    /// and the decoder must be given the same bytes explicitly.
3317    ///
3318    /// A serialized blob parses through the encoder-only path (skips the
3319    /// FSE/HUF decode lookup-table build the encoder never reads); the entropy
3320    /// ENCODER tables — and thus the emitted frame — are identical to a full
3321    /// parse. To reject anything but a serialized dictionary, parse with
3322    /// [`EncoderDictionary::from_bytes`] (upstream `ZSTD_dct_fullDict`) and
3323    /// attach the result.
3324    ///
3325    /// An empty buffer is how the same upstream entry point is told there is
3326    /// no dictionary: it clears whatever was attached and succeeds, returning
3327    /// it, rather than reporting a dictionary too small to use.
3328    pub fn set_dictionary_from_bytes(
3329        &mut self,
3330        raw_dictionary: &[u8],
3331    ) -> Result<Option<EncoderDictionary>, crate::decoding::errors::DictionaryDecodeError> {
3332        if raw_dictionary.is_empty() {
3333            return Ok(self.clear_dictionary());
3334        }
3335        self.attach_dictionary(EncoderDictionary::from_serialized_or_raw_content(
3336            raw_dictionary,
3337        )?)
3338    }
3339
3340    /// Attach an already-parsed [`EncoderDictionary`] without reparsing a raw
3341    /// blob.
3342    ///
3343    /// Accepts an `EncoderDictionary` produced once via
3344    /// [`EncoderDictionary::from_bytes`] / [`EncoderDictionary::from_dictionary`]
3345    /// or handed back by [`Self::clear_dictionary`] / the `set_dictionary*`
3346    /// return value, so callers can reattach or reuse a prepared dictionary
3347    /// across compressions without re-running the dictionary parse each time.
3348    /// Returns the previously-attached dictionary, if any.
3349    pub fn set_encoder_dictionary(
3350        &mut self,
3351        dictionary: EncoderDictionary,
3352    ) -> Result<Option<EncoderDictionary>, crate::decoding::errors::DictionaryDecodeError> {
3353        self.attach_dictionary(dictionary)
3354    }
3355
3356    /// The dictionary frames are compressed with, if one is attached.
3357    ///
3358    /// # Examples
3359    /// ```
3360    /// use structured_zstd::encoding::{CompressionLevel, EncoderDictionary, FrameCompressor};
3361    ///
3362    /// let dictionary = EncoderDictionary::from_serialized_or_raw_content(b"some shared history").unwrap();
3363    /// let mut compressor: FrameCompressor = FrameCompressor::new(CompressionLevel::Default);
3364    /// assert!(compressor.dictionary().is_none());
3365    /// compressor.set_encoder_dictionary(dictionary).unwrap();
3366    /// assert_eq!(compressor.dictionary().map(EncoderDictionary::id), Some(0));
3367    /// ```
3368    pub fn dictionary(&self) -> Option<&EncoderDictionary> {
3369        self.dictionary.as_ref()
3370    }
3371
3372    /// Remove the attached dictionary, returning it as an [`EncoderDictionary`].
3373    pub fn clear_dictionary(&mut self) -> Option<EncoderDictionary> {
3374        // Drop the CDict prime snapshot — it is keyed to the dictionary
3375        // being removed and must not be restored against a different (or no)
3376        // dictionary on the next frame.
3377        self.state.matcher.invalidate_primed_dictionary();
3378        self.dictionary.take()
3379    }
3380
3381    /// Validate `enc`, store it, and return the previously-attached
3382    /// dictionary; its entropy tables come with it, built when it was
3383    /// prepared. Shared by every public attach entry point: `set_dictionary`,
3384    /// `set_dictionary_from_bytes`, and `set_encoder_dictionary`.
3385    fn attach_dictionary(
3386        &mut self,
3387        enc: EncoderDictionary,
3388    ) -> Result<Option<EncoderDictionary>, crate::decoding::errors::DictionaryDecodeError> {
3389        // A zero id is not an error here: it marks a raw-content dictionary,
3390        // which has no header to carry one. The frame then records no
3391        // dictionary ID, so the decoder has to be handed the same bytes.
3392        let dictionary = &enc.inner;
3393        if let Some(index) = dictionary.offset_hist.iter().position(|&rep| rep == 0) {
3394            return Err(
3395                crate::decoding::errors::DictionaryDecodeError::ZeroRepeatOffsetInDictionary {
3396                    index: index as u8,
3397                },
3398            );
3399        }
3400        // A previously-captured CDict prime snapshot belongs to the OLD
3401        // dictionary; drop it so the first frame with the new dictionary
3402        // re-primes (and re-captures) instead of restoring stale tables.
3403        self.state.matcher.invalidate_primed_dictionary();
3404        Ok(self.dictionary.replace(enc))
3405    }
3406}
3407
3408#[cfg(test)]
3409mod tests;