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::match_generator::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}
704
705impl<M: Matcher> CompressState<M> {
706 /// Clears `last_huff_table`, parking the table's buffers in
707 /// `huff_table_spare` for reuse instead of dropping them.
708 #[inline]
709 pub(crate) fn clear_huff_table(&mut self) {
710 if let Some(table) = self.last_huff_table.take() {
711 self.huff_table_spare = Some(table);
712 }
713 }
714
715 /// Replaces `last_huff_table` with `table`, parking any displaced table
716 /// in `huff_table_spare` for reuse.
717 #[inline]
718 pub(crate) fn replace_huff_table(&mut self, table: crate::huff0::huff0_encoder::HuffmanTable) {
719 if let Some(old) = self.last_huff_table.replace(table) {
720 self.huff_table_spare = Some(old);
721 }
722 }
723}
724
725/// Per-frame setup resolved once by [`FrameCompressor::prepare_frame`] and
726/// consumed by the block loop + [`FrameCompressor::finish_frame`]. Lets the
727/// owned `compress()` and the borrowed one-shot path share identical
728/// reset / dict-prime / entropy-seed setup and frame-tail emission.
729struct FramePrep {
730 window_size: u64,
731 use_dictionary_state: bool,
732 source_size_hint_known: bool,
733 initial_size_hint: Option<u64>,
734}
735
736/// Initial capacity for the `all_blocks` accumulator, by source-size hint.
737/// The frame header is written only after all input is read (so
738/// Frame_Content_Size is known), so compressed blocks accumulate in memory
739/// first. Seed-size tiers (mirrors upstream zstd `ZSTD_CStreamOutSize` naming):
740/// - tiny (`<= 4 KiB` hint): payload-bound seed, `>=` anything a tiny input's
741/// compressed output could need.
742/// - small (`<= 64 KiB` hint): absorbs one or two `Vec::extend` doublings
743/// without over-allocating.
744/// - default (one upstream zstd block, `130 KiB`): the value the rest of the encoder
745/// is sized around; larger inputs amortise the first doublings cheaply and
746/// the residue is dominated by internal `compress_block_encoded` buffers.
747///
748/// Shared by the owned (`run_owned_block_loop`) and borrowed
749/// (`run_borrowed_block_loop`) paths so the tier table can't drift between them.
750///
751/// `block_capacity` (the active `targetCBlockSize` cap, or the 128 KiB
752/// format ceiling) bounds every tier: with a small target the first
753/// allocation tracks one capped block + header/checksum slack instead of
754/// keeping the upstream zstd-sized floor that only later growth respects.
755fn initial_all_blocks_cap(initial_size_hint: Option<u64>, block_capacity: usize) -> usize {
756 const TINY_THRESHOLD: u64 = 4 * 1024;
757 const SMALL_THRESHOLD: u64 = 64 * 1024;
758 const TINY_CAP: usize = 4 * 1024;
759 const SMALL_CAP: usize = 16 * 1024;
760 const DEFAULT_CAP: usize = 130 * 1024;
761 let first_block_cap = block_capacity + 3 + 16;
762 match initial_size_hint {
763 Some(h) if h <= TINY_THRESHOLD => TINY_CAP.min(first_block_cap),
764 Some(h) if h <= SMALL_THRESHOLD => SMALL_CAP.min(first_block_cap),
765 _ => DEFAULT_CAP.min(first_block_cap),
766 }
767}
768
769/// Per-block feeder for `run_owned_block_loop`.
770///
771/// `fill_block` appends source bytes to `buf` (which already holds any
772/// carried pre-split suffix) until `buf.len() == block_capacity` or the
773/// source is exhausted, returning `(bytes_appended, reached_eof)`.
774/// `reached_eof` is true iff the block could NOT be filled to
775/// `block_capacity` — the boundary the historical `Read`-loop produced (an
776/// input that is an exact multiple of the block size still yields a
777/// trailing empty last block on the next iteration).
778///
779/// The slice impl exists so the slice entry points
780/// (`compress_independent_frame_into`, `compress_oneshot_*` fallbacks)
781/// append with one `extend_from_slice` — the generic reader impl must
782/// `resize` an initialized target region before `Read::read` can fill it,
783/// which costs a zero-fill memset of the whole block on every frame.
784pub(crate) trait OwnedBlockSource {
785 fn fill_block(
786 &mut self,
787 buf: &mut Vec<u8>,
788 block_capacity: usize,
789 size_hint_remaining: Option<u64>,
790 ) -> (usize, bool);
791}
792
793impl OwnedBlockSource for &[u8] {
794 fn fill_block(
795 &mut self,
796 buf: &mut Vec<u8>,
797 block_capacity: usize,
798 _size_hint_remaining: Option<u64>,
799 ) -> (usize, bool) {
800 let want = block_capacity - buf.len();
801 let take = want.min(self.len());
802 buf.extend_from_slice(&self[..take]);
803 *self = &self[take..];
804 (take, take < want)
805 }
806}
807
808/// Adapter routing a generic [`Read`] source through [`OwnedBlockSource`]:
809/// preserves the historical sizing behaviour — an initialized target region
810/// bounded by the source-size hint, grown (doubling, capped) only when the
811/// hint under-counted.
812pub(crate) struct ReaderBlockSource<Rd>(pub(crate) Rd);
813
814impl<Rd: Read> OwnedBlockSource for ReaderBlockSource<Rd> {
815 fn fill_block(
816 &mut self,
817 buf: &mut Vec<u8>,
818 block_capacity: usize,
819 size_hint_remaining: Option<u64>,
820 ) -> (usize, bool) {
821 let start = buf.len();
822 let mut filled = start;
823 let mut reached_eof = false;
824 // Size the read buffer to the bytes this block actually expects
825 // rather than always zero-filling a full MAX_BLOCK_SIZE: a small
826 // frame otherwise pays a 128 KiB `resize(_, 0)` memset per block
827 // just to read a few KiB (the zero-fill past `filled` is then
828 // truncated away).
829 //
830 // Overflow-free by construction (no `saturating_*` masking):
831 // `filled <= block_capacity` always (the read only ever targets
832 // `[filled..len]` with `len <= block_capacity`, and a carried-over
833 // pre-split suffix is a `split_off` below `block_capacity`), so
834 // `block_capacity - filled` never underflows; pinning `remaining`
835 // to `block_capacity` before the `usize` cast keeps the cast and
836 // the final add within `usize` on every target.
837 let initial_target = match size_hint_remaining {
838 Some(remaining) => {
839 let remaining = remaining.min(block_capacity as u64) as usize;
840 filled + remaining.min(block_capacity - filled)
841 }
842 // Unknown hint, or an inexact hint already met by prior blocks:
843 // read against the full block window.
844 None => block_capacity,
845 };
846 if buf.len() < initial_target {
847 buf.resize(initial_target, 0);
848 }
849 loop {
850 if reached_eof || filled == block_capacity {
851 break;
852 }
853 if filled == buf.len() {
854 // Hint under-counted the block; grow toward block_capacity
855 // (doubling, capped) so reading continues without paying a
856 // full-buffer zero up front. `len <= block_capacity` so the
857 // double stays well within `usize`; `filled < block_capacity`
858 // here (the `== block_capacity` break fired otherwise), so
859 // `filled + 1 <= block_capacity`.
860 let grow_to = (buf.len() * 2).clamp(filled + 1, block_capacity);
861 buf.resize(grow_to, 0);
862 }
863 let read_end = buf.len();
864 let new_bytes = self.0.read(&mut buf[filled..read_end]).unwrap();
865 if new_bytes == 0 {
866 reached_eof = true;
867 break;
868 }
869 filled += new_bytes;
870 }
871 buf.truncate(filled);
872 (filled - start, reached_eof)
873 }
874}
875
876impl<R: Read, W: Write> FrameCompressor<R, W, MatchGeneratorDriver> {
877 /// Create a new `FrameCompressor`
878 pub fn new(compression_level: CompressionLevel) -> Self {
879 Self {
880 uncompressed_data: None,
881 compressed_data: None,
882 compression_level,
883 dictionary: None,
884 dictionary_entropy_cache: None,
885 source_size_hint: None,
886 state: CompressState {
887 matcher: MatchGeneratorDriver::new(1024 * 128, 1),
888 last_huff_table: None,
889 huff_table_spare: None,
890 fse_tables: FseTables::new(),
891 block_scratch: crate::encoding::blocks::CompressedBlockScratch::new(),
892 offset_hist: [1, 4, 8],
893 strategy_tag: crate::encoding::strategy::StrategyTag::for_compression_level(
894 compression_level,
895 ),
896 },
897 magicless: false,
898 content_checksum: false,
899 content_size_flag: true,
900 dict_id_flag: true,
901 target_block_size: None,
902 #[cfg(feature = "hash")]
903 hasher: XxHash64::with_seed(0),
904 #[cfg(feature = "lsm")]
905 frame_emit_info: None,
906 #[cfg(all(feature = "lsm", feature = "hash"))]
907 per_block_checksums_enabled: false,
908 #[cfg(all(feature = "lsm", feature = "hash"))]
909 block_checksums: None,
910 #[cfg(feature = "lsm")]
911 block_decompressed_sizes: alloc::vec::Vec::new(),
912 strategy_override: None,
913 }
914 }
915
916 /// Configure fine-grained compression parameters (#27).
917 ///
918 /// Resets the base [`CompressionLevel`](crate::encoding::CompressionLevel)
919 /// to the parameters' level and installs the per-knob overrides
920 /// (window/hash/chain/search logs, strategy, LDM) applied at the next
921 /// frame. Pass `None`-equivalent (a builder that overrides nothing)
922 /// to fall back to plain level-based compression.
923 ///
924 /// ```rust
925 /// use structured_zstd::encoding::{
926 /// CompressionLevel, CompressionParameters, FrameCompressor, Strategy,
927 /// };
928 /// let params = CompressionParameters::builder(CompressionLevel::Level(19))
929 /// .strategy(Strategy::Btultra2)
930 /// .enable_long_distance_matching(true)
931 /// .build()
932 /// .unwrap();
933 /// let mut compressor: FrameCompressor = FrameCompressor::new(CompressionLevel::Default);
934 /// compressor.set_parameters(¶ms);
935 /// let compressed = compressor.compress_independent_frame(b"some data to compress");
936 /// assert!(!compressed.is_empty());
937 /// ```
938 pub fn set_parameters(&mut self, params: &crate::encoding::CompressionParameters) {
939 self.compression_level = params.level();
940 let overrides = params.overrides();
941 self.strategy_override = overrides.strategy.map(|s| s.tag());
942 // Keep `state.strategy_tag` consistent immediately so the borrowed
943 // one-shot eligibility gate (`borrowed_eligible`) and literal gates
944 // are correct even before the next `compress()` re-sync.
945 self.state.strategy_tag = self.strategy_override.unwrap_or_else(|| {
946 crate::encoding::strategy::StrategyTag::for_compression_level(self.compression_level)
947 });
948 self.state.matcher.set_param_overrides(Some(overrides));
949 }
950
951 /// Whether the borrowed (no per-block history copy) one-shot loop is
952 /// valid for an `input_len`-byte slice under the resolved `prep`.
953 ///
954 /// `Uncompressed` resolves to `StrategyTag::Fast` but must emit stored
955 /// Raw blocks, which the borrowed loop's
956 /// `compress_block_encoded_borrowed` (RLE/raw-fast/compressed) does NOT
957 /// do, so exclude it; it then takes the owned path's dedicated
958 /// Uncompressed arm.
959 ///
960 /// No window-size gate: over-window inputs are handled too. The owned
961 /// path bounds matches to the last `advertised_window` bytes via
962 /// `window_low` and evicts/rehashes its history; the borrowed path
963 /// computes the identical `window_low = block_end - advertised_window`
964 /// and the kernel rejects any hash candidate below it, while the
965 /// per-position `put` during the scan keeps in-window slots current,
966 /// so it produces byte-identical output to the owned (evicting) path
967 /// without ever copying the input into `history`, even when the input
968 /// far exceeds the window.
969 ///
970 /// BUT gate on `input_len <= u32::MAX`: the Fast kernel stores ABSOLUTE
971 /// positions in a `u32` hash table, and the borrowed scan walks
972 /// absolute input offsets up to `block_end == input.len()`. Past 4 GiB
973 /// those offsets truncate / overflow the `u32` position math
974 /// (`base_off + ip0 as u32`, `window_low`), panicking or corrupting.
975 /// The owned/evicting path keeps the scanned window bounded (positions
976 /// stay small), so >4 GiB inputs fall back to it.
977 fn borrowed_eligible(&self, input_len: usize, prep: &FramePrep) -> bool {
978 if matches!(self.compression_level, CompressionLevel::Uncompressed)
979 || input_len > u32::MAX as usize
980 {
981 return false;
982 }
983 if prep.use_dictionary_state {
984 // The borrowed dict scan runs in VIRTUAL `[dict][input]` coordinates,
985 // so the position space is `dict_content.len() + input_len`, not just
986 // `input_len`. A large attached dictionary plus an otherwise-allowed
987 // input can exceed the `u32` floor the kernel asserts — fall back to
988 // the owned (copy) path in that case.
989 let fits_u32 = self
990 .dictionary
991 .as_ref()
992 .and_then(|dict| dict.inner.dict_content.len().checked_add(input_len))
993 .is_some_and(|virtual_len| virtual_len <= u32::MAX as usize);
994 if !fits_u32 {
995 return false;
996 }
997 // Dictionary frames: only the Simple (Fast) backend in attach mode
998 // has a borrowed (no input copy) dict scan. Copy-mode dict frames
999 // and the other backends still take the owned path.
1000 return self.state.matcher.borrowed_dict_supported();
1001 }
1002 // The borrowed (no-copy, in-place over-window) scan exists for the
1003 // Simple (Fast), Dfast, and Row backends, and for the HashChain
1004 // backend's lazy CHAIN parser; BT/optimal (BinaryTree search) stay on
1005 // the owned path. Every borrowed scan applies the per-position
1006 // `window_low = abs_ip - advertised_window` offset cap so over-window
1007 // inputs are matched in place (no input->history copy), matching C's
1008 // continuous-index + windowLow one-shot behaviour.
1009 self.state.matcher.borrowed_supported()
1010 }
1011
1012 /// Compress `input` as one frame's worth of blocks into `out` (appended
1013 /// from its current end): the borrowed in-place loop when
1014 /// [`Self::borrowed_eligible`], else the owned (history-copying) loop fed
1015 /// an in-place `&[u8]` cursor. Returns `total_uncompressed`; the caller
1016 /// emits the frame header (before this call, when the content size is
1017 /// known) or the drain tail.
1018 fn run_one_frame(&mut self, input: &[u8], prep: &FramePrep, out: &mut Vec<u8>) -> u64 {
1019 if self.borrowed_eligible(input.len(), prep) {
1020 self.run_borrowed_block_loop(input, out)
1021 } else {
1022 let mut cursor: &[u8] = input;
1023 self.run_owned_block_loop(&mut cursor, prep.initial_size_hint, true, out)
1024 }
1025 }
1026
1027 /// Compress one contiguous `&[u8]` as a single independent Zstd frame,
1028 /// writing the frame bytes into `out` (its previous contents are
1029 /// replaced and its allocation reused), reusing this compressor's heavy
1030 /// state across calls.
1031 ///
1032 /// This is the reusable-compression-context (CCtx-equivalent) entry
1033 /// point, mirroring C `ZSTD_compress2` over a reused `ZSTD_CCtx`:
1034 /// construct ONE `FrameCompressor` and call this in a loop to emit N
1035 /// independent, self-describing frames (each carrying its own header,
1036 /// blocks, and checksum, decodable in isolation, with no cross-frame
1037 /// match history). Every call resets the per-frame state via
1038 /// [`Self::prepare_frame`]: only the allocations are kept, so the
1039 /// dominant per-frame setup cost (table allocation + dictionary prime)
1040 /// is paid once instead of N times. Passing the same `out` buffer each
1041 /// call additionally reuses the output allocation, matching C's
1042 /// caller-owned `dst` buffer (no per-frame output allocation).
1043 ///
1044 /// Reusing the context + `out` across many small frames (the typical
1045 /// per-block-frame workload) is far cheaper than a fresh
1046 /// [`compress_slice_to_vec`](crate::encoding::compress_slice_to_vec)
1047 /// per block, which allocates and primes from scratch each time.
1048 ///
1049 /// The input is read in place: no [`Self::set_source`] /
1050 /// [`Self::set_drain`] setup is required, and the input lifetime is not
1051 /// baked into the compressor type, so successive calls may pass slices
1052 /// with unrelated lifetimes. When the Fast (Simple) backend is active
1053 /// and no dictionary is set, the matcher references the input directly
1054 /// (no per-block history copy); other backends / dictionary use copy
1055 /// each block into history exactly as the streaming
1056 /// [`compress`](Self::compress) path does. The source-size hint is
1057 /// derived from the input length on every call, so per-frame table
1058 /// sizing tracks each frame's actual size regardless of any earlier
1059 /// hint.
1060 ///
1061 /// A sticky dictionary set via
1062 /// [`set_dictionary`](Self::set_dictionary) (or its variants) is primed
1063 /// into every frame, mirroring `ZSTD_CCtx_loadDictionary` /
1064 /// `ZSTD_CCtx_refCDict`.
1065 ///
1066 /// # Panics
1067 ///
1068 /// Panics on encoder error, matching [`Self::compress`] and
1069 /// [`compress_slice_to_vec`](crate::encoding::compress_slice_to_vec).
1070 pub fn compress_independent_frame_into(&mut self, input: &[u8], out: &mut Vec<u8>) {
1071 // Size the next frame from the actual payload, not a stale hint a
1072 // previous call may have left behind (a wrong hint would change the
1073 // resolved window/header and could flip borrowed eligibility).
1074 self.source_size_hint = Some(input.len() as u64);
1075 let prep = self.prepare_frame();
1076 // Content size is known up front (one-shot), so write the frame
1077 // header FIRST and emit blocks STRAIGHT into `out` — no separate
1078 // `all_blocks` accumulator and no header+blocks copy (which was the
1079 // dominant per-frame memmove + the only un-amortized per-frame alloc
1080 // even when the compressor is reused).
1081 let total_uncompressed = input.len() as u64;
1082 let emit_checksum = cfg!(feature = "hash") && self.content_checksum;
1083 let checksum_len = if emit_checksum { 4 } else { 0 };
1084 out.clear();
1085 // Reserve the header plus ONE block's worst case up front; the block
1086 // loops then grow `out` from the compression ratio observed so far
1087 // (`reserve_for_next_block`). Reserving `compress_bound(input_len)`
1088 // here held a whole-input-sized allocation for the entire frame —
1089 // ~100 MiB peak on a 100 MiB stream whose compressed output is a few
1090 // MiB, where the reference implementation's context peaks at
1091 // window-sized state. Small frames (<= one block) still get their
1092 // full bound in one shot, so the reused-`out` steady state is
1093 // unchanged. 18 = max frame header (magic 4 + descriptor 1 + window
1094 // 1 + dict id 4 + FCS 8).
1095 let first_block_bound = input.len().min(self.block_capacity()) + 3;
1096 out.reserve(18 + first_block_bound + checksum_len);
1097 self.append_frame_header(total_uncompressed, &prep, out);
1098 let header_len = out.len();
1099 let _ = self.run_one_frame(input, &prep, out);
1100 #[cfg(feature = "hash")]
1101 if self.content_checksum {
1102 out.extend_from_slice(&(self.hasher.finish() as u32).to_le_bytes());
1103 }
1104 #[cfg(feature = "lsm")]
1105 {
1106 let blocks_end = out.len() - checksum_len;
1107 self.populate_frame_emit_info(header_len, &out[header_len..blocks_end], emit_checksum);
1108 }
1109 #[cfg(not(feature = "lsm"))]
1110 let _ = header_len;
1111 }
1112
1113 /// Convenience wrapper over [`Self::compress_independent_frame_into`]
1114 /// that allocates and returns a fresh `Vec` per call. Prefer the
1115 /// `_into` form in tight per-block-frame loops to reuse one output
1116 /// buffer across frames (the CCtx-equivalent zero-per-call-alloc
1117 /// output, matching C's caller-owned `dst`).
1118 ///
1119 /// ```rust
1120 /// use structured_zstd::encoding::{FrameCompressor, CompressionLevel};
1121 /// let mut cctx: FrameCompressor = FrameCompressor::new(CompressionLevel::Default);
1122 /// let frame_a = cctx.compress_independent_frame(b"first block payload");
1123 /// let frame_b = cctx.compress_independent_frame(b"second block payload");
1124 /// assert!(!frame_a.is_empty() && !frame_b.is_empty());
1125 /// ```
1126 pub fn compress_independent_frame(&mut self, input: &[u8]) -> Vec<u8> {
1127 let mut out = Vec::new();
1128 self.compress_independent_frame_into(input, &mut out);
1129 out
1130 }
1131
1132 /// Borrowed one-shot block loop: walks `input` in `MAX_BLOCK_SIZE`
1133 /// strides (the Fast backend never pre-splits, so boundaries match the
1134 /// owned loop), scanning each block range in place against the
1135 /// borrowed window via `compress_block_encoded_borrowed` — no
1136 /// per-block `commit_space` copy. Returns `(all_blocks,
1137 /// total_uncompressed)`. Caller guarantees Fast backend + no
1138 /// dictionary; over-window inputs are fine (matches are bounded by
1139 /// `window_low` exactly as the owned evicting path).
1140 fn run_borrowed_block_loop(&mut self, input: &[u8], out: &mut Vec<u8>) -> u64 {
1141 // Blocks are appended to `out` starting here. `out` may already hold
1142 // the frame header (the one-shot compress-into-Vec path writes it
1143 // first, since the content size is known up front, and the loop
1144 // emits blocks straight after it — no separate `all_blocks` Vec and
1145 // no header+blocks copy). Output-size reads below are taken RELATIVE
1146 // to `blocks_start` so a header prefix never skews the upstream zstd split
1147 // `savings` gate (which would change block boundaries / wire output).
1148 let blocks_start = out.len();
1149 let total_uncompressed = input.len() as u64;
1150 // Empty input: emit a single empty last Raw block (mirrors the
1151 // owned loop's empty-file special case).
1152 if input.is_empty() {
1153 let header = BlockHeader {
1154 last_block: true,
1155 block_type: crate::blocks::block::BlockType::Raw,
1156 block_size: 0,
1157 };
1158 header.serialize(out);
1159 #[cfg(feature = "lsm")]
1160 self.block_decompressed_sizes.push(0);
1161 #[cfg(all(feature = "lsm", feature = "hash"))]
1162 if let Some(checksums) = self.block_checksums.as_mut() {
1163 checksums.push(xxh64_block_low32(&[]));
1164 }
1165 return total_uncompressed;
1166 }
1167 // SAFETY: `input` outlives this call (held by the caller across
1168 // the call) and is not mutated. Only the Simple backend is active
1169 // (gated by `compress_oneshot_borrowed`).
1170 unsafe {
1171 self.state.matcher.set_borrowed_window(input);
1172 }
1173 // Panic-safety: clear the borrowed `(ptr, len)` on EVERY exit,
1174 // including an unwind from an `assert!` inside the block loop, so
1175 // a caught-and-reused compressor never retains a dangling window.
1176 // (The next frame's `reset()` also clears it before any read, but
1177 // this guard makes the invariant local and unwind-proof.)
1178 struct ClearBorrowedOnDrop(*mut MatchGeneratorDriver);
1179 impl Drop for ClearBorrowedOnDrop {
1180 fn drop(&mut self) {
1181 // SAFETY: at drop (normal return or unwind) the loop's
1182 // borrows of the matcher have ended, so this is the only
1183 // access. `addr_of_mut!` produced this pointer without an
1184 // intermediate `&mut`, so the interleaved `&mut` uses in
1185 // the loop did not invalidate it.
1186 unsafe { (*self.0).clear_borrowed_window() };
1187 }
1188 }
1189 let _clear_guard = ClearBorrowedOnDrop(core::ptr::addr_of_mut!(self.state.matcher));
1190 let block_capacity = self.block_capacity();
1191 let mut start = 0usize;
1192 while start < input.len() {
1193 reserve_for_next_block(
1194 out,
1195 blocks_start,
1196 start as u64,
1197 input.len() - start,
1198 block_capacity,
1199 );
1200 // Upstream zstd `ZSTD_compress_frameChunk`: size each block via the cheap
1201 // fingerprint pre-splitter so a full 128 KiB block is cut at a
1202 // statistical boundary when it pays. `savings = consumed -
1203 // produced` mirrors the upstream zstd gate (the first block and
1204 // incompressible input keep the full 128 KiB). The borrowed window
1205 // already spans the whole input, so a smaller block is just a
1206 // narrower `(block_start, block_end)` range into it.
1207 let savings = start as i64 - (out.len() - blocks_start) as i64;
1208 // Borrowed path only: warm the pre-split window before the
1209 // cache-cold strided fingerprint read. Gated to exactly the
1210 // conditions under which `optimal_block_size` reads `block`
1211 // (a pre-split level, a full 128 KiB block remaining, the
1212 // block-size cap admits a full block, and `savings >= 3` so the
1213 // splitter actually runs) — so non-pre-split levels, the first
1214 // block, and the trailing partial block pay nothing. See
1215 // `warm_presplit_window`.
1216 if savings >= 3
1217 && input.len() - start >= MAX_BLOCK_SIZE as usize
1218 && block_capacity >= MAX_BLOCK_SIZE as usize
1219 && crate::encoding::match_generator::level_pre_split(self.compression_level)
1220 .is_some()
1221 {
1222 warm_presplit_window(&input[start..start + MAX_BLOCK_SIZE as usize]);
1223 }
1224 let block_len = optimal_block_size(
1225 self.compression_level,
1226 &input[start..],
1227 input.len() - start,
1228 block_capacity,
1229 savings,
1230 );
1231 let end = (start + block_len).min(input.len());
1232 let block = &input[start..end];
1233 let last_block = end == input.len();
1234 #[cfg(feature = "hash")]
1235 if self.content_checksum {
1236 self.hasher.write(block);
1237 }
1238 crate::encoding::levels::compress_block_encoded_borrowed(
1239 &mut self.state,
1240 self.compression_level,
1241 last_block,
1242 block,
1243 start,
1244 end,
1245 out,
1246 #[cfg(feature = "lsm")]
1247 Some(&mut self.block_decompressed_sizes),
1248 #[cfg(all(feature = "lsm", feature = "hash"))]
1249 self.block_checksums.as_mut(),
1250 );
1251 start = end;
1252 }
1253 // `_clear_guard` drops here, clearing the borrowed window.
1254 total_uncompressed
1255 }
1256}
1257
1258impl<R: Read, W: Write, M: Matcher> FrameCompressor<R, W, M> {
1259 /// Create a new `FrameCompressor` with a custom matching algorithm implementation
1260 pub fn new_with_matcher(matcher: M, compression_level: CompressionLevel) -> Self {
1261 Self {
1262 uncompressed_data: None,
1263 compressed_data: None,
1264 dictionary: None,
1265 dictionary_entropy_cache: None,
1266 source_size_hint: None,
1267 state: CompressState {
1268 matcher,
1269 last_huff_table: None,
1270 huff_table_spare: None,
1271 fse_tables: FseTables::new(),
1272 block_scratch: crate::encoding::blocks::CompressedBlockScratch::new(),
1273 offset_hist: [1, 4, 8],
1274 strategy_tag: crate::encoding::strategy::StrategyTag::for_compression_level(
1275 compression_level,
1276 ),
1277 },
1278 compression_level,
1279 magicless: false,
1280 content_checksum: false,
1281 content_size_flag: true,
1282 dict_id_flag: true,
1283 target_block_size: None,
1284 #[cfg(feature = "hash")]
1285 hasher: XxHash64::with_seed(0),
1286 #[cfg(feature = "lsm")]
1287 frame_emit_info: None,
1288 #[cfg(all(feature = "lsm", feature = "hash"))]
1289 per_block_checksums_enabled: false,
1290 #[cfg(all(feature = "lsm", feature = "hash"))]
1291 block_checksums: None,
1292 #[cfg(feature = "lsm")]
1293 block_decompressed_sizes: alloc::vec::Vec::new(),
1294 strategy_override: None,
1295 }
1296 }
1297
1298 /// Enable or disable magicless frame format (`ZSTD_f_zstd1_magicless`).
1299 ///
1300 /// When set to `true`, emitted frames omit the 4-byte magic number
1301 /// prefix. The matching decoder must be configured to expect a
1302 /// magicless stream — wire-format only round-trips with a
1303 /// magicless-aware decoder.
1304 pub fn set_magicless(&mut self, magicless: bool) {
1305 self.magicless = magicless;
1306 }
1307
1308 /// Enable or disable the trailing XXH64 content checksum
1309 /// (semantics of upstream `ZSTD_c_checksumFlag`). Default `false`,
1310 /// matching the upstream library default (`ZSTD_c_checksumFlag = 0`)
1311 /// so out-of-the-box frames carry the same layout and pay the same
1312 /// costs as the reference implementation.
1313 ///
1314 /// When `false`, emitted frames set `Content_Checksum_flag = 0` and carry
1315 /// no trailing digest; such frames are valid (RFC 8878) and decode
1316 /// correctly in any [`ContentChecksum`](crate::decoding::ContentChecksum)
1317 /// mode. Without the `hash` feature no checksum is emitted regardless of
1318 /// this setting.
1319 pub fn set_content_checksum(&mut self, emit: bool) {
1320 self.content_checksum = emit;
1321 }
1322
1323 /// Enable or disable recording `Frame_Content_Size` in the frame header
1324 /// when the total size is known (semantics of upstream
1325 /// `ZSTD_c_contentSizeFlag`). Default `true`, matching upstream. With
1326 /// the flag off the header carries a window descriptor instead (and the
1327 /// single-segment layout, which requires an FCS, is disabled).
1328 pub fn set_content_size_flag(&mut self, emit: bool) {
1329 self.content_size_flag = emit;
1330 }
1331
1332 /// Enable or disable recording the dictionary ID in the frame header
1333 /// when a dictionary is attached (semantics of upstream
1334 /// `ZSTD_c_dictIDFlag`). Default `true`, matching upstream. Frames
1335 /// emitted with the flag off still decode when the decoder is handed
1336 /// the dictionary explicitly.
1337 pub fn set_dictionary_id_flag(&mut self, emit: bool) {
1338 self.dict_id_flag = emit;
1339 }
1340
1341 /// Set an upper bound on emitted block sizes (semantics of upstream
1342 /// `ZSTD_c_targetCBlockSize`): every physical block's payload is capped
1343 /// at `target` bytes (+3-byte block header on the wire), trading some
1344 /// ratio for bounded per-block latency. The value is clamped to
1345 /// `[MIN_TARGET_BLOCK_SIZE, MAX_BLOCK_SIZE]` (the upstream bounds).
1346 /// `None` removes the target.
1347 pub fn set_target_block_size(&mut self, target: Option<u32>) {
1348 self.target_block_size = target.map(|t| {
1349 t.clamp(
1350 crate::common::MIN_TARGET_BLOCK_SIZE,
1351 crate::common::MAX_BLOCK_SIZE,
1352 )
1353 });
1354 }
1355
1356 /// The active block-size cap: the configured target, or the format's
1357 /// 128 KiB block ceiling.
1358 fn block_capacity(&self) -> usize {
1359 self.target_block_size
1360 .map_or(crate::common::MAX_BLOCK_SIZE as usize, |t| t as usize)
1361 }
1362
1363 /// Before calling [FrameCompressor::compress] you need to set the source.
1364 ///
1365 /// This is the data that is compressed and written into the drain.
1366 pub fn set_source(&mut self, uncompressed_data: R) -> Option<R> {
1367 self.uncompressed_data.replace(uncompressed_data)
1368 }
1369
1370 /// Before calling [FrameCompressor::compress] you need to set the drain.
1371 ///
1372 /// As the compressor compresses data, the drain serves as a place for the output to be writte.
1373 pub fn set_drain(&mut self, compressed_data: W) -> Option<W> {
1374 self.compressed_data.replace(compressed_data)
1375 }
1376
1377 /// Provide a hint about the total uncompressed size for the next frame.
1378 ///
1379 /// When set, the encoder selects smaller hash tables and windows for
1380 /// small inputs, matching the C zstd source-size-class behavior.
1381 ///
1382 /// This hint applies only to frame payload bytes (`size`). Dictionary
1383 /// history is primed separately and does not inflate the hinted size or
1384 /// advertised frame window.
1385 /// Must be called before [`compress`](Self::compress).
1386 pub fn set_source_size_hint(&mut self, size: u64) {
1387 self.source_size_hint = Some(size);
1388 }
1389
1390 /// Total heap bytes this compressor's allocations hold, excluding the
1391 /// inline struct: the match-finder tables / history / recycled buffers and
1392 /// the primed-dictionary snapshot (via the matcher), the retained
1393 /// Huffman tables (active + recycled spare), the retained dictionary
1394 /// content, the cached dictionary entropy tables (literals Huffman +
1395 /// LL/ML/OF FSE), and the per-block sidecar buffers. Lets a context
1396 /// report its true footprint through `ZSTD_sizeof_CCtx`.
1397 pub fn heap_size(&self) -> usize {
1398 let mut total = self.state.matcher.heap_size();
1399 total += self
1400 .state
1401 .last_huff_table
1402 .as_ref()
1403 .map_or(0, |table| table.heap_size());
1404 total += self
1405 .state
1406 .huff_table_spare
1407 .as_ref()
1408 .map_or(0, |table| table.heap_size());
1409 total += self
1410 .dictionary
1411 .as_ref()
1412 .map_or(0, |d| d.inner.dict_content.capacity());
1413 total += self
1414 .dictionary_entropy_cache
1415 .as_ref()
1416 .map_or(0, CachedDictionaryEntropy::heap_size);
1417 #[cfg(all(feature = "lsm", feature = "hash"))]
1418 {
1419 total += self
1420 .block_checksums
1421 .as_ref()
1422 .map_or(0, |v| v.capacity() * core::mem::size_of::<u32>());
1423 }
1424 #[cfg(feature = "lsm")]
1425 {
1426 total += self.block_decompressed_sizes.capacity() * core::mem::size_of::<u32>();
1427 }
1428 total
1429 }
1430
1431 /// Compress the uncompressed data from the provided source as one Zstd frame and write it to the provided drain
1432 ///
1433 /// This will repeatedly call [Read::read] on the source to fill up blocks until the source returns 0 on the read call.
1434 /// All compressed blocks are buffered in memory so that the frame header can include the
1435 /// `Frame_Content_Size` field (which requires knowing the total uncompressed size). The
1436 /// entire frame — header, blocks, and optional checksum — is then written to the drain
1437 /// at the end. This means peak memory usage is O(compressed_size).
1438 ///
1439 /// To avoid endlessly encoding from a potentially endless source (like a network socket) you can use the
1440 /// [Read::take] function
1441 /// Per-frame setup values resolved by [`Self::prepare_frame`] and
1442 /// consumed by the block loop + [`Self::finish_frame`]. Lets the
1443 /// owned `compress()` and the borrowed one-shot path share the exact
1444 /// same reset / dict-prime / entropy-seed setup and frame tail.
1445 pub fn compress(&mut self) {
1446 let prep = self.prepare_frame();
1447 // Take the reader out so `run_owned_block_loop` can borrow it
1448 // mutably alongside `&mut self` (the rest of the loop touches
1449 // `self.state` / `self.hasher`, disjoint from the reader). Restored
1450 // before the frame tail so a reused compressor keeps its source.
1451 //
1452 // Deliberately NOT restored on unwind: if the block loop panics the
1453 // source has been partially consumed, so handing it back would let a
1454 // `catch_unwind` caller "successfully" compress the remaining tail
1455 // from an arbitrary midpoint — silent data corruption. Leaving the
1456 // slot empty makes any post-panic reuse fail loudly at the `expect`
1457 // below (matcher/entropy state is equally unre-usable after an
1458 // unwind; the reference implementation likewise requires a context
1459 // reset after an error).
1460 let mut source = self
1461 .uncompressed_data
1462 .take()
1463 .expect("source must be set via set_source before compress()");
1464 // Streaming drain: the content size is only known at EOF, so the
1465 // frame header can't precede the blocks — accumulate them in a local
1466 // buffer and let `finish_frame` write header + blocks to the drain.
1467 let mut all_blocks: Vec<u8> = Vec::with_capacity(initial_all_blocks_cap(
1468 prep.initial_size_hint,
1469 self.block_capacity(),
1470 ));
1471 let mut block_source = ReaderBlockSource(&mut source);
1472 let total_uncompressed = self.run_owned_block_loop(
1473 &mut block_source,
1474 prep.initial_size_hint,
1475 false,
1476 &mut all_blocks,
1477 );
1478 self.uncompressed_data = Some(source);
1479 self.finish_frame(all_blocks, total_uncompressed, &prep);
1480 }
1481
1482 fn prepare_frame(&mut self) -> FramePrep {
1483 // Reset per-frame introspection state so a re-used compressor
1484 // doesn't carry over the previous frame's layout/checksums.
1485 #[cfg(feature = "lsm")]
1486 {
1487 self.frame_emit_info = None;
1488 // Always captured under lsm (drives `decompressed_byte_range`);
1489 // clear, keep the allocation for a reused compressor.
1490 self.block_decompressed_sizes.clear();
1491 }
1492 #[cfg(all(feature = "lsm", feature = "hash"))]
1493 {
1494 if self.per_block_checksums_enabled {
1495 self.block_checksums = Some(alloc::vec::Vec::new());
1496 } else {
1497 self.block_checksums = None;
1498 }
1499 }
1500 let initial_size_hint = self.source_size_hint;
1501 let source_size_hint_known = initial_size_hint.is_some();
1502 let use_dictionary_state =
1503 !matches!(self.compression_level, CompressionLevel::Uncompressed)
1504 && self.state.matcher.supports_dictionary_priming()
1505 && self.dictionary.is_some();
1506 if let Some(size_hint) = self.source_size_hint.take() {
1507 // Keep source-size hint scoped to payload bytes; dictionary priming
1508 // is applied separately and should not force larger matcher sizing.
1509 self.state.matcher.set_source_size_hint(size_hint);
1510 }
1511 // Hand the matcher the dictionary's content size so its binary-tree /
1512 // hash-chain tables shrink to the dictionary's cParams tier (upstream zstd CDict
1513 // economics: the dictionary supplies long matches, so a source-sized live
1514 // table is wasted peak memory). The eviction window stays source-sized so
1515 // the dictionary bytes remain referenceable. Set before `reset` (which
1516 // consumes it) and only when a dictionary will actually be primed.
1517 if use_dictionary_state && let Some(dict) = self.dictionary.as_ref() {
1518 self.state
1519 .matcher
1520 .set_dictionary_size_hint(dict.inner.dict_content.len());
1521 }
1522 // Clearing buffers to allow re-using of the compressor
1523 self.state.matcher.reset(self.compression_level);
1524 self.state.offset_hist = [1, 4, 8];
1525 // Sync `state.strategy_tag` to the level resolved at this reset so
1526 // the literal-compression gates (`min_literals_to_compress` /
1527 // `min_gain` in `encoding::blocks::compressed`) see the correct
1528 // strategy for the next frame. Frame-by-frame level changes go
1529 // through this same `compress()` entry point, so re-syncing here
1530 // covers level switches without touching the matcher dispatch.
1531 // A public-parameter strategy override (#27) wins over the level's
1532 // derived tag so the literal-compression gates and dict-attach
1533 // cutoff below see the strategy the matcher actually runs.
1534 self.state.strategy_tag = self.strategy_override.unwrap_or_else(|| {
1535 crate::encoding::strategy::StrategyTag::for_compression_level(self.compression_level)
1536 });
1537 let cached_entropy = if use_dictionary_state {
1538 self.dictionary_entropy_cache.as_ref()
1539 } else {
1540 None
1541 };
1542 if use_dictionary_state && let Some(dict) = self.dictionary.as_ref() {
1543 // This state drives sequence encoding, while matcher priming below updates
1544 // the match generator's internal repeat-offset history for match finding.
1545 self.state.offset_hist = dict.inner.offset_hist;
1546 // Upstream zstd `ZSTD_shouldAttachDict` (`zstd_compress.c`): a
1547 // precomputed-dictionary table is COPIED into the working context
1548 // only when the source is larger than a per-strategy cutoff; at or
1549 // below it (and for unknown size) the upstream zstd ATTACHES the dictionary
1550 // tables by reference (no per-frame table touch at all). We don't
1551 // have an attach-by-reference path yet, so:
1552 // - large source (> cutoff): reuse the captured prime snapshot
1553 // (a table copy) instead of re-hashing the dictionary — the
1554 // upstream zstd COPY regime, where the copy is cheaper than re-priming;
1555 // - small / unknown source: re-prime (the snapshot copy of the
1556 // whole table would cost MORE than the sparse re-prime here,
1557 // which is exactly why the upstream zstd attaches by reference instead).
1558 // `attachDictSizeCutoffs` per strategy: fast 8K, dfast 16K,
1559 // greedy/lazy/btopt 32K, btultra/btultra2 8K. Expressed as the
1560 // ceil-log bucket (8K = 2^13, 16K = 2^14, 32K = 2^15) so the
1561 // decision uses the SAME bucketed representation as the driver's
1562 // attach/copy gate (`reset_size_log`) — comparing
1563 // `source_size_ceil_log(hint)` on the full u64 avoids the `as usize`
1564 // truncation that could diverge from the driver on 32-bit targets.
1565 // For a power-of-two cutoff `2^k`, `ceil_log2(hint) > k` is exactly
1566 // `hint > 2^k`, so this is identical to the raw `hint > cutoff` on
1567 // 64-bit.
1568 let cutoff_log = match self.state.strategy_tag {
1569 crate::encoding::strategy::StrategyTag::Fast
1570 | crate::encoding::strategy::StrategyTag::BtUltra
1571 | crate::encoding::strategy::StrategyTag::BtUltra2 => 13,
1572 crate::encoding::strategy::StrategyTag::Dfast => 14,
1573 crate::encoding::strategy::StrategyTag::Greedy
1574 | crate::encoding::strategy::StrategyTag::Lazy
1575 | crate::encoding::strategy::StrategyTag::Btlazy2
1576 | crate::encoding::strategy::StrategyTag::BtOpt => 15,
1577 };
1578 let prefer_copy_snapshot = initial_size_hint.is_some_and(|s| {
1579 crate::encoding::match_generator::source_size_ceil_log(s) > cutoff_log
1580 });
1581 let restored = prefer_copy_snapshot
1582 && self
1583 .state
1584 .matcher
1585 .restore_primed_dictionary(self.compression_level);
1586 if !restored {
1587 self.state.matcher.prime_with_dictionary(
1588 dict.inner.dict_content.as_slice(),
1589 dict.inner.offset_hist,
1590 );
1591 if prefer_copy_snapshot {
1592 self.state
1593 .matcher
1594 .capture_primed_dictionary(self.compression_level);
1595 }
1596 }
1597 }
1598 if let Some(cache) = cached_entropy {
1599 // Refill an empty slot from the recycled spare before
1600 // `clone_from`: `Option::clone_from(None ← Some)` falls back to
1601 // a fresh clone (two Vec allocations), while `Some ← Some`
1602 // delegates to the table's buffer-reusing `clone_from`. Frames
1603 // whose last block cleared the table would otherwise re-clone
1604 // the dict seed every frame.
1605 match &cache.huff {
1606 Some(src) => {
1607 if self.state.last_huff_table.is_none() {
1608 self.state.last_huff_table = self.state.huff_table_spare.take();
1609 }
1610 match &mut self.state.last_huff_table {
1611 Some(dst) => dst.clone_from(src),
1612 slot => *slot = Some(src.clone()),
1613 }
1614 }
1615 None => self.state.clear_huff_table(),
1616 }
1617 } else {
1618 self.state.clear_huff_table();
1619 }
1620 // `clone_from` keeps frame-to-frame seeding cheap for reused compressors by
1621 // reusing existing allocations where possible instead of reallocating every frame.
1622 if let Some(cache) = cached_entropy {
1623 self.state
1624 .fse_tables
1625 .ll_previous
1626 .clone_from(&cache.ll_previous);
1627 self.state
1628 .fse_tables
1629 .ml_previous
1630 .clone_from(&cache.ml_previous);
1631 self.state
1632 .fse_tables
1633 .of_previous
1634 .clone_from(&cache.of_previous);
1635 } else {
1636 self.state.fse_tables.ll_previous = None;
1637 self.state.fse_tables.ml_previous = None;
1638 self.state.fse_tables.of_previous = None;
1639 }
1640 let ll_entropy = cached_entropy.and_then(|cache| match cache.ll_previous.as_ref() {
1641 Some(PreviousFseTable::Custom(table)) => Some(table.as_ref()),
1642 _ => None,
1643 });
1644 let ml_entropy = cached_entropy.and_then(|cache| match cache.ml_previous.as_ref() {
1645 Some(PreviousFseTable::Custom(table)) => Some(table.as_ref()),
1646 _ => None,
1647 });
1648 let of_entropy = cached_entropy.and_then(|cache| match cache.of_previous.as_ref() {
1649 Some(PreviousFseTable::Custom(table)) => Some(table.as_ref()),
1650 _ => None,
1651 });
1652 self.state.matcher.seed_dictionary_entropy(
1653 self.state.last_huff_table.as_ref(),
1654 ll_entropy,
1655 ml_entropy,
1656 of_entropy,
1657 );
1658 #[cfg(feature = "hash")]
1659 {
1660 self.hasher = XxHash64::with_seed(0);
1661 }
1662 let window_size = self.state.matcher.window_size();
1663 assert!(
1664 window_size != 0,
1665 "matcher reported window_size == 0, which is invalid"
1666 );
1667 FramePrep {
1668 window_size,
1669 use_dictionary_state,
1670 source_size_hint_known,
1671 initial_size_hint,
1672 }
1673 }
1674
1675 /// Owned streaming block loop: reads blocks from the caller-provided
1676 /// `source` reader, optionally pre-splits, hashes for the content
1677 /// checksum, and emits each block via `compress_block_encoded`,
1678 /// accumulating the block bytes. Returns `(all_blocks,
1679 /// total_uncompressed)`. The source is passed in (rather than read
1680 /// from `self.uncompressed_data`) so the streaming `compress` path can
1681 /// feed the configured reader while the slice paths
1682 /// (`compress_oneshot_borrowed`, `compress_independent_frame`) feed an
1683 /// in-place `&[u8]` cursor without baking its lifetime into the
1684 /// compressor type.
1685 fn run_owned_block_loop<S: OwnedBlockSource>(
1686 &mut self,
1687 source: &mut S,
1688 initial_size_hint: Option<u64>,
1689 // Whether `initial_size_hint` is the input's exact length (the
1690 // one-shot slice paths) or a caller-provided estimate (the streaming
1691 // `Read` path, where `set_source_size_hint` is advisory). An exact
1692 // hint drives the one-shot ratio reservation; an estimate is only
1693 // trusted up to a small lookahead past the bytes actually read.
1694 hint_is_exact: bool,
1695 out: &mut Vec<u8>,
1696 ) -> u64 {
1697 // Compressed blocks are appended to `out` from its current end. The
1698 // streaming drain path passes a fresh buffer (the frame header is
1699 // written to the drain afterward, since Frame_Content_Size is only
1700 // known once the reader hits EOF); the one-shot compress-into-Vec
1701 // path passes `out` already holding the header. The upstream zstd split
1702 // `savings` gate below accumulates block-relative (`before_len`)
1703 // output deltas, so a header prefix never skews it.
1704 let blocks_start = out.len();
1705 let mut total_uncompressed: u64 = 0;
1706 let mut pending_input: Vec<u8> = Vec::new();
1707 let mut reached_eof = false;
1708 let mut savings = 0i64;
1709 // Compress block by block
1710 loop {
1711 // Read up to one upstream zstd block. When the pre-block splitter keeps a
1712 // suffix, top it back up before compressing the next block, matching
1713 // ZSTD_compress_frameChunk() over a contiguous input buffer.
1714 let block_capacity = self.block_capacity();
1715 // Always draw the block buffer from the matcher's recycled pool
1716 // (its capacity already covers the block size, so the resize below
1717 // stays in-place). Any carried pre-split suffix is copied in, and
1718 // `pending_input` is retained as a reusable carry buffer. The prior
1719 // approach `split_off`'d a fresh suffix Vec per pre-split and
1720 // `reserve_exact`-grew it to `block_capacity` every block; on a
1721 // heavily pre-split frame that churned one block-sized allocation
1722 // per split (~12 MB over ~90 splits on a 1 MiB corpus input).
1723 let mut uncompressed_data = self.state.matcher.get_next_space();
1724 uncompressed_data.clear();
1725 uncompressed_data.extend_from_slice(&pending_input);
1726 pending_input.clear();
1727 if !reached_eof {
1728 // Remaining-bytes expectation for the reader source's sizing
1729 // (`None` = unknown, or an inexact hint already met by prior
1730 // blocks). The slice source appends directly and ignores it.
1731 let size_hint_remaining = match initial_size_hint {
1732 Some(hint) if hint > total_uncompressed => Some(hint - total_uncompressed),
1733 _ => None,
1734 };
1735 let (appended, eof) =
1736 source.fill_block(&mut uncompressed_data, block_capacity, size_hint_remaining);
1737 total_uncompressed += appended as u64;
1738 reached_eof = eof;
1739 }
1740 let mut last_block = reached_eof;
1741 let remaining_for_split = if reached_eof {
1742 uncompressed_data.len()
1743 } else {
1744 block_capacity
1745 };
1746 if !matches!(self.compression_level, CompressionLevel::Uncompressed)
1747 && uncompressed_data.len() == block_capacity
1748 {
1749 let block_len = optimal_block_size(
1750 self.compression_level,
1751 &uncompressed_data,
1752 remaining_for_split,
1753 block_capacity,
1754 savings,
1755 );
1756 if block_len < uncompressed_data.len() {
1757 // Carry the kept suffix into the reusable `pending_input`
1758 // buffer (cleared, capacity retained) instead of allocating
1759 // a fresh Vec via `split_off`. Next iteration copies it back
1760 // into a pooled block buffer. The block currently being
1761 // compressed is truncated to the chosen split length.
1762 pending_input.clear();
1763 pending_input.extend_from_slice(&uncompressed_data[block_len..]);
1764 uncompressed_data.truncate(block_len);
1765 last_block = false;
1766 }
1767 }
1768 // As we read, hash that data too (skipped when the content
1769 // checksum is disabled).
1770 #[cfg(feature = "hash")]
1771 if self.content_checksum {
1772 self.hasher.write(&uncompressed_data);
1773 }
1774 // Per-physical-block XXH64 (low 32 bits) for the optional
1775 // per-block checksum sidecar. Hashing happens INSIDE the
1776 // block emitters (RLE / Raw fast-path / Compressed /
1777 // post-split partitions), so the digests vector has
1778 // exactly one entry per physical Block_Header written to
1779 // `all_blocks` — 1:1 with `FrameEmitInfo.blocks`. See
1780 // `enable_per_block_checksums` rustdoc.
1781 // Size the output ahead of this block's emission from the ratio
1782 // observed so far (see `reserve_for_next_block`); with no usable
1783 // size hint, ensure one block's worst case and let the doubling
1784 // growth policy amortize across blocks.
1785 let emitted =
1786 total_uncompressed - uncompressed_data.len() as u64 - pending_input.len() as u64;
1787 match initial_size_hint {
1788 Some(hint) if hint >= total_uncompressed => {
1789 // An advisory hint (streaming path) is only trusted up to
1790 // a small lookahead past the bytes actually read: a hint
1791 // far above the real input would otherwise reserve the
1792 // whole phantom remainder up front.
1793 let hint_remaining = hint - emitted;
1794 let remaining = if hint_is_exact {
1795 hint_remaining
1796 } else {
1797 let buffered = total_uncompressed - emitted;
1798 const HINT_LOOKAHEAD: u64 = 64 * 1024;
1799 hint_remaining.min(buffered + HINT_LOOKAHEAD)
1800 };
1801 reserve_for_next_block(
1802 out,
1803 blocks_start,
1804 emitted,
1805 remaining as usize,
1806 self.block_capacity(),
1807 );
1808 }
1809 _ => {
1810 out.reserve(uncompressed_data.len() + 3 + 16);
1811 }
1812 }
1813 // Special handling is needed for compression of a totally empty file
1814 if uncompressed_data.is_empty() {
1815 let header = BlockHeader {
1816 last_block: true,
1817 block_type: crate::blocks::block::BlockType::Raw,
1818 block_size: 0,
1819 };
1820 header.serialize(out);
1821 #[cfg(feature = "lsm")]
1822 self.block_decompressed_sizes.push(0);
1823 #[cfg(all(feature = "lsm", feature = "hash"))]
1824 if let Some(checksums) = self.block_checksums.as_mut() {
1825 checksums.push(xxh64_block_low32(&[]));
1826 }
1827 break;
1828 }
1829
1830 match self.compression_level {
1831 CompressionLevel::Uncompressed => {
1832 let header = BlockHeader {
1833 last_block,
1834 block_type: crate::blocks::block::BlockType::Raw,
1835 block_size: uncompressed_data.len().try_into().unwrap(),
1836 };
1837 header.serialize(out);
1838 #[cfg(feature = "lsm")]
1839 self.block_decompressed_sizes
1840 .push(uncompressed_data.len() as u32);
1841 #[cfg(all(feature = "lsm", feature = "hash"))]
1842 if let Some(checksums) = self.block_checksums.as_mut() {
1843 checksums.push(xxh64_block_low32(&uncompressed_data));
1844 }
1845 out.extend_from_slice(&uncompressed_data);
1846 savings +=
1847 uncompressed_data.len() as i64 - (3 + uncompressed_data.len()) as i64;
1848 }
1849 CompressionLevel::Fastest
1850 | CompressionLevel::Default
1851 | CompressionLevel::Better
1852 | CompressionLevel::Best
1853 | CompressionLevel::Level(_) => {
1854 let before_len = out.len();
1855 let block_len = uncompressed_data.len();
1856 // A primed dictionary makes "incompressible-looking"
1857 // blocks matchable against the dict, so the raw-fast-
1858 // path inside must be bypassed (it skips matching).
1859 // Mirror prepare_frame's `use_dictionary_state`: a dict
1860 // is only PRIMED (and thus matchable) when the matcher
1861 // supports priming — a non-priming matcher ignores an
1862 // attached dictionary, so the raw-fast-path must stay
1863 // enabled for it. (This arm is already non-Uncompressed.)
1864 let dict_active = self.dictionary.is_some()
1865 && self.state.matcher.supports_dictionary_priming();
1866 compress_block_encoded(
1867 &mut self.state,
1868 self.compression_level,
1869 last_block,
1870 uncompressed_data,
1871 out,
1872 dict_active,
1873 #[cfg(feature = "lsm")]
1874 Some(&mut self.block_decompressed_sizes),
1875 #[cfg(all(feature = "lsm", feature = "hash"))]
1876 self.block_checksums.as_mut(),
1877 );
1878 savings += block_len as i64 - (out.len() - before_len) as i64;
1879 }
1880 }
1881 if last_block && pending_input.is_empty() {
1882 break;
1883 }
1884 }
1885 total_uncompressed
1886 }
1887
1888 /// Append the frame header bytes onto `out` once the total payload size
1889 /// is known (so `Frame_Content_Size` / `single_segment` can be set).
1890 /// Appends rather than returns so the one-shot path serializes straight
1891 /// into the reused output buffer with no per-frame header `Vec`.
1892 fn append_frame_header(&self, total_uncompressed: u64, prep: &FramePrep, out: &mut Vec<u8>) {
1893 // Match the upstream zstd framing policy (`ZSTD_writeFrameHeader`):
1894 // single-segment whenever the content size is known and the whole
1895 // source fits the active window (`contentSizeFlag && windowSize >=
1896 // srcSize`). A single-segment frame REQUIRES an FCS field, so
1897 // suppressing the content size (`content_size_flag` off) forces the
1898 // windowed layout. There is no lower size bound: small payloads
1899 // benefit most, since a windowed frame cannot encode a content size
1900 // below 256 in fewer than 4 FCS bytes (the 1-byte FCS class is
1901 // single-segment-only, see `find_fcs_field_size`), whereas a
1902 // single-segment frame stores it in one byte and omits the window
1903 // descriptor. The single-segment window equals the FCS, so a block
1904 // must never reference past the content: the post-hoc raw fallback in
1905 // the block emitters guarantees any non-shrinking block is stored raw,
1906 // and genuine matches stay within the already-emitted output.
1907 // Dictionary frames qualify too (the dictionary is decoder setup
1908 // state, not part of the regenerated segment), keeping the decoder's
1909 // single-allocation path (our decoder caps reservation to
1910 // min(window, FCS) either way).
1911 let single_segment = self.content_size_flag
1912 && prep.source_size_hint_known
1913 && total_uncompressed <= prep.window_size;
1914 let header = FrameHeader {
1915 frame_content_size: self.content_size_flag.then_some(total_uncompressed),
1916 single_segment,
1917 content_checksum: cfg!(feature = "hash") && self.content_checksum,
1918 dictionary_id: if prep.use_dictionary_state && self.dict_id_flag {
1919 self.dictionary.as_ref().map(|dict| dict.inner.id as u64)
1920 } else {
1921 None
1922 },
1923 window_size: if single_segment {
1924 None
1925 } else {
1926 Some(prep.window_size)
1927 },
1928 magicless: self.magicless,
1929 };
1930 header.serialize(out);
1931 }
1932
1933 /// Write the frame header, accumulated block bytes, and optional
1934 /// trailing content checksum to the configured drain; populate
1935 /// `frame_emit_info` (lsm). Header and blocks are written separately to
1936 /// avoid shifting `all_blocks` to prepend the header. Used by
1937 /// `compress` and `compress_oneshot_borrowed`.
1938 fn finish_frame(&mut self, all_blocks: Vec<u8>, total_uncompressed: u64, prep: &FramePrep) {
1939 let mut header_buf: Vec<u8> = Vec::with_capacity(18);
1940 self.append_frame_header(total_uncompressed, prep, &mut header_buf);
1941 // Snapshot the checksum before borrowing the drain field so the
1942 // `self.hasher` read and the `self.compressed_data` write don't
1943 // both need `&mut self` simultaneously.
1944 #[cfg(feature = "hash")]
1945 let checksum_bytes = self
1946 .content_checksum
1947 .then(|| (self.hasher.finish() as u32).to_le_bytes());
1948 let drain = self.compressed_data.as_mut().unwrap();
1949 drain.write_all(&header_buf).unwrap();
1950 drain.write_all(&all_blocks).unwrap();
1951 // With the `hash` feature AND the content checksum enabled, the header
1952 // set `Content_Checksum_flag` and the 32-bit digest is written at the
1953 // end of the frame. Disabled => no trailing bytes, flag stays 0.
1954 #[cfg(feature = "hash")]
1955 if let Some(checksum_bytes) = checksum_bytes {
1956 drain.write_all(&checksum_bytes).unwrap();
1957 }
1958 #[cfg(feature = "lsm")]
1959 {
1960 let emit_checksum = cfg!(feature = "hash") && self.content_checksum;
1961 self.populate_frame_emit_info(header_buf.len(), &all_blocks, emit_checksum);
1962 }
1963 }
1964
1965 /// Assemble the frame (header + blocks + optional checksum) into the
1966 /// caller-provided `out` buffer, replacing its contents, and populate
1967 /// `frame_emit_info` (lsm). `out` is cleared first (its allocation is
1968 /// reused, the CCtx-equivalent zero-per-call-alloc output path) then
1969 /// grown once to the exact frame size. Used by
1970 /// `compress_independent_frame_into`. The single `all_blocks` copy into
1971 /// `out` is the same one copy `finish_frame` performs writing
1972 /// `all_blocks` into a `Vec` drain, no extra buffering vs the drain
1973 /// path.
1974 /// Walk `all_blocks` to recover per-block layout and store it in
1975 /// `frame_emit_info`. Each Block_Header is 3 bytes LE packing
1976 /// `(block_size << 3) | (block_type << 1) | last_block`. Physical body
1977 /// size differs by type: RLE bodies are always 1 byte (the repeated
1978 /// byte), Raw/Compressed bodies span `block_size`. `header_len` is the
1979 /// serialized frame-header length (frame offset of the first block).
1980 #[cfg(feature = "lsm")]
1981 fn populate_frame_emit_info(
1982 &mut self,
1983 header_len: usize,
1984 all_blocks: &[u8],
1985 emit_checksum: bool,
1986 ) {
1987 use crate::blocks::block::BlockType as BT;
1988 use crate::encoding::frame_emit_info::{FrameBlock, FrameEmitInfo};
1989 // All frame-offset arithmetic below is bounded by u32 on the wire
1990 // (Block_Size is a 21-bit field, frames bounded by MAX_BLOCK_SIZE *
1991 // #blocks). A pathologically large frame whose total emitted size
1992 // exceeds u32::MAX would overflow the cast; bail out by leaving
1993 // `frame_emit_info` at `None` rather than handing the caller a
1994 // silently-truncated layout. The overflow path is statically
1995 // unreachable on every realistic frame so the predictor amortises
1996 // the branch to zero cost.
1997 let frame_header_len: u32 = match u32::try_from(header_len) {
1998 Ok(v) => v,
1999 Err(_) => return,
2000 };
2001 let all_blocks_len_u32: u32 = match u32::try_from(all_blocks.len()) {
2002 Ok(v) => v,
2003 Err(_) => return,
2004 };
2005 let mut blocks: Vec<FrameBlock> = Vec::new();
2006 let mut cursor: usize = 0;
2007 while cursor + 3 <= all_blocks.len() {
2008 let mut header_u32 = [0u8; 4];
2009 header_u32[..3].copy_from_slice(&all_blocks[cursor..cursor + 3]);
2010 let raw = u32::from_le_bytes(header_u32);
2011 let last_block = (raw & 1) != 0;
2012 let block_type = match (raw >> 1) & 0b11 {
2013 0 => BT::Raw,
2014 1 => BT::RLE,
2015 2 => BT::Compressed,
2016 _ => BT::Reserved,
2017 };
2018 let block_size_field = raw >> 3;
2019 // RLE bodies are always 1 byte physical on the wire (the single
2020 // repeated byte); the spec's Block_Size field carries the
2021 // logical repeat count. Raw and Compressed bodies physically
2022 // span block_size_field bytes. Store the physical length in
2023 // body_size so the 'offset + header + body_size' arithmetic
2024 // always lands on the next block boundary, and surface the raw
2025 // spec field separately as block_size_field.
2026 let physical_body: u32 = match block_type {
2027 BT::RLE => 1,
2028 _ => block_size_field,
2029 };
2030 let cursor_u32: u32 = match u32::try_from(cursor) {
2031 Ok(v) => v,
2032 Err(_) => return,
2033 };
2034 let offset_in_frame = match frame_header_len.checked_add(cursor_u32) {
2035 Some(v) => v,
2036 None => return,
2037 };
2038 // Decompressed (regenerated) size, captured per physical block
2039 // during emit (1:1 with the wire blocks scanned here). Raw/RLE are
2040 // wire-derivable (`block_size_field`), so a short sidecar still
2041 // yields the correct value for them. A Compressed block's size is
2042 // NOT on the wire: if the sidecar is missing its entry, fabricating
2043 // 0 would publish a silently-wrong `decompressed_byte_range`. Since
2044 // this metadata is the authoritative mapping for a successful
2045 // encode, bail out (leave `frame_emit_info` at `None`) rather than
2046 // hand back a corrupt layout; the 1:1 push invariant makes this
2047 // unreachable in practice (debug_assert catches a regression).
2048 let decompressed_size = match self.block_decompressed_sizes.get(blocks.len()).copied() {
2049 Some(size) => size,
2050 None if matches!(block_type, BT::Raw | BT::RLE) => block_size_field,
2051 None => {
2052 debug_assert!(
2053 false,
2054 "missing decompressed-size sidecar entry for compressed block {}",
2055 blocks.len()
2056 );
2057 return;
2058 }
2059 };
2060 blocks.push(FrameBlock {
2061 offset_in_frame,
2062 header_size: 3,
2063 body_size: physical_body,
2064 block_size_field,
2065 block_type,
2066 last_block,
2067 decompressed_size,
2068 });
2069 cursor += 3 + physical_body as usize;
2070 if last_block {
2071 break;
2072 }
2073 }
2074 // Fail closed on a structurally incomplete scan: the loop must have
2075 // consumed the whole block section AND ended on a parsed last block.
2076 // A premature `last_block` (bytes left over) or a run-off without any
2077 // last block would otherwise publish an invalid public `FrameEmitInfo`.
2078 // Unreachable for a well-formed self-produced frame (debug_assert
2079 // catches a regression); on release we bail, leaving `frame_emit_info`
2080 // at `None` rather than handing back a corrupt layout.
2081 if cursor != all_blocks.len() || !blocks.last().is_some_and(|b| b.last_block) {
2082 debug_assert!(
2083 false,
2084 "incomplete block scan in populate_frame_emit_info: cursor={} len={} last_block={:?}",
2085 cursor,
2086 all_blocks.len(),
2087 blocks.last().map(|b| b.last_block)
2088 );
2089 return;
2090 }
2091 let checksum_range = if emit_checksum {
2092 let cs_start = match frame_header_len.checked_add(all_blocks_len_u32) {
2093 Some(v) => v,
2094 None => return,
2095 };
2096 let cs_end = match cs_start.checked_add(4) {
2097 Some(v) => v,
2098 None => return,
2099 };
2100 Some(cs_start..cs_end)
2101 } else {
2102 None
2103 };
2104 let body_total = match frame_header_len.checked_add(all_blocks_len_u32) {
2105 Some(v) => v,
2106 None => return,
2107 };
2108 let total_size = if checksum_range.is_some() {
2109 match body_total.checked_add(4) {
2110 Some(v) => v,
2111 None => return,
2112 }
2113 } else {
2114 body_total
2115 };
2116 self.frame_emit_info = Some(FrameEmitInfo {
2117 frame_header_range: 0..frame_header_len,
2118 blocks,
2119 checksum_range,
2120 total_size,
2121 });
2122 }
2123
2124 /// Layout of the most recently emitted frame.
2125 ///
2126 /// Returns `None` if [`compress`](Self::compress) has not been
2127 /// called yet on this compressor. After a successful `compress()`
2128 /// the returned `FrameEmitInfo` describes the frame header range,
2129 /// every emitted block's offset / size / type, and the optional
2130 /// trailing content-checksum range — all in frame-absolute byte
2131 /// offsets matching the bytes written to the drain.
2132 ///
2133 /// Behind the `lsm` Cargo feature.
2134 #[cfg(feature = "lsm")]
2135 pub fn last_frame_emit_info(&self) -> Option<&crate::encoding::frame_emit_info::FrameEmitInfo> {
2136 self.frame_emit_info.as_ref()
2137 }
2138
2139 /// Opt in to per-block XXH64 checksum computation during
2140 /// [`compress`](Self::compress). Default off; zero cost when
2141 /// disabled. The captured digests are accessible via
2142 /// [`last_frame_block_checksums`](Self::last_frame_block_checksums).
2143 ///
2144 /// One checksum is emitted per physical FrameBlock written to
2145 /// the drain: 1:1 cardinality with
2146 /// [`last_frame_emit_info`](Self::last_frame_emit_info)'s
2147 /// `blocks` vector. On the post-split optimization path
2148 /// (Level 16-22 with large window) the per-partition decompressed
2149 /// range is hashed inside the partition loop so the digest count
2150 /// still matches the emitted block count. The decoder collects
2151 /// per-physical-block digests on the same granularity, so
2152 /// element-wise equality holds round-trip.
2153 ///
2154 /// Behind `all(feature = "lsm", feature = "hash")` — the XXH64
2155 /// primitive lives behind the `hash` feature, so this method only
2156 /// compiles when both are enabled.
2157 #[cfg(all(feature = "lsm", feature = "hash"))]
2158 pub fn enable_per_block_checksums(&mut self) {
2159 self.per_block_checksums_enabled = true;
2160 }
2161
2162 /// Per-block XXH64 (low 32 bits) digests captured during the most
2163 /// recent `compress()` call. `None` unless
2164 /// [`enable_per_block_checksums`](Self::enable_per_block_checksums)
2165 /// was called before `compress()`.
2166 ///
2167 /// Behind `all(feature = "lsm", feature = "hash")`.
2168 #[cfg(all(feature = "lsm", feature = "hash"))]
2169 pub fn last_frame_block_checksums(&self) -> Option<&[u32]> {
2170 self.block_checksums.as_deref()
2171 }
2172
2173 /// Get a mutable reference to the source
2174 pub fn source_mut(&mut self) -> Option<&mut R> {
2175 self.uncompressed_data.as_mut()
2176 }
2177
2178 /// Get a mutable reference to the drain
2179 pub fn drain_mut(&mut self) -> Option<&mut W> {
2180 self.compressed_data.as_mut()
2181 }
2182
2183 /// Get a reference to the source
2184 pub fn source(&self) -> Option<&R> {
2185 self.uncompressed_data.as_ref()
2186 }
2187
2188 /// Get a reference to the drain
2189 pub fn drain(&self) -> Option<&W> {
2190 self.compressed_data.as_ref()
2191 }
2192
2193 /// Retrieve the source
2194 pub fn take_source(&mut self) -> Option<R> {
2195 self.uncompressed_data.take()
2196 }
2197
2198 /// Retrieve the drain
2199 pub fn take_drain(&mut self) -> Option<W> {
2200 self.compressed_data.take()
2201 }
2202
2203 /// Before calling [FrameCompressor::compress] you can replace the matcher
2204 pub fn replace_matcher(&mut self, mut match_generator: M) -> M {
2205 core::mem::swap(&mut match_generator, &mut self.state.matcher);
2206 match_generator
2207 }
2208
2209 /// Before calling [FrameCompressor::compress] you can replace the compression level.
2210 ///
2211 /// This also clears any fine-grained parameter overrides installed via
2212 /// [`set_parameters`](Self::set_parameters): reverting to a bare level
2213 /// means plain level-based tuning, not the previous frame's customized
2214 /// strategy / LDM / log overrides. To keep overriding, call
2215 /// [`set_parameters`](Self::set_parameters) again with the new base level.
2216 pub fn set_compression_level(
2217 &mut self,
2218 compression_level: CompressionLevel,
2219 ) -> CompressionLevel {
2220 let old = self.compression_level;
2221 self.compression_level = compression_level;
2222 // Drop sticky overrides so the level switch yields plain geometry.
2223 self.strategy_override = None;
2224 self.state.matcher.clear_param_overrides();
2225 old
2226 }
2227
2228 /// Get the current compression level
2229 pub fn compression_level(&self) -> CompressionLevel {
2230 self.compression_level
2231 }
2232
2233 /// Attach a pre-parsed dictionary to be used for subsequent compressions.
2234 ///
2235 /// In compressed modes, the dictionary id is written only when the active
2236 /// matcher supports dictionary priming.
2237 /// Uncompressed mode and non-priming matchers ignore the attached dictionary
2238 /// at encode time.
2239 pub fn set_dictionary(
2240 &mut self,
2241 dictionary: crate::decoding::Dictionary,
2242 ) -> Result<Option<EncoderDictionary>, crate::decoding::errors::DictionaryDecodeError> {
2243 self.attach_dictionary(EncoderDictionary::from_dictionary(dictionary))
2244 }
2245
2246 /// Parse and attach a serialized dictionary blob.
2247 ///
2248 /// Parses with the encoder-only path (skips the FSE/HUF decode lookup-table
2249 /// build the encoder never reads); the entropy ENCODER tables — and thus
2250 /// the emitted frame — are identical to a full parse.
2251 pub fn set_dictionary_from_bytes(
2252 &mut self,
2253 raw_dictionary: &[u8],
2254 ) -> Result<Option<EncoderDictionary>, crate::decoding::errors::DictionaryDecodeError> {
2255 self.attach_dictionary(EncoderDictionary::from_bytes(raw_dictionary)?)
2256 }
2257
2258 /// Attach an already-parsed [`EncoderDictionary`] without reparsing a raw
2259 /// blob.
2260 ///
2261 /// Accepts an `EncoderDictionary` produced once via
2262 /// [`EncoderDictionary::from_bytes`] / [`EncoderDictionary::from_dictionary`]
2263 /// or handed back by [`Self::clear_dictionary`] / the `set_dictionary*`
2264 /// return value, so callers can reattach or reuse a prepared dictionary
2265 /// across compressions without re-running the dictionary parse each time.
2266 /// Returns the previously-attached dictionary, if any.
2267 pub fn set_encoder_dictionary(
2268 &mut self,
2269 dictionary: EncoderDictionary,
2270 ) -> Result<Option<EncoderDictionary>, crate::decoding::errors::DictionaryDecodeError> {
2271 self.attach_dictionary(dictionary)
2272 }
2273
2274 /// Remove the attached dictionary, returning it as an [`EncoderDictionary`].
2275 pub fn clear_dictionary(&mut self) -> Option<EncoderDictionary> {
2276 self.dictionary_entropy_cache = None;
2277 // Drop the CDict prime snapshot — it is keyed to the dictionary
2278 // being removed and must not be restored against a different (or no)
2279 // dictionary on the next frame.
2280 self.state.matcher.invalidate_primed_dictionary();
2281 self.dictionary.take()
2282 }
2283
2284 /// Validate `enc`, build the encoder entropy cache from it, store it, and
2285 /// return the previously-attached dictionary. Shared by every public
2286 /// attach entry point: `set_dictionary`, `set_dictionary_from_bytes`, and
2287 /// `set_encoder_dictionary`.
2288 fn attach_dictionary(
2289 &mut self,
2290 enc: EncoderDictionary,
2291 ) -> Result<Option<EncoderDictionary>, crate::decoding::errors::DictionaryDecodeError> {
2292 let dictionary = &enc.inner;
2293 if dictionary.id == 0 {
2294 return Err(crate::decoding::errors::DictionaryDecodeError::ZeroDictionaryId);
2295 }
2296 if let Some(index) = dictionary.offset_hist.iter().position(|&rep| rep == 0) {
2297 return Err(
2298 crate::decoding::errors::DictionaryDecodeError::ZeroRepeatOffsetInDictionary {
2299 index: index as u8,
2300 },
2301 );
2302 }
2303 self.dictionary_entropy_cache = Some(CachedDictionaryEntropy::from_dictionary(dictionary));
2304 // A previously-captured CDict prime snapshot belongs to the OLD
2305 // dictionary; drop it so the first frame with the new dictionary
2306 // re-primes (and re-captures) instead of restoring stale tables.
2307 self.state.matcher.invalidate_primed_dictionary();
2308 Ok(self.dictionary.replace(enc))
2309 }
2310}
2311
2312#[cfg(test)]
2313mod tests {
2314 // `format!` is used by ungated tests (e.g. the btlazy2 dict-reuse
2315 // byte-identity test), so the import must not be feature-gated — under
2316 // default features (no `dict_builder`) the gated form left `format!`
2317 // unresolved when the test module is compiled.
2318 use alloc::format;
2319 use alloc::vec;
2320
2321 use super::FrameCompressor;
2322 use crate::common::{MAGIC_NUM, MAX_BLOCK_SIZE};
2323 use crate::decoding::FrameDecoder;
2324 use crate::encoding::{Matcher, Sequence};
2325 use alloc::vec::Vec;
2326
2327 fn generate_data(seed: u64, len: usize) -> Vec<u8> {
2328 let mut state = seed;
2329 let mut data = Vec::with_capacity(len);
2330 for _ in 0..len {
2331 state = state
2332 .wrapping_mul(6364136223846793005)
2333 .wrapping_add(1442695040888963407);
2334 data.push((state >> 33) as u8);
2335 }
2336 data
2337 }
2338
2339 // Cross-implementation parity tests (compress here, decode through the C
2340 // bindings) moved to `ffi-bench/tests/frame_compressor_ffi.rs` so the
2341 // library crate never links libzstd.
2342
2343 struct NoDictionaryMatcher {
2344 last_space: Vec<u8>,
2345 window_size: u64,
2346 }
2347
2348 impl NoDictionaryMatcher {
2349 fn new(window_size: u64) -> Self {
2350 Self {
2351 last_space: Vec::new(),
2352 window_size,
2353 }
2354 }
2355 }
2356
2357 impl Matcher for NoDictionaryMatcher {
2358 fn get_next_space(&mut self) -> Vec<u8> {
2359 vec![0; self.window_size as usize]
2360 }
2361
2362 fn get_last_space(&mut self) -> &[u8] {
2363 self.last_space.as_slice()
2364 }
2365
2366 fn commit_space(&mut self, space: Vec<u8>) {
2367 self.last_space = space;
2368 }
2369
2370 fn skip_matching(&mut self) {}
2371
2372 fn start_matching(&mut self, mut handle_sequence: impl for<'a> FnMut(Sequence<'a>)) {
2373 handle_sequence(Sequence::Literals {
2374 literals: self.last_space.as_slice(),
2375 });
2376 }
2377
2378 fn reset(&mut self, _level: super::CompressionLevel) {
2379 self.last_space.clear();
2380 }
2381
2382 fn window_size(&self) -> u64 {
2383 self.window_size
2384 }
2385 }
2386
2387 #[test]
2388 fn frame_starts_with_magic_num() {
2389 let mock_data = [1_u8, 2, 3].as_slice();
2390 let mut output: Vec<u8> = Vec::new();
2391 let mut compressor = FrameCompressor::new(super::CompressionLevel::Uncompressed);
2392 compressor.set_source(mock_data);
2393 compressor.set_drain(&mut output);
2394
2395 compressor.compress();
2396 assert!(output.starts_with(&MAGIC_NUM.to_le_bytes()));
2397 }
2398
2399 #[test]
2400 fn very_simple_raw_compress() {
2401 let mock_data = [1_u8, 2, 3].as_slice();
2402 let mut output: Vec<u8> = Vec::new();
2403 let mut compressor = FrameCompressor::new(super::CompressionLevel::Uncompressed);
2404 compressor.set_source(mock_data);
2405 compressor.set_drain(&mut output);
2406
2407 compressor.compress();
2408 }
2409
2410 #[test]
2411 fn very_simple_compress() {
2412 let mut mock_data = vec![0; 1 << 17];
2413 mock_data.extend(vec![1; (1 << 17) - 1]);
2414 mock_data.extend(vec![2; (1 << 18) - 1]);
2415 mock_data.extend(vec![2; 1 << 17]);
2416 mock_data.extend(vec![3; (1 << 17) - 1]);
2417 let mut output: Vec<u8> = Vec::new();
2418 let mut compressor = FrameCompressor::new(super::CompressionLevel::Uncompressed);
2419 compressor.set_source(mock_data.as_slice());
2420 compressor.set_drain(&mut output);
2421
2422 compressor.compress();
2423
2424 let mut decoder = FrameDecoder::new();
2425 let mut decoded = Vec::with_capacity(mock_data.len());
2426 decoder.decode_all_to_vec(&output, &mut decoded).unwrap();
2427 assert_eq!(mock_data, decoded);
2428 }
2429
2430 #[test]
2431 fn rle_compress() {
2432 let mock_data = vec![0; 1 << 19];
2433 let mut output: Vec<u8> = Vec::new();
2434 let mut compressor = FrameCompressor::new(super::CompressionLevel::Uncompressed);
2435 compressor.set_source(mock_data.as_slice());
2436 compressor.set_drain(&mut output);
2437
2438 compressor.compress();
2439
2440 let mut decoder = FrameDecoder::new();
2441 let mut decoded = Vec::with_capacity(mock_data.len());
2442 decoder.decode_all_to_vec(&output, &mut decoded).unwrap();
2443 assert_eq!(mock_data, decoded);
2444 }
2445
2446 #[test]
2447 fn aaa_compress() {
2448 let mock_data = vec![0, 1, 3, 4, 5];
2449 let mut output: Vec<u8> = Vec::new();
2450 let mut compressor = FrameCompressor::new(super::CompressionLevel::Uncompressed);
2451 compressor.set_source(mock_data.as_slice());
2452 compressor.set_drain(&mut output);
2453
2454 compressor.compress();
2455
2456 let mut decoder = FrameDecoder::new();
2457 let mut decoded = Vec::with_capacity(mock_data.len());
2458 decoder.decode_all_to_vec(&output, &mut decoded).unwrap();
2459 assert_eq!(mock_data, decoded);
2460 }
2461
2462 #[test]
2463 fn dictionary_compression_sets_required_dict_id_and_roundtrips() {
2464 let dict_raw = include_bytes!("../../dict_tests/dictionary");
2465 let dict_for_encoder = crate::decoding::Dictionary::decode_dict(dict_raw).unwrap();
2466 let dict_for_decoder = crate::decoding::Dictionary::decode_dict(dict_raw).unwrap();
2467
2468 let mut data = Vec::new();
2469 for _ in 0..8 {
2470 data.extend_from_slice(&dict_for_decoder.dict_content[..2048]);
2471 }
2472
2473 let mut with_dict = Vec::new();
2474 let mut compressor = FrameCompressor::new(super::CompressionLevel::Fastest);
2475 let previous = compressor
2476 .set_dictionary_from_bytes(dict_raw)
2477 .expect("dictionary bytes should parse");
2478 assert!(
2479 previous.is_none(),
2480 "first dictionary insert should return None"
2481 );
2482 assert_eq!(
2483 compressor
2484 .set_dictionary(dict_for_encoder)
2485 .expect("valid dictionary should attach")
2486 .expect("set_dictionary_from_bytes inserted previous dictionary")
2487 .id(),
2488 dict_for_decoder.id
2489 );
2490 compressor.set_source(data.as_slice());
2491 compressor.set_drain(&mut with_dict);
2492 compressor.compress();
2493
2494 let (frame_header, _) = crate::decoding::frame::read_frame_header(with_dict.as_slice())
2495 .expect("encoded stream should have a frame header");
2496 assert_eq!(frame_header.dictionary_id(), Some(dict_for_decoder.id));
2497
2498 let mut decoder = FrameDecoder::new();
2499 let mut missing_dict_target = Vec::with_capacity(data.len());
2500 let err = decoder
2501 .decode_all_to_vec(&with_dict, &mut missing_dict_target)
2502 .unwrap_err();
2503 assert!(
2504 matches!(
2505 &err,
2506 crate::decoding::errors::FrameDecoderError::DictNotProvided { .. }
2507 ),
2508 "dict-compressed stream should require dictionary id, got: {err:?}"
2509 );
2510
2511 let mut decoder = FrameDecoder::new();
2512 decoder.add_dict(dict_for_decoder).unwrap();
2513 let mut decoded = Vec::with_capacity(data.len());
2514 decoder.decode_all_to_vec(&with_dict, &mut decoded).unwrap();
2515 assert_eq!(decoded, data);
2516 }
2517
2518 #[cfg(all(feature = "dict_builder", feature = "std"))]
2519 #[test]
2520 fn dictionary_compression_roundtrips_with_dict_builder_dictionary() {
2521 use std::io::Cursor;
2522
2523 let mut training = Vec::new();
2524 for idx in 0..256u32 {
2525 training.extend_from_slice(
2526 format!("tenant=demo table=orders key={idx} region=eu\n").as_bytes(),
2527 );
2528 }
2529 let mut raw_dict = Vec::new();
2530 crate::dictionary::create_raw_dict_from_source(
2531 Cursor::new(training.as_slice()),
2532 training.len(),
2533 &mut raw_dict,
2534 4096,
2535 )
2536 .expect("dict_builder training should succeed");
2537 assert!(
2538 !raw_dict.is_empty(),
2539 "dict_builder produced an empty dictionary"
2540 );
2541
2542 let dict_id = 0xD1C7_0008;
2543 let encoder_dict =
2544 crate::decoding::Dictionary::from_raw_content(dict_id, raw_dict.clone()).unwrap();
2545 let decoder_dict =
2546 crate::decoding::Dictionary::from_raw_content(dict_id, raw_dict.clone()).unwrap();
2547
2548 let mut payload = Vec::new();
2549 for idx in 0..96u32 {
2550 payload.extend_from_slice(
2551 format!(
2552 "tenant=demo table=orders op=put key={idx} value=aaaaabbbbbcccccdddddeeeee\n"
2553 )
2554 .as_bytes(),
2555 );
2556 }
2557
2558 let mut without_dict = Vec::new();
2559 let mut baseline = FrameCompressor::new(super::CompressionLevel::Fastest);
2560 baseline.set_source(payload.as_slice());
2561 baseline.set_drain(&mut without_dict);
2562 baseline.compress();
2563
2564 let mut with_dict = Vec::new();
2565 let mut compressor = FrameCompressor::new(super::CompressionLevel::Fastest);
2566 compressor
2567 .set_dictionary(encoder_dict)
2568 .expect("valid dict_builder dictionary should attach");
2569 compressor.set_source(payload.as_slice());
2570 compressor.set_drain(&mut with_dict);
2571 compressor.compress();
2572
2573 let (frame_header, _) = crate::decoding::frame::read_frame_header(with_dict.as_slice())
2574 .expect("encoded stream should have a frame header");
2575 assert_eq!(frame_header.dictionary_id(), Some(dict_id));
2576 let mut decoder = FrameDecoder::new();
2577 decoder.add_dict(decoder_dict).unwrap();
2578 let mut decoded = Vec::with_capacity(payload.len());
2579 decoder.decode_all_to_vec(&with_dict, &mut decoded).unwrap();
2580 assert_eq!(decoded, payload);
2581 assert!(
2582 with_dict.len() < without_dict.len(),
2583 "trained dictionary should improve compression for this small payload"
2584 );
2585 }
2586
2587 #[test]
2588 fn set_dictionary_from_bytes_seeds_entropy_tables_for_first_block() {
2589 let dict_raw = include_bytes!("../../dict_tests/dictionary");
2590 let mut output = Vec::new();
2591 let input = b"";
2592
2593 let mut compressor = FrameCompressor::new(super::CompressionLevel::Fastest);
2594 let previous = compressor
2595 .set_dictionary_from_bytes(dict_raw)
2596 .expect("dictionary bytes should parse");
2597 assert!(previous.is_none());
2598
2599 compressor.set_source(input.as_slice());
2600 compressor.set_drain(&mut output);
2601 compressor.compress();
2602
2603 assert!(
2604 compressor.state.last_huff_table.is_some(),
2605 "dictionary entropy should seed previous huffman table before first block"
2606 );
2607 assert!(
2608 compressor.state.fse_tables.ll_previous.is_some(),
2609 "dictionary entropy should seed previous ll table before first block"
2610 );
2611 assert!(
2612 compressor.state.fse_tables.ml_previous.is_some(),
2613 "dictionary entropy should seed previous ml table before first block"
2614 );
2615 assert!(
2616 compressor.state.fse_tables.of_previous.is_some(),
2617 "dictionary entropy should seed previous of table before first block"
2618 );
2619 }
2620
2621 // `set_content_size_flag(false)`: the header must omit the FCS field
2622 // (and the single-segment layout that requires it) while the frame
2623 // still round-trips through our decoder.
2624 #[test]
2625 fn content_size_flag_off_omits_fcs_and_roundtrips() {
2626 let payload = alloc::vec![0x42u8; 4096];
2627
2628 let mut compressor: FrameCompressor =
2629 FrameCompressor::new(super::CompressionLevel::Fastest);
2630 let mut with_fcs = Vec::new();
2631 compressor.compress_independent_frame_into(&payload, &mut with_fcs);
2632
2633 compressor.set_content_size_flag(false);
2634 let mut without_fcs = Vec::new();
2635 compressor.compress_independent_frame_into(&payload, &mut without_fcs);
2636
2637 let parsed_with = crate::decoding::frame::read_frame_header(with_fcs.as_slice())
2638 .expect("flag-on frame header must parse")
2639 .0;
2640 assert_eq!(parsed_with.frame_content_size(), 4096);
2641
2642 let parsed_without = crate::decoding::frame::read_frame_header(without_fcs.as_slice())
2643 .expect("flag-off frame header must parse")
2644 .0;
2645 // 0 is the decoder's "unknown content size" sentinel...
2646 assert_eq!(
2647 parsed_without.frame_content_size(),
2648 0,
2649 "FCS must be omitted with the content-size flag off"
2650 );
2651 // ...and the descriptor must confirm the field is ABSENT (0 bytes),
2652 // not present with an explicit zero value.
2653 assert_eq!(
2654 parsed_without
2655 .descriptor
2656 .frame_content_size_bytes()
2657 .expect("descriptor must parse"),
2658 0,
2659 "the FCS field itself must be omitted, not written as zero"
2660 );
2661
2662 let mut decoder = crate::decoding::FrameDecoder::new();
2663 // `decode_all_to_vec` fills existing capacity (no FCS to pre-size
2664 // from with the flag off), so reserve the expected payload upfront.
2665 let mut decoded = Vec::with_capacity(payload.len() + 64);
2666 decoder
2667 .decode_all_to_vec(&without_fcs, &mut decoded)
2668 .expect("flag-off frame must decode");
2669 assert_eq!(decoded, payload);
2670 }
2671
2672 // `set_dictionary_id_flag(false)`: a dict-compressed frame must omit
2673 // the dictionary ID and still decode when the dictionary is handed to
2674 // the decoder explicitly.
2675 #[test]
2676 fn dict_id_flag_off_omits_dictionary_id_and_roundtrips() {
2677 let dict_raw = include_bytes!("../../dict_tests/dictionary");
2678 let payload = b"dictionary-keyed payload dictionary-keyed payload".repeat(8);
2679
2680 let mut compressor: FrameCompressor =
2681 FrameCompressor::new(super::CompressionLevel::Fastest);
2682 compressor
2683 .set_dictionary_from_bytes(dict_raw)
2684 .expect("dictionary bytes should parse");
2685 compressor.set_dictionary_id_flag(false);
2686 let mut frame = Vec::new();
2687 compressor.compress_independent_frame_into(&payload, &mut frame);
2688
2689 let parsed = crate::decoding::frame::read_frame_header(frame.as_slice())
2690 .expect("frame header must parse")
2691 .0;
2692 assert_eq!(
2693 parsed.dictionary_id(),
2694 None,
2695 "dictionary id must be omitted with the dict-id flag off"
2696 );
2697
2698 // With the ID omitted the decoder cannot look the dictionary up by
2699 // header; hand it explicitly (the `reset_with_dict_handle` path).
2700 let mut sd = crate::decoding::StreamingDecoder::new_with_dictionary_bytes(
2701 frame.as_slice(),
2702 dict_raw,
2703 )
2704 .expect("decoder must accept the dictionary");
2705 let mut dec = Vec::new();
2706 std::io::Read::read_to_end(&mut sd, &mut dec)
2707 .expect("frame must decode with the dictionary handed explicitly");
2708 assert_eq!(dec, payload);
2709 }
2710
2711 // The output reservation must track the observed compression ratio, not
2712 // the whole-input `compress_bound`: a multi-MiB compressible stream's
2713 // output buffer stays at output scale (the old up-front bound held an
2714 // input-sized allocation for the whole frame). Incompressible input may
2715 // still re-estimate to ~the full bound — that is the genuine worst case.
2716 #[test]
2717 fn compressible_stream_output_capacity_stays_at_output_scale() {
2718 // 4 MiB of highly repetitive log-like lines.
2719 let line = b"ts=2026-03-26T21:39:28Z level=INFO msg=\"flush memtable\" tenant=demo\n";
2720 let mut input = Vec::with_capacity(4 << 20);
2721 while input.len() < (4 << 20) {
2722 let take = line.len().min((4 << 20) - input.len());
2723 input.extend_from_slice(&line[..take]);
2724 }
2725
2726 let mut compressor: FrameCompressor =
2727 FrameCompressor::new(super::CompressionLevel::Fastest);
2728 let mut out = Vec::new();
2729 compressor.compress_independent_frame_into(&input, &mut out);
2730
2731 assert!(!out.is_empty());
2732 assert!(
2733 out.capacity() < input.len() / 4,
2734 "capacity {} must stay at output scale (input {}, output {})",
2735 out.capacity(),
2736 input.len(),
2737 out.len()
2738 );
2739
2740 // Round-trip: the adaptive reservation must not affect the bytes.
2741 let mut decoder = crate::decoding::FrameDecoder::new();
2742 let mut decoded = Vec::with_capacity(input.len() + 64);
2743 decoder
2744 .decode_all_to_vec(&out, &mut decoded)
2745 .expect("frame must decode");
2746 assert_eq!(decoded, input);
2747 }
2748
2749 // A dictionary frame with a known content size that fits the window
2750 // must take the single-segment layout (reference parity): the
2751 // dictionary is decoder setup state, not part of the regenerated
2752 // segment, so it must not force the windowed multi-segment layout.
2753 #[test]
2754 fn dict_frame_with_known_size_is_single_segment() {
2755 let dict_raw = include_bytes!("../../dict_tests/dictionary");
2756 let payload = b"dictionary-keyed payload dictionary-keyed payload".repeat(64);
2757
2758 let mut compressor: FrameCompressor =
2759 FrameCompressor::new(super::CompressionLevel::Fastest);
2760 compressor
2761 .set_dictionary_from_bytes(dict_raw)
2762 .expect("dictionary bytes should parse");
2763 let mut frame = Vec::new();
2764 compressor.compress_independent_frame_into(&payload, &mut frame);
2765
2766 let parsed = crate::decoding::frame::read_frame_header(frame.as_slice())
2767 .expect("frame header must parse")
2768 .0;
2769 assert!(
2770 parsed.descriptor.single_segment_flag(),
2771 "dict frame with known size <= window must be single-segment"
2772 );
2773 assert!(parsed.dictionary_id().is_some());
2774 assert_eq!(parsed.frame_content_size(), payload.len() as u64);
2775
2776 // Round-trip through our own decoder with the dictionary.
2777 let mut decoder = crate::decoding::FrameDecoder::new();
2778 decoder
2779 .add_dict_from_bytes(dict_raw)
2780 .expect("decoder must accept the dictionary");
2781 let mut decoded = Vec::with_capacity(payload.len() + 64);
2782 decoder
2783 .decode_all_to_vec(&frame, &mut decoded)
2784 .expect("single-segment dict frame must decode");
2785 assert_eq!(decoded, payload);
2786 }
2787
2788 // Regression test: `heap_size()` must count the retained Huffman tables
2789 // (the active `last_huff_table` and the recycled `huff_table_spare`).
2790 // A reused context that parks a table would otherwise under-report its
2791 // footprint through the public size API.
2792 #[test]
2793 fn heap_size_counts_active_and_spare_huffman_tables() {
2794 let mut compressor: FrameCompressor =
2795 FrameCompressor::new(super::CompressionLevel::Fastest);
2796 let base = compressor.heap_size();
2797
2798 let active = crate::huff0::huff0_encoder::HuffmanTable::build_from_data(
2799 b"abacabadabacabaeabacabadabacaba",
2800 );
2801 let active_bytes = active.heap_size();
2802 assert!(active_bytes > 0, "built table must own heap buffers");
2803 compressor.state.last_huff_table = Some(active);
2804 assert_eq!(
2805 compressor.heap_size(),
2806 base + active_bytes,
2807 "heap_size must include the active last_huff_table"
2808 );
2809
2810 let spare = crate::huff0::huff0_encoder::HuffmanTable::build_from_data(
2811 b"the quick brown fox jumps over the lazy dog",
2812 );
2813 let spare_bytes = spare.heap_size();
2814 assert!(spare_bytes > 0, "built table must own heap buffers");
2815 compressor.state.huff_table_spare = Some(spare);
2816 assert_eq!(
2817 compressor.heap_size(),
2818 base + active_bytes + spare_bytes,
2819 "heap_size must include the parked huff_table_spare"
2820 );
2821 }
2822
2823 #[test]
2824 fn set_encoder_dictionary_reattaches_prepared_dict_without_reparse() {
2825 let dict_raw = include_bytes!("../../dict_tests/dictionary");
2826 let payload = b"tenant=demo table=orders op=put key=1 value=aaaaabbbbbcccccdddddeeeee\n\
2827 tenant=demo table=orders op=put key=2 value=aaaaabbbbbcccccdddddeeeee\n";
2828
2829 // Prepare the EncoderDictionary once, then attach it via the prepared-
2830 // dictionary API (no raw-blob reparse at attach time).
2831 let prepared =
2832 super::EncoderDictionary::from_bytes(dict_raw).expect("dict bytes should parse");
2833 let dict_id = prepared.id();
2834
2835 let mut with_dict = Vec::new();
2836 let mut compressor = FrameCompressor::new(super::CompressionLevel::Fastest);
2837 let previous = compressor
2838 .set_encoder_dictionary(prepared)
2839 .expect("prepared dictionary should attach");
2840 assert!(previous.is_none());
2841 compressor.set_source(payload.as_slice());
2842 compressor.set_drain(&mut with_dict);
2843 compressor.compress();
2844 // clear_dictionary hands the prepared dictionary back (last use of
2845 // `compressor`, so its `&mut with_dict` drain borrow ends here).
2846 let returned = compressor
2847 .clear_dictionary()
2848 .expect("dictionary was attached");
2849 assert_eq!(returned.id(), dict_id);
2850
2851 // The reattached dictionary drives the frame: its id is advertised and
2852 // the stream round-trips through a decoder primed with the same dict.
2853 let (frame_header, _) = crate::decoding::frame::read_frame_header(with_dict.as_slice())
2854 .expect("encoded stream should have a frame header");
2855 assert_eq!(frame_header.dictionary_id(), Some(dict_id));
2856 let decoder_dict = crate::decoding::Dictionary::decode_dict(dict_raw).unwrap();
2857 let mut decoder = FrameDecoder::new();
2858 decoder.add_dict(decoder_dict).unwrap();
2859 let mut decoded = Vec::with_capacity(payload.len());
2860 decoder.decode_all_to_vec(&with_dict, &mut decoded).unwrap();
2861 assert_eq!(decoded.as_slice(), payload.as_slice());
2862
2863 // The dictionary handed back by clear_dictionary reattaches to another
2864 // compressor without touching the raw bytes again, producing an
2865 // identical frame.
2866 let mut with_dict2 = Vec::new();
2867 let mut compressor2 = FrameCompressor::new(super::CompressionLevel::Fastest);
2868 compressor2
2869 .set_encoder_dictionary(returned)
2870 .expect("returned dictionary should reattach");
2871 compressor2.set_source(payload.as_slice());
2872 compressor2.set_drain(&mut with_dict2);
2873 compressor2.compress();
2874 assert_eq!(
2875 with_dict2, with_dict,
2876 "reattached prepared dict must produce an identical frame"
2877 );
2878 }
2879
2880 #[test]
2881 fn dict_primed_matcher_snapshot_reused_across_frames_is_byte_identical() {
2882 // CDict-equivalent: a compressor reused across frames with the same
2883 // dictionary restores the primed matcher snapshot on frames 2..N
2884 // (a table copy) instead of re-hashing the dictionary. The restored
2885 // state must reproduce the first-frame (freshly-primed) output
2886 // byte-for-byte, and every frame must round-trip through a
2887 // dict-primed decoder.
2888 let dict_raw = include_bytes!("../../dict_tests/dictionary");
2889 // Source must exceed the Fast strategy's 8 KiB attach cutoff so the
2890 // copy-snapshot (restore) path is taken on frame 2 — at or below the
2891 // cutoff the upstream zstd attaches by reference and we fall back to re-prime,
2892 // which would not exercise restore.
2893 let mut payload = Vec::new();
2894 while payload.len() < 16 * 1024 {
2895 payload.extend_from_slice(
2896 b"tenant=demo table=orders op=put key=1 value=aaaaabbbbbcccccdddddeeeee\n",
2897 );
2898 }
2899
2900 let prepared =
2901 super::EncoderDictionary::from_bytes(dict_raw).expect("dict bytes should parse");
2902 let dict_id = prepared.id();
2903 let mut compressor: FrameCompressor =
2904 FrameCompressor::new(super::CompressionLevel::Fastest);
2905 compressor
2906 .set_encoder_dictionary(prepared)
2907 .expect("prepared dictionary should attach");
2908
2909 // Frame 1 primes + captures the snapshot; frame 2 restores it.
2910 let frame1 = compressor.compress_independent_frame(payload.as_slice());
2911 let frame2 = compressor.compress_independent_frame(payload.as_slice());
2912 assert_eq!(
2913 frame1, frame2,
2914 "restored prime snapshot must reproduce the freshly-primed frame byte-for-byte"
2915 );
2916
2917 // Both frames advertise the dict id and round-trip through a
2918 // dict-primed decoder.
2919 for frame in [&frame1, &frame2] {
2920 let (hdr, _) =
2921 crate::decoding::frame::read_frame_header(frame.as_slice()).expect("frame header");
2922 assert_eq!(hdr.dictionary_id(), Some(dict_id));
2923 let mut decoder = FrameDecoder::new();
2924 decoder
2925 .add_dict(crate::decoding::Dictionary::decode_dict(dict_raw).unwrap())
2926 .unwrap();
2927 let mut decoded = Vec::with_capacity(payload.len());
2928 decoder.decode_all_to_vec(frame, &mut decoded).unwrap();
2929 assert_eq!(decoded.as_slice(), payload.as_slice());
2930 }
2931 }
2932
2933 #[test]
2934 fn dict_primed_matcher_cache_reused_across_small_attach_frames_is_byte_identical() {
2935 // CDict-equivalent ATTACH path (small source, at/below the Fast 8 KiB
2936 // attach cutoff): frames 2..N re-prime — re-committing the dict bytes
2937 // to history — but reuse the already-built dict table instead of
2938 // re-hashing it. The cached-table frame must reproduce the
2939 // freshly-primed first frame byte-for-byte, and a fresh single-frame
2940 // compressor (no prior dict cache) must produce the identical bytes
2941 // too, proving the cache changes timing, not output.
2942 let dict_raw = include_bytes!("../../dict_tests/dictionary");
2943 // Stay under the 8 KiB cutoff so the attach (re-prime) path is taken
2944 // every frame rather than the copy-snapshot restore.
2945 let mut payload = Vec::new();
2946 while payload.len() < 2 * 1024 {
2947 payload.extend_from_slice(b"tenant=demo op=put key=1 value=aaaaabbbbbcccccddddd\n");
2948 }
2949
2950 let prepared =
2951 super::EncoderDictionary::from_bytes(dict_raw).expect("dict bytes should parse");
2952 let dict_id = prepared.id();
2953 let mut compressor: FrameCompressor =
2954 FrameCompressor::new(super::CompressionLevel::Fastest);
2955 compressor
2956 .set_encoder_dictionary(prepared)
2957 .expect("prepared dictionary should attach");
2958
2959 // Frame 1 builds + marks the dict table; frame 2 reuses it.
2960 let frame1 = compressor.compress_independent_frame(payload.as_slice());
2961 let frame2 = compressor.compress_independent_frame(payload.as_slice());
2962 assert_eq!(
2963 frame1, frame2,
2964 "reused dict table (attach path) must reproduce the freshly-built frame byte-for-byte"
2965 );
2966
2967 // A fresh compressor (cold dict cache) must emit the same bytes — the
2968 // cache is a timing optimization, never a content change.
2969 let fresh_prepared =
2970 super::EncoderDictionary::from_bytes(dict_raw).expect("dict bytes should parse");
2971 let mut fresh: FrameCompressor = FrameCompressor::new(super::CompressionLevel::Fastest);
2972 fresh
2973 .set_encoder_dictionary(fresh_prepared)
2974 .expect("prepared dictionary should attach");
2975 let fresh_frame = fresh.compress_independent_frame(payload.as_slice());
2976 assert_eq!(
2977 fresh_frame, frame1,
2978 "cold-cache compressor must match the warm-cache frame byte-for-byte"
2979 );
2980
2981 for frame in [&frame1, &frame2] {
2982 let (hdr, _) =
2983 crate::decoding::frame::read_frame_header(frame.as_slice()).expect("frame header");
2984 assert_eq!(hdr.dictionary_id(), Some(dict_id));
2985 let mut decoder = FrameDecoder::new();
2986 decoder
2987 .add_dict(crate::decoding::Dictionary::decode_dict(dict_raw).unwrap())
2988 .unwrap();
2989 let mut decoded = Vec::with_capacity(payload.len());
2990 decoder.decode_all_to_vec(frame, &mut decoded).unwrap();
2991 assert_eq!(decoded.as_slice(), payload.as_slice());
2992 }
2993 }
2994
2995 #[test]
2996 fn dict_fast_epoch_reset_many_frames_and_attach_copy_alternation_byte_identical() {
2997 // The Fast attach path invalidates the main hash table between
2998 // frames with an epoch-bias advance instead of a memset. Two things
2999 // need proving against a fresh-compressor reference:
3000 // 1. the bias accumulates across MANY reused frames without ever
3001 // letting a stale entry through (every frame byte-identical);
3002 // 2. crossing the 8 KiB attach/copy cutoff in both directions
3003 // (attach → copy clears the bias for the raw-slice kernel,
3004 // copy → attach re-enters epoch mode) stays byte-identical.
3005 let dict_raw = include_bytes!("../../dict_tests/dictionary");
3006 let mut small = Vec::new();
3007 while small.len() < 2 * 1024 {
3008 small.extend_from_slice(b"tenant=demo op=put key=1 value=aaaaabbbbbcccccddddd\n");
3009 }
3010 // Over the Fast 8 KiB attach cutoff → copy-mode frame.
3011 let mut large = Vec::new();
3012 while large.len() < 64 * 1024 {
3013 large.extend_from_slice(b"tenant=demo op=scan range=[k0,k9) limit=500 order=asc\n");
3014 }
3015
3016 let mut reused: FrameCompressor = FrameCompressor::new(super::CompressionLevel::Fastest);
3017 reused
3018 .set_encoder_dictionary(
3019 super::EncoderDictionary::from_bytes(dict_raw).expect("dict bytes should parse"),
3020 )
3021 .expect("prepared dictionary should attach");
3022
3023 let reference = |payload: &[u8]| -> alloc::vec::Vec<u8> {
3024 let mut fresh: FrameCompressor = FrameCompressor::new(super::CompressionLevel::Fastest);
3025 fresh
3026 .set_encoder_dictionary(
3027 super::EncoderDictionary::from_bytes(dict_raw)
3028 .expect("dict bytes should parse"),
3029 )
3030 .expect("prepared dictionary should attach");
3031 fresh.compress_independent_frame(payload)
3032 };
3033
3034 let small_expected = reference(&small);
3035 let large_expected = reference(&large);
3036
3037 // 1. Long attach-only run: every frame advances the epoch bias.
3038 for i in 0..32 {
3039 let frame = reused.compress_independent_frame(small.as_slice());
3040 assert_eq!(
3041 frame, small_expected,
3042 "attach frame {i} diverged from the fresh-compressor reference"
3043 );
3044 }
3045 // 2. Cutoff alternation: attach → copy → attach → copy.
3046 for i in 0..4 {
3047 let frame = reused.compress_independent_frame(large.as_slice());
3048 assert_eq!(
3049 frame, large_expected,
3050 "copy frame {i} diverged from the fresh-compressor reference"
3051 );
3052 let frame = reused.compress_independent_frame(small.as_slice());
3053 assert_eq!(
3054 frame, small_expected,
3055 "attach frame after copy {i} diverged from the fresh-compressor reference"
3056 );
3057 }
3058 }
3059
3060 #[test]
3061 fn dict_primed_btlazy2_reused_across_attach_and_copy_boundary_is_byte_identical() {
3062 // Btlazy2 (Level 15) uses the 32 KiB dict attach/copy cutoff in
3063 // prepare_frame. Exercise BOTH sides of that boundary on a reused
3064 // compressor: a sub-cutoff payload (re-prime/attach path) and an
3065 // over-cutoff payload (copy-snapshot restore path). In each case the
3066 // warm-cache second frame must reproduce the cold-cache first frame
3067 // byte-for-byte (the dict cache is a timing optimization, never a
3068 // content change), and every frame must round-trip.
3069 let dict_raw = include_bytes!("../../dict_tests/dictionary");
3070 let dict_id = super::EncoderDictionary::from_bytes(dict_raw)
3071 .expect("dict bytes should parse")
3072 .id();
3073 // Distinct lines so the payload does not trivially self-compress; the
3074 // BT finder + dict dual-probe both get exercised.
3075 let make_payload = |target: usize| {
3076 let mut p = Vec::with_capacity(target);
3077 let mut i = 0u64;
3078 while p.len() < target {
3079 p.extend_from_slice(
3080 format!(
3081 "tenant=demo op=put key={i} value=aaaaabbbbbcccccddddd-{}\n",
3082 i % 97
3083 )
3084 .as_bytes(),
3085 );
3086 i += 1;
3087 }
3088 p
3089 };
3090 // Below the 32 KiB cutoff (attach/re-prime) and above it (copy-snapshot).
3091 for target in [16 * 1024usize, 64 * 1024usize] {
3092 let payload = make_payload(target);
3093 let mut warm: FrameCompressor =
3094 FrameCompressor::new(super::CompressionLevel::Level(15));
3095 warm.set_encoder_dictionary(
3096 super::EncoderDictionary::from_bytes(dict_raw).expect("dict parse"),
3097 )
3098 .expect("dict attach");
3099 // Frame 1 builds + marks the dict tables; frame 2 reuses them.
3100 let frame1 = warm.compress_independent_frame(payload.as_slice());
3101 let frame2 = warm.compress_independent_frame(payload.as_slice());
3102 assert_eq!(
3103 frame1, frame2,
3104 "reused dict cache must reproduce the freshly-primed frame byte-for-byte \
3105 (Level 15, target={target})"
3106 );
3107 // Cold-cache compressor: must match the warm-cache bytes.
3108 let mut cold: FrameCompressor =
3109 FrameCompressor::new(super::CompressionLevel::Level(15));
3110 cold.set_encoder_dictionary(
3111 super::EncoderDictionary::from_bytes(dict_raw).expect("dict parse"),
3112 )
3113 .expect("dict attach");
3114 let cold_frame = cold.compress_independent_frame(payload.as_slice());
3115 assert_eq!(
3116 cold_frame, frame1,
3117 "cold-cache compressor must match warm-cache frame (Level 15, target={target})"
3118 );
3119 // Round-trip through a decoder primed with the same dict.
3120 for frame in [&frame1, &frame2] {
3121 let (hdr, _) = crate::decoding::frame::read_frame_header(frame.as_slice())
3122 .expect("frame header");
3123 assert_eq!(hdr.dictionary_id(), Some(dict_id));
3124 let mut decoder = FrameDecoder::new();
3125 decoder
3126 .add_dict(crate::decoding::Dictionary::decode_dict(dict_raw).unwrap())
3127 .unwrap();
3128 let mut decoded = Vec::with_capacity(payload.len());
3129 decoder.decode_all_to_vec(frame, &mut decoded).unwrap();
3130 assert_eq!(decoded.as_slice(), payload.as_slice());
3131 }
3132 }
3133 }
3134
3135 #[test]
3136 fn dict_primed_btultra2_restore_is_floor_safe_and_byte_identical() {
3137 // Regression guard for the dictionary primed-snapshot RESTORE path on
3138 // the binary-tree (btultra2 / Level 22) backend — the path a minimal /
3139 // decoupled prepared-dict refactor rewrites.
3140 //
3141 // The trap it pins: a reused compressor compresses frame A (which fills
3142 // the live hash/chain tables with frame-A positions and advances the
3143 // window floor), then frame B of the SAME resolved shape (same size →
3144 // same PrimedKey → the snapshot RESTORE path) but DIFFERENT content. The
3145 // restore must reinstate the clean post-prime dict state with NO live
3146 // frame-A entries surviving above the restored floor; a restore that
3147 // leaks stale frame-A positions would surface FALSE matches and produce
3148 // a different (or undecodable) frame B. The invariant: a snapshot
3149 // restore is a pure timing optimization and MUST be byte-identical to a
3150 // cold compressor compressing frame B from scratch, and must round-trip.
3151 let dict_raw = include_bytes!("../../dict_tests/dictionary");
3152 let dict_id = super::EncoderDictionary::from_bytes(dict_raw)
3153 .expect("dict bytes should parse")
3154 .id();
3155 // 48 KiB > the btultra2 8 KiB attach cutoff → the copy-snapshot
3156 // capture/restore path. Two distinct payloads of the SAME size so frame
3157 // B resolves to frame A's snapshot key and takes the restore path.
3158 let make_payload = |seed: u64, target: usize| {
3159 let mut p = Vec::with_capacity(target);
3160 let mut i = seed;
3161 while p.len() < target {
3162 p.extend_from_slice(
3163 format!(
3164 "tenant=demo op=put key={i} value=aaaaabbbbbcccccddddd-{}\n",
3165 i % 89
3166 )
3167 .as_bytes(),
3168 );
3169 i = i.wrapping_add(1);
3170 }
3171 p.truncate(target);
3172 p
3173 };
3174 let size = 48 * 1024usize;
3175 let frame_a = make_payload(0, size);
3176 let frame_b = make_payload(1_000_000, size);
3177
3178 let mut warm: FrameCompressor = FrameCompressor::new(super::CompressionLevel::Level(22));
3179 warm.set_encoder_dictionary(
3180 super::EncoderDictionary::from_bytes(dict_raw).expect("dict parse"),
3181 )
3182 .expect("dict attach");
3183 // Frame A: cold cache — primes the dict + captures the snapshot, and
3184 // fills the live tables with frame-A positions.
3185 let _wa = warm.compress_independent_frame(frame_a.as_slice());
3186 // Frame B: warm cache — takes the snapshot RESTORE path (same size).
3187 let warm_b = warm.compress_independent_frame(frame_b.as_slice());
3188
3189 // Cold compressor compressing frame B from scratch: the ground truth.
3190 let mut cold: FrameCompressor = FrameCompressor::new(super::CompressionLevel::Level(22));
3191 cold.set_encoder_dictionary(
3192 super::EncoderDictionary::from_bytes(dict_raw).expect("dict parse"),
3193 )
3194 .expect("dict attach");
3195 let cold_b = cold.compress_independent_frame(frame_b.as_slice());
3196
3197 assert_eq!(
3198 warm_b, cold_b,
3199 "frame B via snapshot restore must be byte-identical to a cold compress \
3200 (a restore that leaks frame-A live-table entries would diverge here)"
3201 );
3202
3203 // Round-trip frame B through a dict-primed decoder.
3204 let (hdr, _) =
3205 crate::decoding::frame::read_frame_header(warm_b.as_slice()).expect("frame header");
3206 assert_eq!(hdr.dictionary_id(), Some(dict_id));
3207 let mut decoder = FrameDecoder::new();
3208 decoder
3209 .add_dict(crate::decoding::Dictionary::decode_dict(dict_raw).unwrap())
3210 .unwrap();
3211 let mut decoded = Vec::with_capacity(frame_b.len());
3212 decoder
3213 .decode_all_to_vec(warm_b.as_slice(), &mut decoded)
3214 .unwrap();
3215 assert_eq!(decoded.as_slice(), frame_b.as_slice());
3216 }
3217
3218 #[test]
3219 fn dict_primed_btultra2_ldm_restore_is_byte_identical() {
3220 // Same restore-path byte-identity guard as
3221 // `dict_primed_btultra2_restore_is_floor_safe_and_byte_identical`, but
3222 // with long-distance matching ENABLED. The BtMatcher's LDM producer is
3223 // part of the snapshot; a refactor that decouples it (so the snapshot
3224 // does not retain the empty LDM table) must reinstate an equivalent
3225 // empty producer on restore. This pins that the warm-cache (restore)
3226 // frame stays byte-identical to a cold compress when LDM is on.
3227 let dict_raw = include_bytes!("../../dict_tests/dictionary");
3228 let dict_id = super::EncoderDictionary::from_bytes(dict_raw)
3229 .expect("dict bytes should parse")
3230 .id();
3231 let make_payload = |seed: u64, target: usize| {
3232 let mut p = Vec::with_capacity(target);
3233 let mut i = seed;
3234 while p.len() < target {
3235 p.extend_from_slice(
3236 format!(
3237 "tenant=demo op=put key={i} value=aaaaabbbbbcccccddddd-{}\n",
3238 i % 89
3239 )
3240 .as_bytes(),
3241 );
3242 i = i.wrapping_add(1);
3243 }
3244 p.truncate(target);
3245 p
3246 };
3247 let ldm_params =
3248 crate::encoding::CompressionParameters::builder(super::CompressionLevel::Level(22))
3249 .enable_long_distance_matching(true)
3250 .build()
3251 .expect("LDM-only params build");
3252 let size = 48 * 1024usize;
3253 let frame_a = make_payload(0, size);
3254 let frame_b = make_payload(1_000_000, size);
3255
3256 let mut warm: FrameCompressor = FrameCompressor::new(super::CompressionLevel::Level(22));
3257 warm.set_parameters(&ldm_params);
3258 warm.set_encoder_dictionary(
3259 super::EncoderDictionary::from_bytes(dict_raw).expect("dict parse"),
3260 )
3261 .expect("dict attach");
3262 let _wa = warm.compress_independent_frame(frame_a.as_slice());
3263 let warm_b = warm.compress_independent_frame(frame_b.as_slice());
3264
3265 let mut cold: FrameCompressor = FrameCompressor::new(super::CompressionLevel::Level(22));
3266 cold.set_parameters(&ldm_params);
3267 cold.set_encoder_dictionary(
3268 super::EncoderDictionary::from_bytes(dict_raw).expect("dict parse"),
3269 )
3270 .expect("dict attach");
3271 let cold_b = cold.compress_independent_frame(frame_b.as_slice());
3272
3273 assert_eq!(
3274 warm_b, cold_b,
3275 "LDM-on frame B via snapshot restore must be byte-identical to a cold compress"
3276 );
3277
3278 let (hdr, _) =
3279 crate::decoding::frame::read_frame_header(warm_b.as_slice()).expect("frame header");
3280 assert_eq!(hdr.dictionary_id(), Some(dict_id));
3281 let mut decoder = FrameDecoder::new();
3282 decoder
3283 .add_dict(crate::decoding::Dictionary::decode_dict(dict_raw).unwrap())
3284 .unwrap();
3285 let mut decoded = Vec::with_capacity(frame_b.len());
3286 decoder
3287 .decode_all_to_vec(warm_b.as_slice(), &mut decoded)
3288 .unwrap();
3289 assert_eq!(decoded.as_slice(), frame_b.as_slice());
3290 }
3291
3292 #[test]
3293 fn set_dictionary_from_bytes_matches_full_decode_byte_for_byte() {
3294 // The encoder-only dict parse (`decode_dict_for_encoding`, used by
3295 // `set_dictionary_from_bytes`) skips the FSE/HUF decoder-table build and
3296 // the enrich passes. The encoder entropy tables are derived purely from
3297 // the symbol probabilities / Huffman weights, so the compressed output
3298 // MUST be byte-identical to the full-decode path. This pins the
3299 // load-bearing equivalence so a future FSE/HUF parsing refactor that
3300 // still round-trips but silently diverges on the probabilities/weights
3301 // fails loudly here instead of producing a different (but valid) frame.
3302 let dict_raw = include_bytes!("../../dict_tests/dictionary");
3303 let payload = b"tenant=demo table=orders op=put key=1 value=aaaaabbbbbcccccdddddeeeee\n\
3304 tenant=demo table=orders op=put key=2 value=aaaaabbbbbcccccdddddeeeee\n";
3305
3306 // Path A: encoder-only parse straight from the raw blob.
3307 let mut from_bytes_out = Vec::new();
3308 {
3309 let mut compressor = FrameCompressor::new(super::CompressionLevel::Fastest);
3310 compressor
3311 .set_dictionary_from_bytes(dict_raw)
3312 .expect("dictionary bytes should parse");
3313 compressor.set_source(payload.as_slice());
3314 compressor.set_drain(&mut from_bytes_out);
3315 compressor.compress();
3316 }
3317
3318 // Path B: full decode (builds the decoder tables too), then attach for
3319 // encoding via the `Dictionary` setter.
3320 let full_decode = crate::decoding::Dictionary::decode_dict(dict_raw)
3321 .expect("dictionary bytes should fully decode");
3322 let mut full_decode_out = Vec::new();
3323 {
3324 let mut compressor = FrameCompressor::new(super::CompressionLevel::Fastest);
3325 compressor
3326 .set_dictionary(full_decode)
3327 .expect("full-decode dictionary should attach");
3328 compressor.set_source(payload.as_slice());
3329 compressor.set_drain(&mut full_decode_out);
3330 compressor.compress();
3331 }
3332
3333 assert_eq!(
3334 from_bytes_out, full_decode_out,
3335 "encoder-only dict parse must produce byte-identical output to the full decode"
3336 );
3337 }
3338
3339 #[test]
3340 fn set_dictionary_rejects_zero_dictionary_id() {
3341 let invalid = crate::decoding::Dictionary {
3342 id: 0,
3343 fse: crate::decoding::scratch::FSEScratch::new(),
3344 huf: crate::decoding::scratch::HuffmanScratch::new(),
3345 dict_content: vec![1, 2, 3],
3346 offset_hist: [1, 4, 8],
3347 };
3348
3349 let mut compressor: FrameCompressor<
3350 &[u8],
3351 Vec<u8>,
3352 crate::encoding::match_generator::MatchGeneratorDriver,
3353 > = FrameCompressor::new(super::CompressionLevel::Fastest);
3354 let result = compressor.set_dictionary(invalid);
3355 assert!(matches!(
3356 result,
3357 Err(crate::decoding::errors::DictionaryDecodeError::ZeroDictionaryId)
3358 ));
3359 }
3360
3361 #[test]
3362 fn set_dictionary_rejects_zero_repeat_offsets() {
3363 let invalid = crate::decoding::Dictionary {
3364 id: 1,
3365 fse: crate::decoding::scratch::FSEScratch::new(),
3366 huf: crate::decoding::scratch::HuffmanScratch::new(),
3367 dict_content: vec![1, 2, 3],
3368 offset_hist: [0, 4, 8],
3369 };
3370
3371 let mut compressor: FrameCompressor<
3372 &[u8],
3373 Vec<u8>,
3374 crate::encoding::match_generator::MatchGeneratorDriver,
3375 > = FrameCompressor::new(super::CompressionLevel::Fastest);
3376 let result = compressor.set_dictionary(invalid);
3377 assert!(matches!(
3378 result,
3379 Err(
3380 crate::decoding::errors::DictionaryDecodeError::ZeroRepeatOffsetInDictionary {
3381 index: 0
3382 }
3383 )
3384 ));
3385 }
3386
3387 #[test]
3388 fn uncompressed_mode_does_not_require_dictionary() {
3389 let dict_id = 0xABCD_0001;
3390 let dict =
3391 crate::decoding::Dictionary::from_raw_content(dict_id, b"shared-history".to_vec())
3392 .expect("raw dictionary should be valid");
3393
3394 let payload = b"plain-bytes-that-should-stay-raw";
3395 let mut output = Vec::new();
3396 let mut compressor = FrameCompressor::new(super::CompressionLevel::Uncompressed);
3397 compressor
3398 .set_dictionary(dict)
3399 .expect("dictionary should attach in uncompressed mode");
3400 compressor.set_source(payload.as_slice());
3401 compressor.set_drain(&mut output);
3402 compressor.compress();
3403
3404 let (frame_header, _) = crate::decoding::frame::read_frame_header(output.as_slice())
3405 .expect("encoded frame should have a header");
3406 assert_eq!(
3407 frame_header.dictionary_id(),
3408 None,
3409 "raw/uncompressed frames must not advertise dictionary dependency"
3410 );
3411
3412 let mut decoder = FrameDecoder::new();
3413 let mut decoded = Vec::with_capacity(payload.len());
3414 decoder.decode_all_to_vec(&output, &mut decoded).unwrap();
3415 assert_eq!(decoded, payload);
3416 }
3417
3418 #[test]
3419 fn default_level_tiny_raw_dict_compresses_cleanly() {
3420 // Coverage for the dfast dict-attach fast path with a
3421 // sub-min-match raw-content dictionary: the dict-table probe in
3422 // `start_matching_fast_loop` is gated on the dict table actually
3423 // existing (`table().is_some()`), not merely on `is_attached()`,
3424 // so a dictionary whose hashable region is shorter than the
3425 // short-hash lookahead (where `prime_dict_tables_for_range`
3426 // returns before allocating the tables) never dereferences a
3427 // null dict pointer. Compressing at the default (dfast) level
3428 // with such a dict must succeed.
3429 let dict_id = 0xABCD_0009;
3430 let dict = crate::decoding::Dictionary::from_raw_content(dict_id, b"abc".to_vec())
3431 .expect("raw dictionary should be valid");
3432 let payload = b"the quick brown fox jumps over the lazy dog, repeatedly and at length";
3433 let mut output = Vec::new();
3434 let mut compressor = FrameCompressor::new(super::CompressionLevel::Default);
3435 compressor
3436 .set_dictionary(dict)
3437 .expect("tiny raw dictionary should attach");
3438 compressor.set_source(payload.as_slice());
3439 compressor.set_drain(&mut output);
3440 compressor.compress();
3441 assert!(!output.is_empty(), "compression should produce a frame");
3442
3443 // The emitted frame must advertise the attached dictionary id, proving
3444 // the tiny-dict path stayed active (the payload round-trips either way,
3445 // so without this the test would also pass on a silent no-dict frame).
3446 let (frame_header, _) = crate::decoding::frame::read_frame_header(output.as_slice())
3447 .expect("encoded frame should have a readable header");
3448 assert_eq!(
3449 frame_header.dictionary_id(),
3450 Some(dict_id),
3451 "tiny raw dict frame should still advertise its dictionary id",
3452 );
3453
3454 // Full roundtrip: decode the dict-compressed frame with the SAME
3455 // dictionary attached and confirm byte-exact recovery — proves the
3456 // tiny-dict fast path produces a correct frame, not just a non-empty
3457 // one.
3458 let decode_dict = crate::decoding::Dictionary::from_raw_content(dict_id, b"abc".to_vec())
3459 .expect("raw dictionary should be valid");
3460 let mut decoder = FrameDecoder::new();
3461 decoder
3462 .add_dict(decode_dict)
3463 .expect("decoder dict should attach");
3464 let mut decoded = Vec::with_capacity(payload.len());
3465 decoder
3466 .decode_all_to_vec(&output, &mut decoded)
3467 .expect("dict roundtrip should decode");
3468 assert_eq!(decoded, payload, "tiny-dict roundtrip mismatch");
3469 }
3470
3471 /// Exercises the dictionary dual-probe (live + immutable dict tables)
3472 /// in the Fast / dfast / Row match finders with a dict whose content
3473 /// the payload actually reuses, so each backend's dict long/short
3474 /// probe (and the dfast `ip+1` dict-long retry) is reached and the
3475 /// dict-compressed frame round-trips through a decoder primed with the
3476 /// same dict. The 3-byte-dict test above only proves the null-table
3477 /// guard; this proves the full attach path produces correct frames.
3478 #[test]
3479 fn dict_attach_roundtrips_across_backends_with_matching_payload() {
3480 let dict_id = 0xD1C7_0001;
3481 // Distinct lines so the payload does NOT self-compress: each line
3482 // appears exactly once in the payload, so without the dictionary there
3483 // are no in-frame back-references to exploit. The dictionary holds the
3484 // SAME lines, so the only way the output shrinks is if the dict probe
3485 // actually fires. A no-dict baseline below pins that the dict path ran
3486 // (self-compressible payloads would round-trip + stay small via
3487 // in-frame matches alone, proving nothing).
3488 let line = |i: u32| {
3489 alloc::format!(
3490 "ts=2026-03-26T21:{:02}:{:02}Z level=INFO msg=\"event {i:05}\" tenant=t{i} region=eu\n",
3491 i / 60 % 60,
3492 i % 60,
3493 )
3494 .into_bytes()
3495 };
3496 let mut dict_content = Vec::new();
3497 for i in 0..256u32 {
3498 dict_content.extend_from_slice(&line(i));
3499 }
3500 // Payload = the same distinct lines in a different (stride) order, each
3501 // once → no self-repeats, every line is a dictionary match.
3502 let mut payload = Vec::new();
3503 let mut i = 0u32;
3504 for _ in 0..256u32 {
3505 payload.extend_from_slice(&line(i));
3506 i = (i + 97) % 256; // coprime stride → permutation, no adjacency
3507 }
3508
3509 let compress_at = |level, dict: Option<Vec<u8>>| -> Vec<u8> {
3510 let mut compressor = FrameCompressor::new(level);
3511 if let Some(bytes) = dict {
3512 let d = crate::decoding::Dictionary::from_raw_content(dict_id, bytes)
3513 .expect("raw dictionary should be valid");
3514 compressor
3515 .set_dictionary(d)
3516 .expect("dictionary should attach");
3517 }
3518 let mut out = Vec::new();
3519 compressor.set_source(payload.as_slice());
3520 compressor.set_drain(&mut out);
3521 compressor.compress();
3522 out
3523 };
3524
3525 for level in [
3526 super::CompressionLevel::Level(-5), // Fast (negative)
3527 super::CompressionLevel::Level(1), // Fast
3528 super::CompressionLevel::Default, // dfast (L3)
3529 super::CompressionLevel::Level(8), // Row-backed lazy2
3530 ] {
3531 let out = compress_at(level, Some(dict_content.clone()));
3532 let no_dict = compress_at(level, None);
3533 // The dict path MUST measurably beat no-dict on this
3534 // non-self-compressible payload — otherwise the dict probe never
3535 // fired and the roundtrip below would prove nothing.
3536 assert!(
3537 out.len() < no_dict.len(),
3538 "level {level:?}: dict-primed output ({}) must beat no-dict ({}) — dict probe did not fire",
3539 out.len(),
3540 no_dict.len(),
3541 );
3542
3543 let ddict =
3544 crate::decoding::Dictionary::from_raw_content(dict_id, dict_content.clone())
3545 .expect("raw dictionary should be valid");
3546 let mut decoder = FrameDecoder::new();
3547 decoder.add_dict(ddict).expect("decoder dict should attach");
3548 let mut decoded = Vec::with_capacity(payload.len());
3549 decoder
3550 .decode_all_to_vec(&out, &mut decoded)
3551 .unwrap_or_else(|e| panic!("level {level:?}: dict roundtrip decode failed: {e:?}"));
3552 assert_eq!(decoded, payload, "level {level:?}: dict roundtrip mismatch");
3553 }
3554 }
3555
3556 /// Reusing one compressor across independent frames with DIFFERENT
3557 /// dictionaries must drop the per-backend dict cache on each swap
3558 /// (Simple/Dfast/Row keep the attach index across frames). Without the
3559 /// invalidation a later frame would reuse the previous dict's rows.
3560 /// Each frame round-trips through a decoder primed with its own dict.
3561 #[test]
3562 fn dict_swap_across_reused_compressor_roundtrips() {
3563 // Distinct lines per dict (not a single repeated line) so payloads do
3564 // NOT self-compress: each line appears once, so a frame only shrinks if
3565 // the dict probe fires, and — crucially for the invalidation check — if
3566 // frame B reused dict A's stale rows it would emit offsets into A's
3567 // distinct content, which decode under dict B reconstructs as WRONG
3568 // bytes (caught by the roundtrip). A single repeated line would hide
3569 // pollution behind in-frame matches.
3570 let lines = |tag: &str| -> (Vec<u8>, Vec<u8>) {
3571 let line =
3572 |i: u32| alloc::format!("{tag} record {i:05} field=value{i} end\n").into_bytes();
3573 let mut dict = Vec::new();
3574 for i in 0..256u32 {
3575 dict.extend_from_slice(&line(i));
3576 }
3577 let mut payload = Vec::new();
3578 let mut i = 0u32;
3579 for _ in 0..256u32 {
3580 payload.extend_from_slice(&line(i));
3581 i = (i + 97) % 256;
3582 }
3583 (dict, payload)
3584 };
3585 let (dict_a, payload_a) = lines("alpha");
3586 let (dict_b, payload_b) = lines("bravo");
3587
3588 for level in [
3589 super::CompressionLevel::Default,
3590 super::CompressionLevel::Level(8),
3591 ] {
3592 let no_dict = |payload: &[u8]| -> usize {
3593 let mut c: FrameCompressor = FrameCompressor::new(level);
3594 c.compress_independent_frame(payload).len()
3595 };
3596 let no_dict_a = no_dict(&payload_a);
3597 let no_dict_b = no_dict(&payload_b);
3598
3599 let mut compressor: FrameCompressor = FrameCompressor::new(level);
3600 for (dict_bytes, payload, no_dict_len) in [
3601 (&dict_a, &payload_a, no_dict_a),
3602 (&dict_b, &payload_b, no_dict_b),
3603 ] {
3604 let dict =
3605 crate::decoding::Dictionary::from_raw_content(0xD1C7_0002, dict_bytes.clone())
3606 .expect("raw dictionary should be valid");
3607 compressor
3608 .set_dictionary(dict)
3609 .expect("dictionary should attach");
3610 let out = compressor.compress_independent_frame(payload.as_slice());
3611 assert!(
3612 out.len() < no_dict_len,
3613 "level {level:?}: dict frame ({}) must beat no-dict ({}) — dict probe did not fire",
3614 out.len(),
3615 no_dict_len,
3616 );
3617
3618 let ddict =
3619 crate::decoding::Dictionary::from_raw_content(0xD1C7_0002, dict_bytes.clone())
3620 .expect("raw dictionary should be valid");
3621 let mut decoder = FrameDecoder::new();
3622 decoder.add_dict(ddict).expect("decoder dict should attach");
3623 let mut decoded = Vec::with_capacity(payload.len());
3624 decoder
3625 .decode_all_to_vec(&out, &mut decoded)
3626 .unwrap_or_else(|e| panic!("level {level:?}: dict-swap decode failed: {e:?}"));
3627 assert_eq!(
3628 decoded, *payload,
3629 "level {level:?}: dict-swap roundtrip mismatch (stale dict rows?)"
3630 );
3631 }
3632 }
3633 }
3634
3635 #[test]
3636 fn dictionary_roundtrip_stays_valid_after_output_exceeds_window() {
3637 use crate::encoding::match_generator::MatchGeneratorDriver;
3638
3639 let dict_id = 0xABCD_0002;
3640 let dict = crate::decoding::Dictionary::from_raw_content(dict_id, b"abcdefgh".to_vec())
3641 .expect("raw dictionary should be valid");
3642 let dict_for_decoder =
3643 crate::decoding::Dictionary::from_raw_content(dict_id, b"abcdefgh".to_vec())
3644 .expect("raw dictionary should be valid");
3645
3646 // Payload must exceed the encoder's advertised window (512 KiB
3647 // for Fastest after `window_log = 19` alignment with upstream zstd's
3648 // L1 fast row in `clevels.h`) so the test actually exercises
3649 // cross-window-boundary behavior.
3650 let payload = b"abcdefgh".repeat(512 * 1024 / 8 + 64);
3651 let matcher = MatchGeneratorDriver::new(1024, 1);
3652
3653 let mut no_dict_output = Vec::new();
3654 let mut no_dict_compressor =
3655 FrameCompressor::new_with_matcher(matcher, super::CompressionLevel::Fastest);
3656 no_dict_compressor.set_source(payload.as_slice());
3657 no_dict_compressor.set_drain(&mut no_dict_output);
3658 no_dict_compressor.compress();
3659 let (no_dict_frame_header, _) =
3660 crate::decoding::frame::read_frame_header(no_dict_output.as_slice())
3661 .expect("baseline frame should have a header");
3662 let no_dict_window = no_dict_frame_header
3663 .window_size()
3664 .expect("window size should be present");
3665
3666 let mut output = Vec::new();
3667 let matcher = MatchGeneratorDriver::new(1024, 1);
3668 let mut compressor =
3669 FrameCompressor::new_with_matcher(matcher, super::CompressionLevel::Fastest);
3670 compressor
3671 .set_dictionary(dict)
3672 .expect("dictionary should attach");
3673 compressor.set_source(payload.as_slice());
3674 compressor.set_drain(&mut output);
3675 compressor.compress();
3676
3677 let (frame_header, _) = crate::decoding::frame::read_frame_header(output.as_slice())
3678 .expect("encoded frame should have a header");
3679 let advertised_window = frame_header
3680 .window_size()
3681 .expect("window size should be present");
3682 assert_eq!(
3683 advertised_window, no_dict_window,
3684 "dictionary priming must not inflate advertised window size"
3685 );
3686 assert!(
3687 payload.len() > advertised_window as usize,
3688 "test must cross the advertised window boundary"
3689 );
3690
3691 let mut decoder = FrameDecoder::new();
3692 decoder.add_dict(dict_for_decoder).unwrap();
3693 let mut decoded = Vec::with_capacity(payload.len());
3694 decoder.decode_all_to_vec(&output, &mut decoded).unwrap();
3695 assert_eq!(decoded, payload);
3696 }
3697
3698 #[test]
3699 fn source_size_hint_with_dictionary_keeps_roundtrip_and_nonincreasing_window() {
3700 let dict_id = 0xABCD_0004;
3701 let dict_content = b"abcd".repeat(1024); // 4 KiB dictionary history
3702 let dict = crate::decoding::Dictionary::from_raw_content(dict_id, dict_content).unwrap();
3703 let dict_for_decoder =
3704 crate::decoding::Dictionary::from_raw_content(dict_id, b"abcd".repeat(1024)).unwrap();
3705 let payload = b"abcdabcdabcdabcd".repeat(128);
3706
3707 let mut hinted_output = Vec::new();
3708 let mut hinted = FrameCompressor::new(super::CompressionLevel::Fastest);
3709 hinted.set_dictionary(dict).unwrap();
3710 hinted.set_source_size_hint(1);
3711 hinted.set_source(payload.as_slice());
3712 hinted.set_drain(&mut hinted_output);
3713 hinted.compress();
3714
3715 let mut no_hint_output = Vec::new();
3716 let mut no_hint = FrameCompressor::new(super::CompressionLevel::Fastest);
3717 no_hint
3718 .set_dictionary(
3719 crate::decoding::Dictionary::from_raw_content(dict_id, b"abcd".repeat(1024))
3720 .unwrap(),
3721 )
3722 .unwrap();
3723 no_hint.set_source(payload.as_slice());
3724 no_hint.set_drain(&mut no_hint_output);
3725 no_hint.compress();
3726
3727 let hinted_window = crate::decoding::frame::read_frame_header(hinted_output.as_slice())
3728 .expect("encoded frame should have a header")
3729 .0
3730 .window_size()
3731 .expect("window size should be present");
3732 let no_hint_window = crate::decoding::frame::read_frame_header(no_hint_output.as_slice())
3733 .expect("encoded frame should have a header")
3734 .0
3735 .window_size()
3736 .expect("window size should be present");
3737 assert!(
3738 hinted_window <= no_hint_window,
3739 "source-size hint should not increase advertised window with dictionary priming",
3740 );
3741
3742 let mut decoder = FrameDecoder::new();
3743 decoder.add_dict(dict_for_decoder).unwrap();
3744 let mut decoded = Vec::with_capacity(payload.len());
3745 decoder
3746 .decode_all_to_vec(&hinted_output, &mut decoded)
3747 .unwrap();
3748 assert_eq!(decoded, payload);
3749 }
3750
3751 /// A dictionary segment embedded ONCE in otherwise-incompressible
3752 /// input must be matched against the dictionary. Before the fix the
3753 /// raw-fast-path (which skips matching) fired on the
3754 /// incompressible-looking block and the dictionary was never searched,
3755 /// so `with_dict` came out the same size as `no_dict` (the embedded
3756 /// match was lost). Now the block compresses against the dict.
3757 #[test]
3758 fn dictionary_segment_in_incompressible_input_is_matched() {
3759 // Deterministic LCG bytes: high-entropy, so the only compressible
3760 // content is the embedded dictionary segment.
3761 fn lcg(seed: u64, n: usize) -> alloc::vec::Vec<u8> {
3762 let mut s = seed;
3763 (0..n)
3764 .map(|_| {
3765 s = s
3766 .wrapping_mul(6364136223846793005)
3767 .wrapping_add(1442695040888963407);
3768 (s >> 56) as u8
3769 })
3770 .collect()
3771 }
3772 let dict_id = 0x00DC_7777;
3773 let r = lcg(1, 512); // the dictionary content
3774 let mut payload = lcg(2, 2000); // incompressible filler before
3775 payload.extend_from_slice(&r); // the single dict-matchable segment
3776 payload.extend_from_slice(&lcg(3, 1500)); // filler after
3777
3778 // Precondition: the payload must actually look incompressible so
3779 // that the raw-fast-path WOULD fire (and skip matching) without
3780 // the fix. If the heuristic ever changes and this no longer holds,
3781 // the test below would pass vacuously — assert it up front.
3782 assert!(
3783 crate::encoding::incompressible::block_looks_incompressible(&payload),
3784 "test payload must look incompressible to exercise the raw-fast-path",
3785 );
3786
3787 let compress =
3788 |level: super::CompressionLevel, dict: Option<&[u8]>| -> alloc::vec::Vec<u8> {
3789 let mut out = alloc::vec::Vec::new();
3790 let mut c = FrameCompressor::new(level);
3791 if let Some(d) = dict {
3792 c.set_dictionary(
3793 crate::decoding::Dictionary::from_raw_content(dict_id, d.to_vec()).unwrap(),
3794 )
3795 .unwrap();
3796 }
3797 c.set_source(payload.as_slice());
3798 c.set_drain(&mut out);
3799 c.compress();
3800 out
3801 };
3802
3803 for lvl in [
3804 super::CompressionLevel::Level(2),
3805 super::CompressionLevel::Level(6),
3806 super::CompressionLevel::Level(19),
3807 ] {
3808 let with_dict = compress(lvl, Some(&r));
3809 let no_dict = compress(lvl, None);
3810 // The 512-byte dict segment should be matched, saving most of
3811 // its length (generous slack for sequence/header coding).
3812 assert!(
3813 with_dict.len() + 300 < no_dict.len(),
3814 "{lvl:?}: dict segment not matched (with_dict={}, no_dict={})",
3815 with_dict.len(),
3816 no_dict.len(),
3817 );
3818 // The dict-compressed frame must round-trip through the decoder.
3819 let mut decoder = FrameDecoder::new();
3820 decoder
3821 .add_dict(
3822 crate::decoding::Dictionary::from_raw_content(dict_id, r.clone()).unwrap(),
3823 )
3824 .unwrap();
3825 let mut decoded = Vec::with_capacity(payload.len());
3826 decoder.decode_all_to_vec(&with_dict, &mut decoded).unwrap();
3827 assert_eq!(decoded, payload, "{lvl:?}: dict round-trip mismatch");
3828
3829 // A dictionary that does NOT appear in the input must not make
3830 // the output larger than the no-dict (raw) encoding: the
3831 // post-compress raw fallback covers incompressible-with-dict.
3832 let unrelated = lcg(99, 512);
3833 let with_bad_dict = compress(lvl, Some(&unrelated));
3834 assert!(
3835 with_bad_dict.len() <= no_dict.len() + 16,
3836 "{lvl:?}: unhelpful dict expanded output (with={}, no_dict={})",
3837 with_bad_dict.len(),
3838 no_dict.len(),
3839 );
3840 }
3841 }
3842
3843 #[test]
3844 fn source_size_hint_with_dictionary_keeps_roundtrip_for_larger_payload() {
3845 let dict_id = 0xABCD_0005;
3846 let dict_content = b"abcd".repeat(1024); // 4 KiB dictionary history
3847 let dict = crate::decoding::Dictionary::from_raw_content(dict_id, dict_content).unwrap();
3848 let dict_for_decoder =
3849 crate::decoding::Dictionary::from_raw_content(dict_id, b"abcd".repeat(1024)).unwrap();
3850 let payload = b"abcd".repeat(1024); // 4 KiB payload
3851 let payload_len = payload.len() as u64;
3852
3853 let mut hinted_output = Vec::new();
3854 let mut hinted = FrameCompressor::new(super::CompressionLevel::Fastest);
3855 hinted.set_dictionary(dict).unwrap();
3856 hinted.set_source_size_hint(payload_len);
3857 hinted.set_source(payload.as_slice());
3858 hinted.set_drain(&mut hinted_output);
3859 hinted.compress();
3860
3861 let mut no_hint_output = Vec::new();
3862 let mut no_hint = FrameCompressor::new(super::CompressionLevel::Fastest);
3863 no_hint
3864 .set_dictionary(
3865 crate::decoding::Dictionary::from_raw_content(dict_id, b"abcd".repeat(1024))
3866 .unwrap(),
3867 )
3868 .unwrap();
3869 no_hint.set_source(payload.as_slice());
3870 no_hint.set_drain(&mut no_hint_output);
3871 no_hint.compress();
3872
3873 let hinted_window = crate::decoding::frame::read_frame_header(hinted_output.as_slice())
3874 .expect("encoded frame should have a header")
3875 .0
3876 .window_size()
3877 .expect("window size should be present");
3878 let no_hint_window = crate::decoding::frame::read_frame_header(no_hint_output.as_slice())
3879 .expect("encoded frame should have a header")
3880 .0
3881 .window_size()
3882 .expect("window size should be present");
3883 assert!(
3884 hinted_window <= no_hint_window,
3885 "source-size hint should not increase advertised window with dictionary priming",
3886 );
3887
3888 let mut decoder = FrameDecoder::new();
3889 decoder.add_dict(dict_for_decoder).unwrap();
3890 let mut decoded = Vec::with_capacity(payload.len());
3891 decoder
3892 .decode_all_to_vec(&hinted_output, &mut decoded)
3893 .unwrap();
3894 assert_eq!(decoded, payload);
3895 }
3896
3897 #[test]
3898 fn custom_matcher_without_dictionary_priming_does_not_advertise_dict_id() {
3899 let dict_id = 0xABCD_0003;
3900 let dict = crate::decoding::Dictionary::from_raw_content(dict_id, b"abcdefgh".to_vec())
3901 .expect("raw dictionary should be valid");
3902 let payload = b"abcdefghabcdefgh";
3903
3904 let mut output = Vec::new();
3905 let matcher = NoDictionaryMatcher::new(64);
3906 let mut compressor =
3907 FrameCompressor::new_with_matcher(matcher, super::CompressionLevel::Fastest);
3908 compressor
3909 .set_dictionary(dict)
3910 .expect("dictionary should attach");
3911 compressor.set_source(payload.as_slice());
3912 compressor.set_drain(&mut output);
3913 compressor.compress();
3914
3915 let (frame_header, _) = crate::decoding::frame::read_frame_header(output.as_slice())
3916 .expect("encoded frame should have a header");
3917 assert_eq!(
3918 frame_header.dictionary_id(),
3919 None,
3920 "matchers that do not support dictionary priming must not advertise dictionary dependency"
3921 );
3922
3923 let mut decoder = FrameDecoder::new();
3924 let mut decoded = Vec::with_capacity(payload.len());
3925 decoder.decode_all_to_vec(&output, &mut decoded).unwrap();
3926 assert_eq!(decoded, payload);
3927 }
3928
3929 #[cfg(feature = "hash")]
3930 #[test]
3931 fn checksum_two_frames_reused_compressor() {
3932 // Compress the same data twice using the same compressor and verify that:
3933 // 1. The checksum written in each frame matches what the decoder calculates.
3934 // 2. The hasher is correctly reset between frames (no cross-contamination).
3935 // If the hasher were NOT reset, the second frame's calculated checksum
3936 // would differ from the one stored in the frame data, causing assert_eq to fail.
3937 let data: Vec<u8> = (0u8..=255).cycle().take(1024).collect();
3938
3939 let mut compressor = FrameCompressor::new(super::CompressionLevel::Uncompressed);
3940
3941 // --- Frame 1 ---
3942 let mut compressed1 = Vec::new();
3943 compressor.set_source(data.as_slice());
3944 compressor.set_drain(&mut compressed1);
3945 compressor.compress();
3946
3947 // --- Frame 2 (reuse the same compressor) ---
3948 let mut compressed2 = Vec::new();
3949 compressor.set_source(data.as_slice());
3950 compressor.set_drain(&mut compressed2);
3951 compressor.compress();
3952
3953 fn decode_and_collect(compressed: &[u8]) -> (Vec<u8>, Option<u32>, Option<u32>) {
3954 let mut decoder = FrameDecoder::new();
3955 let mut source = compressed;
3956 decoder.reset(&mut source).unwrap();
3957 while !decoder.is_finished() {
3958 decoder
3959 .decode_blocks(&mut source, crate::decoding::BlockDecodingStrategy::All)
3960 .unwrap();
3961 }
3962 let mut decoded = Vec::new();
3963 decoder.collect_to_writer(&mut decoded).unwrap();
3964 (
3965 decoded,
3966 decoder.get_checksum_from_data(),
3967 decoder.get_calculated_checksum(),
3968 )
3969 }
3970
3971 let (decoded1, chksum_from_data1, chksum_calculated1) = decode_and_collect(&compressed1);
3972 assert_eq!(decoded1, data, "frame 1: decoded data mismatch");
3973 assert_eq!(
3974 chksum_from_data1, chksum_calculated1,
3975 "frame 1: checksum mismatch"
3976 );
3977
3978 let (decoded2, chksum_from_data2, chksum_calculated2) = decode_and_collect(&compressed2);
3979 assert_eq!(decoded2, data, "frame 2: decoded data mismatch");
3980 assert_eq!(
3981 chksum_from_data2, chksum_calculated2,
3982 "frame 2: checksum mismatch"
3983 );
3984
3985 // Same data compressed twice must produce the same checksum.
3986 // If state leaked across frames, the second calculated checksum would differ.
3987 assert_eq!(
3988 chksum_from_data1, chksum_from_data2,
3989 "frame 1 and frame 2 should have the same checksum (same data, hash must reset per frame)"
3990 );
3991 }
3992
3993 #[cfg(feature = "lsm")]
3994 #[test]
3995 fn frame_emit_info_decompressed_ranges_match_decoded_output() {
3996 // Part A correctness: the per-block `decompressed_size` captured during
3997 // encode (and the `decompressed_byte_range` prefix sum derived from it)
3998 // must describe the real decoded output exactly — one entry per
3999 // physical block, contiguous, summing to the full decompressed length.
4000 // A multi-block compressible payload exercises the Compressed-block
4001 // path (whose regenerated size is NOT on the wire, so it relies on the
4002 // encode-side capture this test guards).
4003 let data = emit_info_fixture_data();
4004
4005 // Cover both the single-block-per-chunk path (Default) and the
4006 // Level(16..=22) post-split path (multiple physical partitions per
4007 // input chunk), since lsm-tree compresses at zstd:22 and post-split
4008 // is the riskiest capture site (per-partition `src_size`).
4009 for level in [
4010 super::CompressionLevel::Default,
4011 super::CompressionLevel::Level(22),
4012 ] {
4013 let mut compressed = Vec::new();
4014 let mut compressor = FrameCompressor::new(level);
4015 // Pledge the source size so the high-level (22) window shrinks to
4016 // fit the payload, keeping the frame compact (no oversized window
4017 // descriptor for a small input). Still >= 128 KiB, so post-split
4018 // eligibility is preserved.
4019 compressor.set_source_size_hint(data.len() as u64);
4020 compressor.set_source(data.as_slice());
4021 compressor.set_drain(&mut compressed);
4022 compressor.compress();
4023
4024 let info = compressor
4025 .last_frame_emit_info()
4026 .expect("emit info populated after compress")
4027 .clone();
4028
4029 // Reference: full decode of the same frame.
4030 let mut decoder = FrameDecoder::new();
4031 let mut source = compressed.as_slice();
4032 decoder.reset(&mut source).unwrap();
4033 while !decoder.is_finished() {
4034 decoder
4035 .decode_blocks(&mut source, crate::decoding::BlockDecodingStrategy::All)
4036 .unwrap();
4037 }
4038 let mut decoded = Vec::new();
4039 decoder.collect_to_writer(&mut decoded).unwrap();
4040 assert_eq!(decoded, data, "sanity: frame must round-trip ({level:?})");
4041
4042 assert!(
4043 info.blocks.len() >= 2,
4044 "fixture must span multiple blocks to exercise the mapping ({level:?}, got {})",
4045 info.blocks.len()
4046 );
4047 assert!(
4048 info.blocks.last().unwrap().last_block,
4049 "final block must carry last_block ({level:?})"
4050 );
4051
4052 // Pin the Level(22) post-split path: the owned loop feeds the
4053 // encoder MAX_BLOCK_SIZE input chunks, so without post-split the
4054 // block count cannot exceed the chunk count. More blocks than
4055 // chunks proves at least one chunk was split into multiple physical
4056 // partitions (the per-partition `src_size` capture under test).
4057 if matches!(level, super::CompressionLevel::Level(22)) {
4058 let max_block = crate::common::MAX_BLOCK_SIZE as usize;
4059 let n_chunks = data.len().div_ceil(max_block);
4060 assert!(
4061 info.blocks.len() > n_chunks,
4062 "Level(22) must exercise post-split: {} blocks for {} input chunks",
4063 info.blocks.len(),
4064 n_chunks
4065 );
4066 }
4067
4068 // Per-block ranges: contiguous, zero-based, summing to the full output.
4069 let mut expected_start = 0u64;
4070 for i in 0..info.blocks.len() {
4071 let range = info
4072 .decompressed_byte_range(i)
4073 .expect("in-bounds block has a range");
4074 assert_eq!(
4075 range.start, expected_start,
4076 "block {i} range must start where the previous ended ({level:?})"
4077 );
4078 assert_eq!(
4079 u64::from(info.blocks[i].decompressed_size),
4080 range.end - range.start,
4081 "block {i} decompressed_size must equal its range width ({level:?})"
4082 );
4083 // Validate the mapping against REAL per-block bytes, not just
4084 // prefix-sum consistency: decode block `i` alone and require it
4085 // to equal the corresponding slice of the full decode. A
4086 // sidecar that swapped sizes between adjacent blocks (same sum,
4087 // same contiguity) would fail here.
4088 let mut psrc = compressed.as_slice();
4089 let mut pdec = FrameDecoder::new();
4090 pdec.reset(&mut psrc).unwrap();
4091 let pd = pdec
4092 .decode_blocks_partial(&mut psrc, i as u32, i as u32 + 1, None, false)
4093 .unwrap();
4094 assert!(
4095 pd.stopped_at.is_none(),
4096 "block {i} must decode cleanly ({level:?})"
4097 );
4098 assert_eq!(
4099 pd.data.as_slice(),
4100 &decoded[range.start as usize..range.end as usize],
4101 "block {i} partial-decode bytes must equal the full-decode slice ({level:?})"
4102 );
4103 expected_start = range.end;
4104 }
4105 assert_eq!(
4106 expected_start,
4107 decoded.len() as u64,
4108 "block decompressed sizes must sum to the full decoded length ({level:?})"
4109 );
4110 assert_eq!(
4111 info.decompressed_byte_range(info.blocks.len()),
4112 None,
4113 "out-of-range index yields None ({level:?})"
4114 );
4115 }
4116 }
4117
4118 /// ~400 KiB semi-repetitive payload (long runs interleaved with a stride
4119 /// phrase) that compresses into several multi-block frames across levels.
4120 #[cfg(feature = "lsm")]
4121 fn emit_info_fixture_data() -> Vec<u8> {
4122 let mut data: Vec<u8> = Vec::with_capacity(400 * 1024);
4123 let mut x = 0x9E37_79B9u32;
4124 while data.len() < 400 * 1024 {
4125 x ^= x << 13;
4126 x ^= x >> 17;
4127 x ^= x << 5;
4128 let run = 16 + (x as usize % 48);
4129 let byte = (x >> 24) as u8;
4130 for _ in 0..run {
4131 data.push(byte);
4132 }
4133 data.extend_from_slice(b"the quick brown fox jumps over the lazy dog\n");
4134 }
4135 data
4136 }
4137
4138 #[cfg(feature = "lsm")]
4139 #[test]
4140 fn frame_emit_info_decompressed_ranges_match_on_borrowed_oneshot_path() {
4141 // The borrowed one-shot path (`compress_independent_frame` ->
4142 // `run_borrowed_block_loop` -> `compress_block_encoded_borrowed`)
4143 // threads the decompressed-size sidecar through a DIFFERENT emit site
4144 // than the owned/streaming loop, so it needs its own per-block mapping
4145 // check. A Fast level keeps the encoder on the borrowed-eligible
4146 // (Simple matcher) path.
4147 let data = emit_info_fixture_data();
4148
4149 let mut compressor: FrameCompressor =
4150 FrameCompressor::new(super::CompressionLevel::Fastest);
4151 let compressed = compressor.compress_independent_frame(data.as_slice());
4152 let info = compressor
4153 .last_frame_emit_info()
4154 .expect("emit info populated after compress_independent_frame")
4155 .clone();
4156 // Pin the compressed-block path: without this the fixture could regress
4157 // into the raw-fast fallback and still pass via the Raw wire-size
4158 // fallback in populate_frame_emit_info, never exercising the borrowed
4159 // compressed-block sidecar capture this test targets.
4160 assert!(
4161 info.blocks
4162 .iter()
4163 .any(|b| matches!(b.block_type, crate::blocks::block::BlockType::Compressed)),
4164 "borrowed-path fixture must emit at least one compressed block"
4165 );
4166 assert!(
4167 info.blocks.len() >= 2,
4168 "borrowed fixture must span multiple blocks (got {})",
4169 info.blocks.len()
4170 );
4171 assert!(info.blocks.last().unwrap().last_block);
4172
4173 // Full decode reference.
4174 let mut decoder = FrameDecoder::new();
4175 let mut source = compressed.as_slice();
4176 decoder.reset(&mut source).unwrap();
4177 while !decoder.is_finished() {
4178 decoder
4179 .decode_blocks(&mut source, crate::decoding::BlockDecodingStrategy::All)
4180 .unwrap();
4181 }
4182 let mut decoded = Vec::new();
4183 decoder.collect_to_writer(&mut decoded).unwrap();
4184 assert_eq!(decoded, data, "borrowed one-shot frame must round-trip");
4185
4186 // Each block's mapping must match real per-block bytes.
4187 let mut expected_start = 0u64;
4188 for i in 0..info.blocks.len() {
4189 let range = info.decompressed_byte_range(i).unwrap();
4190 assert_eq!(range.start, expected_start, "block {i} range contiguity");
4191 let mut psrc = compressed.as_slice();
4192 let mut pdec = FrameDecoder::new();
4193 pdec.reset(&mut psrc).unwrap();
4194 let pd = pdec
4195 .decode_blocks_partial(&mut psrc, i as u32, i as u32 + 1, None, false)
4196 .unwrap();
4197 assert!(pd.stopped_at.is_none(), "block {i} must decode cleanly");
4198 assert_eq!(
4199 pd.data.as_slice(),
4200 &decoded[range.start as usize..range.end as usize],
4201 "borrowed block {i} partial-decode bytes must equal the full-decode slice"
4202 );
4203 expected_start = range.end;
4204 }
4205 assert_eq!(
4206 expected_start,
4207 decoded.len() as u64,
4208 "ranges sum to full length"
4209 );
4210 }
4211
4212 // The fuzz-artifact interop replay (C-compress -> our-decode and
4213 // our-compress -> C-decode) moved to `ffi-bench/tests/fuzz_interop.rs` so
4214 // the library crate never links libzstd.
4215
4216 /// Homogeneous input — every byte the same — must NOT be split:
4217 /// both border histograms are identical (all 512 hits on a single
4218 /// slot), so `presplit_fingerprints_differ` returns `false` and the
4219 /// function takes the early-return path at
4220 /// `zstd_preSplit.c:214` returning `blockSize`.
4221 #[test]
4222 fn split_block_from_borders_keeps_homogeneous_block() {
4223 let block = vec![0xAAu8; MAX_BLOCK_SIZE as usize];
4224 let split = super::split_block_from_borders(&block);
4225 assert_eq!(split, MAX_BLOCK_SIZE as usize);
4226 }
4227
4228 /// Heterogeneous input — first half all zeros, second half a
4229 /// counter sequence — has clearly distinguishable border
4230 /// histograms, so the borders heuristic decides to split.
4231 ///
4232 /// The transition sits at exactly the block midpoint, so the
4233 /// middle 512-byte sample (`block[mid-256..mid+256]`) is half
4234 /// zeros + half counter values. That makes it roughly
4235 /// equidistant from both border fingerprints — the
4236 /// `abs_diff(dist_from_begin, dist_from_end) < min_distance`
4237 /// branch fires and the heuristic returns the midpoint (64 KiB)
4238 /// per `zstd_preSplit.c:222`. The test asserts the exact value
4239 /// rather than just "one of {32K, 64K, 96K}" so a regression
4240 /// to a different quantised arm cannot silently slip through.
4241 #[test]
4242 fn split_block_from_borders_returns_midpoint_for_centred_transition() {
4243 let mut block = vec![0u8; MAX_BLOCK_SIZE as usize];
4244 for (i, byte) in block
4245 .iter_mut()
4246 .enumerate()
4247 .skip(MAX_BLOCK_SIZE as usize / 2)
4248 {
4249 *byte = (i % 251 + 1) as u8;
4250 }
4251 let split = super::split_block_from_borders(&block);
4252 assert_eq!(
4253 split,
4254 64 * 1024,
4255 "centred-transition fixture must take the symmetric \
4256 midpoint arm (`abs_diff < min_distance`), got {split}"
4257 );
4258 }
4259
4260 /// `level_pre_split` resolves the per-level split knob through the
4261 /// `LevelParams` table, mirroring the upstream zstd `splitLevels[]` by strategy
4262 /// (`ZSTD_optimalBlockSize`): fast → 0 (from-borders), dfast → 1,
4263 /// greedy/lazy → 2, lazy2/btlazy2 (Lazy tag at depth 2) → 3,
4264 /// btopt/btultra/btultra2 → 4. `Uncompressed` has no numeric level so it
4265 /// stays `None`.
4266 #[test]
4267 fn pre_split_level_dispatches_by_compression_level() {
4268 use crate::encoding::CompressionLevel;
4269 use crate::encoding::match_generator::level_pre_split;
4270 assert_eq!(level_pre_split(CompressionLevel::Uncompressed), None);
4271 // Fastest = level 1 (fast) → 0 (from-borders).
4272 assert_eq!(level_pre_split(CompressionLevel::Fastest), Some(0));
4273 // Default = level 3 (dfast) → 1.
4274 assert_eq!(level_pre_split(CompressionLevel::Default), Some(1));
4275 // Better is a pure alias for level 7 (lazy): same as Level(7).
4276 assert_eq!(
4277 level_pre_split(CompressionLevel::Better),
4278 level_pre_split(CompressionLevel::Level(7)),
4279 );
4280 // Best resolves to the level-13 table row (btlazy2): pin it to that
4281 // numeric route so the named path can't drift from the pre-split
4282 // table.
4283 assert_eq!(
4284 level_pre_split(CompressionLevel::Best),
4285 level_pre_split(CompressionLevel::Level(13)),
4286 );
4287 assert_eq!(level_pre_split(CompressionLevel::Level(2)), Some(0)); // fast
4288 assert_eq!(level_pre_split(CompressionLevel::Level(4)), Some(1)); // dfast
4289 assert_eq!(level_pre_split(CompressionLevel::Level(5)), Some(2)); // greedy
4290 assert_eq!(level_pre_split(CompressionLevel::Level(7)), Some(2)); // lazy (depth 1)
4291 // lazy2 / btlazy2 use the rate-1 full-scan splitter (4), not the
4292 // rate-5 sampler (3): the sampler phantom-splits homogeneous periodic
4293 // input (see `pre_split` comment + `periodic_stream_not_oversplit`).
4294 assert_eq!(level_pre_split(CompressionLevel::Level(8)), Some(4)); // lazy2 lower bound
4295 assert_eq!(level_pre_split(CompressionLevel::Level(11)), Some(4)); // lazy2 (depth 2)
4296 assert_eq!(level_pre_split(CompressionLevel::Level(12)), Some(4)); // lazy2 upper bound
4297 assert_eq!(level_pre_split(CompressionLevel::Level(13)), Some(4)); // btlazy2 lower bound
4298 assert_eq!(level_pre_split(CompressionLevel::Level(15)), Some(4)); // btlazy2 (depth 2)
4299 assert_eq!(level_pre_split(CompressionLevel::Level(16)), Some(4)); // btopt
4300 assert_eq!(level_pre_split(CompressionLevel::Level(22)), Some(4)); // btultra2
4301 }
4302
4303 /// Regression: a homogeneous but periodic multi-block stream must not be
4304 /// pre-split into tiny blocks at the lazy2 / btlazy2 levels. The rate-5
4305 /// chunk sampler used to phantom-split such input at every 8 KB chunk,
4306 /// cascading a large stream into hundreds of tiny blocks whose per-block
4307 /// headers ballooned the output (~5x vs the lazy level next door). With
4308 /// the rate-1 full-scan splitter the periodic stream is seen as uniform
4309 /// and stays a few full blocks. We assert the lazy2 (L8) and btlazy2 (L15)
4310 /// outputs stay within 2x of the lazy (L7) output on the same input, and
4311 /// that every output round-trips.
4312 #[test]
4313 fn periodic_stream_not_oversplit() {
4314 use crate::encoding::{CompressionLevel, compress_slice_to_vec};
4315 const LINES: &[&str] = &[
4316 "ts=2026-03-26T21:39:28Z level=INFO msg=\"flush memtable\" tenant=demo table=orders region=eu-west\n",
4317 "ts=2026-03-26T21:39:29Z level=INFO msg=\"rotate segment\" tenant=demo table=orders region=eu-west\n",
4318 "ts=2026-03-26T21:39:30Z level=INFO msg=\"compact level\" tenant=demo table=orders region=eu-west\n",
4319 "ts=2026-03-26T21:39:31Z level=INFO msg=\"write block\" tenant=demo table=orders region=eu-west\n",
4320 ];
4321 // 512 KB = 4 upstream zstd blocks, enough for the cascade to manifest.
4322 let target = 512 * 1024usize;
4323 let mut data = Vec::with_capacity(target);
4324 let mut i = 0;
4325 while data.len() < target {
4326 let line = LINES[i % LINES.len()].as_bytes();
4327 let take = line.len().min(target - data.len());
4328 data.extend_from_slice(&line[..take]);
4329 i += 1;
4330 }
4331 let l7 = compress_slice_to_vec(&data, CompressionLevel::Level(7)); // lazy depth1
4332 let l8 = compress_slice_to_vec(&data, CompressionLevel::Level(8)); // lazy2
4333 let l15 = compress_slice_to_vec(&data, CompressionLevel::Level(15)); // btlazy2
4334 assert!(
4335 l8.len() < l7.len() * 2,
4336 "lazy2 over-split periodic stream: l7={} l8={}",
4337 l7.len(),
4338 l8.len()
4339 );
4340 assert!(
4341 l15.len() < l7.len() * 2,
4342 "btlazy2 over-split periodic stream: l7={} l15={}",
4343 l7.len(),
4344 l15.len()
4345 );
4346 for out in [&l7, &l8, &l15] {
4347 let mut decoder = FrameDecoder::new();
4348 let mut round = Vec::with_capacity(data.len());
4349 decoder
4350 .decode_all_to_vec(out, &mut round)
4351 .expect("decode periodic stream");
4352 assert_eq!(round, data, "periodic stream roundtrip mismatch");
4353 }
4354 }
4355
4356 /// End-to-end: a 256 KB payload whose SECOND 128 KB upstream zstd block carries
4357 /// an intra-block fingerprint transition, compressed at Level(5)
4358 /// (greedy, the pre-split path this revision routes through the cheap
4359 /// chunk splitter), round-trips through the crate's own decoder.
4360 ///
4361 /// The transition lives in the second block on purpose: the upstream zstd
4362 /// `savings < 3` gate skips splitting the first block (savings start at
4363 /// 0), so the first block is a homogeneous compressible run that banks
4364 /// savings, and the second block is the one whose intra-block transition
4365 /// `split_block_by_chunks()` resolves into a sub-block boundary (the
4366 /// `pending_input.split_off(...)` path). The test asserts that split
4367 /// decision directly so it cannot silently stop exercising the path if
4368 /// the fixture or params drift, then proves the emitted split frame
4369 /// round-trips. Level 13 (lazy) no longer pre-splits, hence Level 5.
4370 #[test]
4371 fn greedy_chunk_split_roundtrips_through_own_decoder() {
4372 use crate::encoding::CompressionLevel;
4373 let mut data = vec![0u8; 256 * 1024];
4374 // First 128 KB: homogeneous low-entropy run (compressible, banks
4375 // the savings the upstream zstd gate needs). Second 128 KB: low-entropy run
4376 // for its first half, then a counter sequence: a clear intra-block
4377 // fingerprint transition at the 192 KB midpoint for the chunk
4378 // splitter to find.
4379 for (i, byte) in data.iter_mut().enumerate() {
4380 *byte = if i < 192 * 1024 {
4381 (i & 0x07) as u8
4382 } else {
4383 (i % 251 + 1) as u8
4384 };
4385 }
4386
4387 // Directly assert the chunk splitter resolves the second block's
4388 // intra-block transition into a sub-block boundary once savings have
4389 // accrued (the compressible first block banks well over the gate).
4390 let second_block = &data[128 * 1024..];
4391 let split = super::optimal_block_size(
4392 CompressionLevel::Level(5),
4393 second_block,
4394 second_block.len(),
4395 MAX_BLOCK_SIZE as usize,
4396 100,
4397 );
4398 assert!(
4399 split < MAX_BLOCK_SIZE as usize,
4400 "second upstream zstd block must chunk-split at its intra-block transition, got {split}",
4401 );
4402
4403 let mut compressed = Vec::new();
4404 let mut compressor = FrameCompressor::new(CompressionLevel::Level(5));
4405 compressor.set_source(data.as_slice());
4406 compressor.set_drain(&mut compressed);
4407 compressor.compress();
4408
4409 let mut decoder = FrameDecoder::new();
4410 let mut source = compressed.as_slice();
4411 decoder
4412 .reset(&mut source)
4413 .expect("frame header should parse");
4414 while !decoder.is_finished() {
4415 decoder
4416 .decode_blocks(&mut source, crate::decoding::BlockDecodingStrategy::All)
4417 .expect("decode should succeed");
4418 }
4419 let mut decoded = Vec::with_capacity(data.len());
4420 decoder.collect_to_writer(&mut decoded).unwrap();
4421 assert_eq!(decoded, data, "roundtrip must reproduce the input verbatim");
4422 }
4423
4424 /// Outside-diff coverage for the FAST one-shot path.
4425 /// `compress_slice_to_vec` / `compress_independent_frame` on a Fast level
4426 /// routes through `run_borrowed_block_loop` (not the owned loop the test
4427 /// above covers), which must honour `optimal_block_size` and emit a
4428 /// sub-`MAX_BLOCK_SIZE` boundary rather than fixed 128 KiB blocks. A
4429 /// 256 KiB input is two 128 KiB blocks when unsplit; a chunk boundary in
4430 /// the second block yields >= 3 decoded blocks, asserted on the round-trip.
4431 #[test]
4432 fn fast_oneshot_borrowed_split_emits_subblock() {
4433 use crate::encoding::CompressionLevel;
4434 // First 192 KiB: homogeneous zero run (banks the savings the split
4435 // gate needs). The second 128 KiB block flips to a counter sequence
4436 // at its 64 KiB midpoint (the 192 KiB mark) — a fingerprint
4437 // transition the Fast from-borders splitter (split level 0) resolves
4438 // into a sub-block boundary.
4439 let mut data = vec![0u8; 256 * 1024];
4440 for (i, byte) in data.iter_mut().enumerate() {
4441 if i >= 192 * 1024 {
4442 *byte = (i % 251 + 1) as u8;
4443 }
4444 }
4445
4446 // Pin the splitter decision for the Fast path directly (mirrors the
4447 // greedy test): the second upstream zstd block must resolve to a sub-block
4448 // boundary, so the >= 3 block count below cannot pass vacuously.
4449 let second_block = &data[128 * 1024..];
4450 assert!(
4451 super::optimal_block_size(
4452 CompressionLevel::Fastest,
4453 second_block,
4454 second_block.len(),
4455 MAX_BLOCK_SIZE as usize,
4456 100,
4457 ) < MAX_BLOCK_SIZE as usize,
4458 "fixture must resolve to a sub-block split in the second upstream zstd block",
4459 );
4460
4461 // Drive the borrowed one-shot route explicitly (Fast level ->
4462 // run_borrowed_block_loop via compress_independent_frame).
4463 let mut compressor: FrameCompressor = FrameCompressor::new(CompressionLevel::Fastest);
4464 let frame = compressor.compress_independent_frame(&data);
4465
4466 let mut decoder = FrameDecoder::new();
4467 let mut source = frame.as_slice();
4468 decoder
4469 .reset(&mut source)
4470 .expect("frame header should parse");
4471 while !decoder.is_finished() {
4472 decoder
4473 .decode_blocks(&mut source, crate::decoding::BlockDecodingStrategy::All)
4474 .expect("decode should succeed");
4475 }
4476 let mut decoded = Vec::with_capacity(data.len());
4477 decoder.collect_to_writer(&mut decoded).unwrap();
4478 assert_eq!(decoded, data, "roundtrip must reproduce the input verbatim");
4479 assert!(
4480 decoder.blocks_decoded() >= 3,
4481 "fast one-shot borrowed path must split the second upstream zstd block \
4482 (256 KiB unsplit = 2 blocks), got {} blocks",
4483 decoder.blocks_decoded(),
4484 );
4485 }
4486
4487 /// Regression: `set_compression_level` followed by `compress()` must
4488 /// refresh `state.strategy_tag` through the reset-time sync so the
4489 /// literal-compression gates (`min_literals_to_compress`,
4490 /// `min_gain`) use the NEW level's strategy. Picks a level pair
4491 /// that genuinely crosses strategy bands — `Fastest` resolves to
4492 /// `Fast`, `Level(20)` resolves to `BtUltra2` — so a missed sync
4493 /// would leave the construction-time tag visible and trip the
4494 /// assertion. `CompressionLevel::Best` would also pass type-wise
4495 /// but resolves to `Lazy` today, which keeps `min_literals_to_compress`
4496 /// in the same `shift=3 → 64-byte` band as `Fast` and weakens the
4497 /// signal that the gate floor actually moved.
4498 #[cfg(feature = "std")]
4499 #[test]
4500 fn set_compression_level_then_compress_refreshes_strategy_tag() {
4501 use super::CompressionLevel;
4502 use crate::encoding::strategy::StrategyTag;
4503
4504 let data = vec![0xABu8; 256];
4505 let mut out = Vec::new();
4506 let mut compressor = FrameCompressor::new(CompressionLevel::Fastest);
4507 let initial_tag = compressor.state.strategy_tag;
4508 assert_eq!(
4509 initial_tag,
4510 StrategyTag::for_compression_level(CompressionLevel::Fastest),
4511 "construction-time strategy_tag must reflect initial level",
4512 );
4513
4514 // Switch to a level whose resolved strategy lives in a different
4515 // band, then run a full compress cycle — the matcher.reset()
4516 // inside `compress` is the only site that can refresh the tag.
4517 let new_level = CompressionLevel::Level(20);
4518 compressor.set_compression_level(new_level);
4519 compressor.set_source(data.as_slice());
4520 compressor.set_drain(&mut out);
4521 compressor.compress();
4522
4523 let new_tag = compressor.state.strategy_tag;
4524 let expected = StrategyTag::for_compression_level(new_level);
4525 assert_eq!(
4526 new_tag, expected,
4527 "strategy_tag must follow set_compression_level → compress, \
4528 got {new_tag:?} expected {expected:?}",
4529 );
4530 assert_eq!(
4531 expected,
4532 StrategyTag::BtUltra2,
4533 "test fixture invariant: Level(20) must resolve to BtUltra2 \
4534 so the post-switch tag visibly crosses the band boundary",
4535 );
4536 assert_ne!(
4537 new_tag, initial_tag,
4538 "test fixture invariant: chosen levels must resolve to \
4539 different StrategyTag variants",
4540 );
4541 }
4542
4543 /// Magicless mode (`ZSTD_f_zstd1_magicless`): encoded frame
4544 /// MUST NOT start with the 4-byte magic prefix, AND must
4545 /// round-trip through a magicless-aware decoder.
4546 #[test]
4547 fn magicless_frame_omits_magic_and_roundtrips() {
4548 use crate::common::MAGIC_NUM;
4549 let input: alloc::vec::Vec<u8> = (0..512u32).map(|i| (i ^ 0xA5) as u8).collect();
4550
4551 // Encode with magicless = true.
4552 let mut output: Vec<u8> = Vec::new();
4553 let mut compressor = FrameCompressor::new(super::CompressionLevel::Default);
4554 compressor.set_magicless(true);
4555 compressor.set_source(input.as_slice());
4556 compressor.set_drain(&mut output);
4557 compressor.compress();
4558
4559 // 1. Encoded output must NOT begin with the zstd magic number.
4560 assert!(
4561 !output.starts_with(&MAGIC_NUM.to_le_bytes()),
4562 "magicless frame must omit the 4-byte magic prefix",
4563 );
4564
4565 // 2. A magicless-aware decoder must round-trip the payload.
4566 let mut decoder = crate::decoding::FrameDecoder::new();
4567 decoder.set_magicless(true);
4568 let mut cursor: &[u8] = output.as_slice();
4569 decoder.init(&mut cursor).expect("magicless init");
4570 decoder
4571 .decode_blocks(&mut cursor, crate::decoding::BlockDecodingStrategy::All)
4572 .expect("decode_blocks");
4573 let mut decoded: Vec<u8> = Vec::new();
4574 decoder
4575 .collect_to_writer(&mut decoded)
4576 .expect("collect_to_writer");
4577 assert_eq!(decoded, input, "magicless roundtrip must preserve bytes");
4578
4579 // 3. A standard (magicful) decoder MUST reject a magicless
4580 // frame at the header-read step — the first 4 bytes are
4581 // the frame-header descriptor + window / dictionary / FCS
4582 // metadata, not the magic. We accept either
4583 // `BadMagicNumber` (typical case: first 4 bytes don't
4584 // match `MAGIC_NUM` and don't fall in the skippable-frame
4585 // magic range) or `SkipFrame` (rare: the first 4 bytes
4586 // coincidentally land in `0x184D2A50..=0x184D2A5F`). Both
4587 // prove the standard decoder did not treat the bytes as a
4588 // real magicful frame.
4589 use crate::decoding::errors::{FrameDecoderError, ReadFrameHeaderError};
4590 let mut std_decoder = crate::decoding::FrameDecoder::new();
4591 let std_init = std_decoder.init(output.as_slice());
4592 match std_init {
4593 Err(FrameDecoderError::ReadFrameHeaderError(
4594 ReadFrameHeaderError::BadMagicNumber(_) | ReadFrameHeaderError::SkipFrame { .. },
4595 )) => {}
4596 other => panic!(
4597 "standard decoder must reject a magicless frame with \
4598 ReadFrameHeaderError::BadMagicNumber or SkipFrame, got {other:?}",
4599 ),
4600 }
4601 }
4602
4603 /// A reused `FrameCompressor` must emit byte-identical frames to a
4604 /// fresh compressor per input across both the borrowed (Fast) and
4605 /// owned (Dfast/Lazy/Greedy/Uncompressed) backends. This proves
4606 /// `prepare_frame` fully resets the per-frame state (matcher window,
4607 /// content hasher, FSE/Huffman seeds) between independent frames; a
4608 /// missed reset would corrupt frame N>=2's header checksum or matches.
4609 /// Each emitted frame must also round-trip.
4610 #[test]
4611 fn compress_independent_frame_reuse_matches_fresh_and_roundtrips() {
4612 use crate::encoding::{CompressionLevel, compress_slice_to_vec};
4613 let levels = [
4614 CompressionLevel::Uncompressed,
4615 CompressionLevel::Fastest,
4616 CompressionLevel::Default,
4617 CompressionLevel::Better,
4618 CompressionLevel::Best,
4619 CompressionLevel::Level(5),
4620 ];
4621 let inputs: Vec<Vec<u8>> = vec![
4622 Vec::new(),
4623 vec![0x00],
4624 b"the quick brown fox jumps over the lazy dog\n".to_vec(),
4625 vec![0x7Eu8; 50_000], // highly compressible
4626 generate_data(0xABCD, 70_000), // pseudo-random
4627 generate_data(0x1234, 200_000),
4628 ];
4629 for level in levels {
4630 let mut cctx: FrameCompressor = FrameCompressor::new(level);
4631 for data in &inputs {
4632 let reused = cctx.compress_independent_frame(data);
4633 let fresh = compress_slice_to_vec(data, level);
4634 assert_eq!(
4635 reused,
4636 fresh,
4637 "reused frame != fresh frame for len={} level={:?}",
4638 data.len(),
4639 level,
4640 );
4641 let mut decoder = FrameDecoder::new();
4642 let mut decoded = Vec::with_capacity(data.len());
4643 decoder.decode_all_to_vec(&reused, &mut decoded).unwrap();
4644 assert_eq!(
4645 decoded,
4646 *data,
4647 "roundtrip failed for len={} level={:?}",
4648 data.len(),
4649 level,
4650 );
4651 }
4652 }
4653 }
4654
4655 /// `compress_independent_frame_into` must replace (not append to) the
4656 /// caller's buffer each call, so a smaller frame after a larger one
4657 /// yields exactly the smaller frame, and the reused buffer's content
4658 /// matches a fresh compression of the same input.
4659 #[test]
4660 fn compress_independent_frame_into_replaces_buffer_contents() {
4661 use crate::encoding::{CompressionLevel, compress_slice_to_vec};
4662 let large = vec![0x11u8; 40_000];
4663 let small = b"short payload".to_vec();
4664 let mut cctx: FrameCompressor = FrameCompressor::new(CompressionLevel::Default);
4665 let mut out = Vec::new();
4666 cctx.compress_independent_frame_into(&large, &mut out);
4667 let frame_large = out.clone();
4668 // Reusing the same buffer for a smaller frame must clear it first.
4669 cctx.compress_independent_frame_into(&small, &mut out);
4670 assert_eq!(
4671 out,
4672 compress_slice_to_vec(&small, CompressionLevel::Default),
4673 "reused buffer must hold exactly the second frame",
4674 );
4675 // The first frame, captured before reuse, still round-trips.
4676 let mut decoder = FrameDecoder::new();
4677 let mut decoded = Vec::with_capacity(large.len());
4678 decoder
4679 .decode_all_to_vec(&frame_large, &mut decoded)
4680 .unwrap();
4681 assert_eq!(decoded, large);
4682 }
4683
4684 /// A sticky dictionary set once on a reused compressor must be primed
4685 /// into every independent frame (mirroring `ZSTD_CCtx_loadDictionary`):
4686 /// each frame decodes with the dictionary and is byte-identical to a
4687 /// fresh compressor carrying the same dictionary. This proves
4688 /// `prepare_frame` re-primes the dictionary (matcher content + offset
4689 /// history + entropy seed) every call rather than only on the first.
4690 #[test]
4691 fn compress_independent_frame_reuses_sticky_dictionary() {
4692 use crate::encoding::CompressionLevel;
4693 let dict_raw = include_bytes!("../../dict_tests/dictionary");
4694 let dict_content = crate::decoding::Dictionary::decode_dict(dict_raw).unwrap();
4695 let mut payload_a = Vec::new();
4696 for _ in 0..8 {
4697 payload_a.extend_from_slice(&dict_content.dict_content[..2048]);
4698 }
4699 let payload_b = b"a different second frame payload, still dict-attached".to_vec();
4700 let inputs = [payload_a, payload_b];
4701
4702 let mut cctx: FrameCompressor = FrameCompressor::new(CompressionLevel::Fastest);
4703 cctx.set_dictionary_from_bytes(dict_raw)
4704 .expect("dictionary bytes should parse");
4705
4706 for data in &inputs {
4707 let reused = cctx.compress_independent_frame(data);
4708 // Fresh compressor carrying the same sticky dictionary.
4709 let mut fresh_enc: FrameCompressor = FrameCompressor::new(CompressionLevel::Fastest);
4710 fresh_enc
4711 .set_dictionary_from_bytes(dict_raw)
4712 .expect("dictionary bytes should parse");
4713 let fresh = fresh_enc.compress_independent_frame(data);
4714 assert_eq!(
4715 reused,
4716 fresh,
4717 "reused dict frame != fresh dict frame, len={}",
4718 data.len(),
4719 );
4720 // Round-trip with the dictionary on the decode side.
4721 let dict_for_decoder = crate::decoding::Dictionary::decode_dict(dict_raw).unwrap();
4722 let mut decoder = FrameDecoder::new();
4723 decoder.add_dict(dict_for_decoder).unwrap();
4724 let mut decoded = Vec::with_capacity(data.len());
4725 decoder.decode_all_to_vec(&reused, &mut decoded).unwrap();
4726 assert_eq!(&decoded, data, "dict roundtrip failed, len={}", data.len());
4727 }
4728 }
4729}