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