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