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