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