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