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