structured_zstd/encoding/streaming_encoder.rs
1use alloc::format;
2use alloc::string::{String, ToString};
3use alloc::vec::Vec;
4use core::borrow::BorrowMut;
5use core::marker::PhantomData;
6
7use crate::common::MAX_BLOCK_SIZE;
8#[cfg(feature = "hash")]
9use core::hash::Hasher;
10#[cfg(feature = "hash")]
11use twox_hash::XxHash64;
12
13use crate::encoding::fastpath::EncodeTier;
14use crate::encoding::levels::compress_block_encoded;
15use crate::encoding::{
16 CompressionLevel, EncoderDictionary, MatchGeneratorDriver, Matcher, block_header::BlockHeader,
17 frame_compressor::CompressState, frame_compressor::FrameTuning, frame_compressor::FseTables,
18 frame_compressor::PreviousFseTable, frame_header::FrameHeader,
19};
20use crate::io::{Error, ErrorKind, Write};
21
22/// Incremental frame encoder that implements [`Write`].
23///
24/// Data can be provided with multiple `write()` calls. Full blocks are compressed
25/// automatically, `flush()` emits the currently buffered partial block as non-last,
26/// and `finish()` closes the frame and returns the wrapped writer.
27///
28/// One encoder writes one frame into the drain it owns, through a
29/// [`CompressionContext`] it owns ([`new`](StreamingEncoder::new)) or borrows
30/// ([`with_context`](Self::with_context)). Borrowing is how frame after frame
31/// is compressed with the same settings, dictionary and match-finder
32/// allocations: the context outlives each encoder and is ready for the next
33/// frame once [`finish`](Self::finish) returns.
34pub struct StreamingEncoder<
35 W: Write,
36 M: Matcher = MatchGeneratorDriver,
37 C: BorrowMut<CompressionContext<M>> = CompressionContext<M>,
38> {
39 drain: Option<W>,
40 context: C,
41 matcher: PhantomData<M>,
42}
43
44/// A reusable streaming compression context: the settings, the attached
45/// dictionary, the match finder and its buffers, kept from one frame to the
46/// next, with the output handed in on each call rather than owned. The
47/// counterpart of upstream zstd's `ZSTD_CCtx` driven by `ZSTD_compressStream2`.
48///
49/// Settings apply from the next frame on and must be made before its first
50/// [`write`](Self::write); [`finish_frame`](Self::finish_frame) closes the
51/// frame and readies the context for another. A pledged size belongs to one
52/// frame; every other setting, the dictionary included, stays until replaced.
53/// A frame that fails leaves the context failed: every later call reports
54/// that failure, and a new context is needed.
55///
56/// Reusing a context produces the same frames as a fresh
57/// [`StreamingEncoder`] per frame, without rebuilding the match finder's
58/// tables or re-attaching the dictionary for each of them.
59///
60/// # Examples
61/// ```
62/// use structured_zstd::encoding::{CompressionContext, CompressionLevel};
63///
64/// let mut context = CompressionContext::new(CompressionLevel::Default);
65/// let mut frames = Vec::new();
66/// for payload in [&b"first frame"[..], b"second frame"] {
67/// let mut frame = Vec::new();
68/// context.set_pledged_content_size(payload.len() as u64).unwrap();
69/// context.write(&mut frame, payload).unwrap();
70/// context.finish_frame(&mut frame).unwrap();
71/// frames.push(frame);
72/// }
73/// use std::io::Read;
74/// let mut decoder = structured_zstd::decoding::StreamingDecoder::new(&frames[1][..]);
75/// let mut decoded = Vec::new();
76/// decoder.read_to_end(&mut decoded).unwrap();
77/// assert_eq!(decoded, b"second frame");
78/// ```
79pub struct CompressionContext<M: Matcher = MatchGeneratorDriver> {
80 compression_level: CompressionLevel,
81 state: CompressState<M>,
82 /// The block being written out. An uncompressed frame also buffers its
83 /// input here, behind the header, so a raw block is assembled where it
84 /// is written from; a compressed frame buffers input in the match
85 /// finder's history, uncommitted until its block is cut.
86 encoded_scratch: Vec<u8>,
87 errored: bool,
88 last_error_kind: Option<ErrorKind>,
89 last_error_message: Option<String>,
90 frame_started: bool,
91 /// Upper bound on emitted block sizes (upstream `ZSTD_c_targetCBlockSize`
92 /// semantics; see `FrameCompressor::set_target_block_size`). `None` =
93 /// the format's 128 KiB ceiling.
94 target_block_size: Option<u32>,
95 /// Largest block a frame may carry (upstream `ZSTD_c_maxBlockSize`); see
96 /// [`Self::set_max_block_size`].
97 max_block_size: Option<u32>,
98 /// The pledged size of the frame in progress; cleared when it ends.
99 pledged_content_size: Option<u64>,
100 /// Advisory source-size hint from [`set_source_size_hint`](Self::set_source_size_hint).
101 /// Unlike `pledged_content_size` it carries no end-of-frame enforcement, but
102 /// it still feeds the small-input gates (matcher sizing AND the Fast HUF
103 /// fast-path gate) so `set_source_size_hint(small)` reduces work the same way
104 /// a pledge does. The HUF gate reads `pledged_content_size.or(source_size_hint)`.
105 /// A parameter like upstream `ZSTD_c_srcSizeHint`, so it outlives a frame.
106 source_size_hint: Option<u64>,
107 /// Whether a pledged size is written into the header's
108 /// `Frame_Content_Size` field (upstream `ZSTD_c_contentSizeFlag`).
109 /// Pledge *enforcement* is independent of this flag — upstream
110 /// validates consumed bytes against the pledge at frame end even
111 /// when the header omits the field. Default `true`.
112 content_size_flag: bool,
113 bytes_consumed: u64,
114 /// Upstream `ZSTD_compress_frameChunk` `savings`: bytes consumed minus
115 /// bytes produced so far in this frame; the block pre-splitter only cuts
116 /// full blocks once the frame has saved enough.
117 savings: i64,
118 tuning: FrameTuning,
119 /// `ZSTD_f_zstd1_magicless` — omit the 4-byte magic number prefix.
120 /// Default false. See [`Self::set_magicless`].
121 magicless: bool,
122 /// Whether to emit a trailing XXH64 content checksum and set the frame
123 /// header's `Content_Checksum_flag` (upstream `ZSTD_c_checksumFlag`).
124 /// Default `false`, matching the upstream library default; combined with
125 /// the `hash` feature, so without `hash` no checksum is emitted
126 /// regardless. See [`Self::set_content_checksum`].
127 content_checksum: bool,
128 /// Dictionary applied to each frame (upstream zstd `ZSTD_CCtx_loadDictionary`
129 /// on a streaming context), with the entropy tables it seeds. `None` = no
130 /// dictionary. Set before a frame's first write.
131 dictionary: Option<EncoderDictionary>,
132 /// Whether the frame header records the attached dictionary's ID
133 /// (upstream `ZSTD_c_dictIDFlag`). Default `true`. Raw-content
134 /// dictionaries (upstream `ZSTD_CCtx_refPrefix`) carry a synthetic
135 /// non-zero ID that must not reach the wire, so their attach path
136 /// turns this off. See [`Self::set_dictionary_id_flag`].
137 dictionary_id_flag: bool,
138 #[cfg(feature = "hash")]
139 hasher: XxHash64,
140}
141
142impl<W: Write> StreamingEncoder<W, MatchGeneratorDriver> {
143 /// Creates a streaming encoder backed by the default match generator.
144 ///
145 /// The encoder writes compressed bytes into `drain` and applies `compression_level`
146 /// to all subsequently written blocks.
147 pub fn new(drain: W, compression_level: CompressionLevel) -> Self {
148 Self::with_context(drain, CompressionContext::new(compression_level))
149 }
150}
151
152impl<W: Write, C: BorrowMut<CompressionContext>> StreamingEncoder<W, MatchGeneratorDriver, C> {
153 /// Configure fine-grained compression parameters; see
154 /// [`CompressionContext::set_parameters`]. Must be called before the first
155 /// [`write`](Write::write).
156 pub fn set_parameters(
157 &mut self,
158 params: &crate::encoding::CompressionParameters,
159 ) -> Result<(), Error> {
160 self.context.borrow_mut().set_parameters(params)
161 }
162}
163
164impl<W: Write, M: Matcher> StreamingEncoder<W, M> {
165 /// Creates a streaming encoder with an explicitly provided matcher implementation.
166 ///
167 /// This constructor is primarily intended for tests and advanced callers that need
168 /// custom match-window behavior.
169 pub fn new_with_matcher(matcher: M, drain: W, compression_level: CompressionLevel) -> Self {
170 Self::with_context(
171 drain,
172 CompressionContext::new_with_matcher(matcher, compression_level),
173 )
174 }
175}
176
177impl<W: Write, M: Matcher, C: BorrowMut<CompressionContext<M>>> StreamingEncoder<W, M, C> {
178 /// Write one frame into `drain` through `context`: owned, or borrowed from
179 /// a caller that keeps it for the next frame with every setting, the
180 /// dictionary and the allocations it has.
181 ///
182 /// # Examples
183 /// ```
184 /// use std::io::Write;
185 /// use structured_zstd::encoding::{CompressionContext, CompressionLevel, StreamingEncoder};
186 ///
187 /// let mut context = CompressionContext::new(CompressionLevel::Default);
188 /// for payload in [&b"first frame"[..], b"second frame"] {
189 /// let mut encoder = StreamingEncoder::with_context(Vec::new(), &mut context);
190 /// encoder.write_all(payload).unwrap();
191 /// let frame = encoder.finish().unwrap();
192 /// assert!(!frame.is_empty());
193 /// }
194 /// ```
195 pub fn with_context(drain: W, context: C) -> Self {
196 Self {
197 drain: Some(drain),
198 context,
199 matcher: PhantomData,
200 }
201 }
202
203 /// Bound each block's payload; see
204 /// [`CompressionContext::set_target_block_size`]. Must be set before the
205 /// first write.
206 pub fn set_target_block_size(&mut self, target: Option<u32>) -> Result<(), Error> {
207 self.context.borrow_mut().set_target_block_size(target)
208 }
209
210 /// Cap each block's size; see [`CompressionContext::set_max_block_size`].
211 /// Must be set before the first write.
212 pub fn set_max_block_size(&mut self, size: Option<u32>) -> Result<(), Error> {
213 self.context.borrow_mut().set_max_block_size(size)
214 }
215
216 /// Enable or disable the trailing XXH64 content checksum; see
217 /// [`CompressionContext::set_content_checksum`]. Must be called before the
218 /// first write.
219 pub fn set_content_checksum(&mut self, emit: bool) -> Result<(), Error> {
220 self.context.borrow_mut().set_content_checksum(emit)
221 }
222
223 /// Enable or disable the magicless frame format; see
224 /// [`CompressionContext::set_magicless`]. Must be called before the first
225 /// write.
226 pub fn set_magicless(&mut self, magicless: bool) -> Result<(), Error> {
227 self.context.borrow_mut().set_magicless(magicless)
228 }
229
230 /// Pledge the total uncompressed content size of the frame; see
231 /// [`CompressionContext::set_pledged_content_size`]. Must be called before
232 /// the first write.
233 pub fn set_pledged_content_size(&mut self, size: u64) -> Result<(), Error> {
234 self.context.borrow_mut().set_pledged_content_size(size)
235 }
236
237 /// Control whether a pledged size reaches the header; see
238 /// [`CompressionContext::set_content_size_flag`]. Must be called before
239 /// the first write.
240 pub fn set_content_size_flag(&mut self, emit: bool) -> Result<(), Error> {
241 self.context.borrow_mut().set_content_size_flag(emit)
242 }
243
244 /// Provide an advisory size for the frame; see
245 /// [`CompressionContext::set_source_size_hint`]. Must be called before the
246 /// first write.
247 pub fn set_source_size_hint(&mut self, size: u64) -> Result<(), Error> {
248 self.context.borrow_mut().set_source_size_hint(size)
249 }
250
251 /// Attach a dictionary blob to the frame; see
252 /// [`CompressionContext::set_dictionary_from_bytes`]. Must be called before
253 /// the first write.
254 pub fn set_dictionary_from_bytes(&mut self, raw_dictionary: &[u8]) -> Result<(), Error> {
255 self.context
256 .borrow_mut()
257 .set_dictionary_from_bytes(raw_dictionary)
258 }
259
260 /// Whether the header records the dictionary ID; see
261 /// [`CompressionContext::set_dictionary_id_flag`]. Must be set before the
262 /// first write.
263 pub fn set_dictionary_id_flag(&mut self, emit: bool) -> Result<(), Error> {
264 self.context.borrow_mut().set_dictionary_id_flag(emit)
265 }
266
267 /// Attach an already-parsed [`EncoderDictionary`] to the frame; see
268 /// [`CompressionContext::set_encoder_dictionary`]. Must be called before
269 /// the first write.
270 pub fn set_encoder_dictionary(&mut self, dict: EncoderDictionary) -> Result<(), Error> {
271 self.context.borrow_mut().set_encoder_dictionary(dict)
272 }
273
274 /// Returns an immutable reference to the wrapped output drain.
275 ///
276 /// The drain remains available for the encoder lifetime; [`finish`](Self::finish)
277 /// consumes the encoder and returns ownership of the drain.
278 pub fn get_ref(&self) -> &W {
279 self.drain
280 .as_ref()
281 .expect("streaming encoder drain is present until finish consumes self")
282 }
283
284 /// Total heap bytes this encoder's allocations hold, excluding the inline
285 /// struct and the drain `W` (whose footprint the owner can measure through
286 /// [`get_ref`](Self::get_ref)); see [`CompressionContext::heap_size`].
287 pub fn heap_size(&self) -> usize {
288 self.context.borrow().heap_size()
289 }
290
291 /// Returns a mutable reference to the wrapped output drain.
292 ///
293 /// It is inadvisable to directly write to the underlying writer, as doing
294 /// so would corrupt the zstd frame being assembled by the encoder.
295 ///
296 /// The drain remains available for the encoder lifetime; [`finish`](Self::finish)
297 /// consumes the encoder and returns ownership of the drain.
298 pub fn get_mut(&mut self) -> &mut W {
299 self.drain
300 .as_mut()
301 .expect("streaming encoder drain is present until finish consumes self")
302 }
303
304 /// Finalizes the current zstd frame and returns the wrapped output drain.
305 ///
306 /// If no payload was written yet, this still emits a valid empty frame.
307 /// Calling this method consumes the encoder; a borrowed context is then
308 /// ready for the next frame, also when this fails: the frame goes with the
309 /// drain (see [`CompressionContext::abandon_frame`]).
310 pub fn finish(mut self) -> Result<W, Error> {
311 let mut drain = self
312 .drain
313 .take()
314 .expect("streaming encoder drain must be present when finishing");
315 self.context.borrow_mut().finish_frame(&mut drain)?;
316 Ok(drain)
317 }
318
319 fn drain_mut(&mut self) -> (&mut W, &mut CompressionContext<M>) {
320 let drain = self
321 .drain
322 .as_mut()
323 .expect("streaming encoder drain is present until finish consumes self");
324 (drain, self.context.borrow_mut())
325 }
326}
327
328/// The frame belongs to the drain this encoder writes into: an encoder that
329/// goes without finishing it (dropped mid-frame, or a `finish` that failed)
330/// takes it along, so a borrowed context starts the next encoder's frame
331/// afresh instead of continuing this one into another drain.
332impl<W: Write, M: Matcher, C: BorrowMut<CompressionContext<M>>> Drop for StreamingEncoder<W, M, C> {
333 fn drop(&mut self) {
334 self.context.borrow_mut().abandon_frame();
335 }
336}
337
338impl<W: Write, M: Matcher, C: BorrowMut<CompressionContext<M>>> Write
339 for StreamingEncoder<W, M, C>
340{
341 fn write(&mut self, buf: &[u8]) -> Result<usize, Error> {
342 let (drain, context) = self.drain_mut();
343 context.write(drain, buf)
344 }
345
346 fn flush(&mut self) -> Result<(), Error> {
347 let (drain, context) = self.drain_mut();
348 context.flush(drain)
349 }
350}
351
352impl CompressionContext<MatchGeneratorDriver> {
353 /// Creates a context backed by the default match generator, compressing
354 /// at `compression_level`.
355 pub fn new(compression_level: CompressionLevel) -> Self {
356 Self::new_with_matcher(
357 MatchGeneratorDriver::new(MAX_BLOCK_SIZE as usize, 1),
358 compression_level,
359 )
360 }
361
362 /// Configure fine-grained compression parameters (#27): resets the level to
363 /// the parameters' level and installs the per-knob overrides (window / hash
364 /// / chain / search logs, strategy, long-distance matching) applied at the
365 /// next frame. Mirrors [`FrameCompressor::set_parameters`](crate::encoding::FrameCompressor::set_parameters).
366 /// Must be called before the frame's first [`write`](Self::write). Only the
367 /// built-in `MatchGeneratorDriver` exposes the override knobs, so this
368 /// lives on the default-matcher impl.
369 pub fn set_parameters(
370 &mut self,
371 params: &crate::encoding::CompressionParameters,
372 ) -> Result<(), Error> {
373 self.ensure_settable("compression parameters must be set before the first write")?;
374 self.compression_level = params.level();
375 let overrides = params.overrides();
376 // Persist the strategy override so `ensure_frame_started`'s level-based
377 // resync does not discard it (matching `FrameCompressor::set_parameters`).
378 self.tuning = FrameTuning::from_overrides(&overrides);
379 self.state.strategy_tag = self.tuning.strategy.map_or_else(
380 || {
381 crate::encoding::strategy::StrategyTag::for_compression_level(
382 self.compression_level,
383 )
384 },
385 |(tag, _)| tag,
386 );
387 self.state.huf_optimal_search = crate::encoding::frame_compressor::huf_search_enabled(
388 self.state.strategy_tag,
389 self.pledged_content_size.or(self.source_size_hint),
390 );
391 self.state.matcher.set_param_overrides(Some(overrides));
392 Ok(())
393 }
394
395 /// Ask `producer` for each block's parse in place of the built-in
396 /// matchers; `None` removes it. With `fallback` a block the producer fails
397 /// on is matched as usual; without, it goes out uncompressed and
398 /// [`Self::take_sequence_producer_failure`] reports it. With long-distance
399 /// matching on, the long matches are taken outright and the producer
400 /// parses the literals between them.
401 pub fn set_sequence_producer(
402 &mut self,
403 producer: Option<alloc::boxed::Box<dyn crate::encoding::sequences::SequenceProducer>>,
404 fallback: bool,
405 ) {
406 self.state.matcher.set_sequence_producer(producer, fallback);
407 }
408
409 /// Whether a block went out uncompressed because the producer failed on it
410 /// with no fallback, since this was last asked.
411 pub fn take_sequence_producer_failure(&mut self) -> bool {
412 self.state.matcher.take_producer_failure()
413 }
414}
415
416impl<M: Matcher> CompressionContext<M> {
417 /// Creates a context with an explicitly provided matcher implementation.
418 ///
419 /// This constructor is primarily intended for tests and advanced callers that need
420 /// custom match-window behavior.
421 pub fn new_with_matcher(matcher: M, compression_level: CompressionLevel) -> Self {
422 Self {
423 compression_level,
424 state: CompressState {
425 matcher,
426 sequence_capture: None,
427 copy_kernel: crate::encoding::fastpath::select_kernel(),
428 last_huff_table: None,
429 huff_table_spare: None,
430 huff_weights: Default::default(),
431 seen_content: Default::default(),
432 fse_tables: FseTables::new(),
433 block_scratch: crate::encoding::blocks::CompressedBlockScratch::new(),
434 workspace: crate::encoding::workspace::Workspace::new(),
435 offset_hist: [1, 4, 8],
436 strategy_tag: crate::encoding::strategy::StrategyTag::for_compression_level(
437 compression_level,
438 ),
439 pre_split: crate::encoding::levels::config::level_pre_split(compression_level)
440 .map(|tier| tier as u8),
441 huf_optimal_search: true,
442 literal_compression_disabled: matches!(
443 compression_level,
444 CompressionLevel::Level(n) if n < 0
445 ),
446 post_split_allowed: true,
447 },
448 encoded_scratch: Vec::new(),
449 errored: false,
450 last_error_kind: None,
451 last_error_message: None,
452 frame_started: false,
453 target_block_size: None,
454 max_block_size: None,
455 pledged_content_size: None,
456 source_size_hint: None,
457 content_size_flag: true,
458 bytes_consumed: 0,
459 savings: 0,
460 tuning: FrameTuning::default(),
461 magicless: false,
462 content_checksum: false,
463 dictionary: None,
464 dictionary_id_flag: true,
465 #[cfg(feature = "hash")]
466 hasher: XxHash64::with_seed(0),
467 }
468 }
469
470 /// Compress the next frames at `level`, with the level's own tuning: any
471 /// parameter override installed by
472 /// [`set_parameters`](CompressionContext::set_parameters) is dropped, as
473 /// [`FrameCompressor::set_compression_level`](crate::encoding::FrameCompressor::set_compression_level)
474 /// drops it. Must be called before the frame's first [`write`](Self::write).
475 ///
476 /// # Examples
477 /// ```
478 /// use structured_zstd::encoding::{CompressionContext, CompressionLevel};
479 ///
480 /// let mut context = CompressionContext::new(CompressionLevel::Fastest);
481 /// context.set_compression_level(CompressionLevel::Better).unwrap();
482 /// let mut frame = Vec::new();
483 /// context.write(&mut frame, b"compressed at the new level").unwrap();
484 /// context.finish_frame(&mut frame).unwrap();
485 /// ```
486 pub fn set_compression_level(&mut self, level: CompressionLevel) -> Result<(), Error> {
487 self.ensure_settable("the compression level must be set before the first write")?;
488 self.compression_level = level;
489 self.tuning = FrameTuning::default();
490 self.state.matcher.clear_param_overrides();
491 Ok(())
492 }
493
494 /// Set an upper bound on each physical block's payload (semantics of
495 /// upstream `ZSTD_c_targetCBlockSize`): every block carries at most
496 /// `target` payload bytes, +3-byte block header on the wire — the
497 /// upstream knob is likewise a convergence target for block sizing,
498 /// not a cap on header-inclusive wire bytes. Clamped to
499 /// `[MIN_TARGET_BLOCK_SIZE, MAX_BLOCK_SIZE]`; mirrors
500 /// `FrameCompressor::set_target_block_size`. Must be set before the
501 /// frame's first write.
502 pub fn set_target_block_size(&mut self, target: Option<u32>) -> Result<(), Error> {
503 self.ensure_settable("the block-size target must be set before the first write")?;
504 self.target_block_size = target.map(|t| {
505 t.clamp(
506 crate::common::MIN_TARGET_BLOCK_SIZE,
507 crate::common::MAX_BLOCK_SIZE,
508 )
509 });
510 Ok(())
511 }
512
513 /// Cap every block at `size` bytes (upstream `ZSTD_c_maxBlockSize`),
514 /// clamped to `[MIN_MAX_BLOCK_SIZE, MAX_BLOCK_SIZE]`; mirrors
515 /// `FrameCompressor::set_max_block_size`. Must be set before the frame's
516 /// first write.
517 pub fn set_max_block_size(&mut self, size: Option<u32>) -> Result<(), Error> {
518 self.ensure_settable("the maximum block size must be set before the first write")?;
519 self.max_block_size = size.map(|s| {
520 s.clamp(
521 crate::common::MIN_MAX_BLOCK_SIZE,
522 crate::common::MAX_BLOCK_SIZE,
523 )
524 });
525 Ok(())
526 }
527
528 /// Enable or disable the trailing XXH64 content checksum
529 /// (upstream `ZSTD_c_checksumFlag`). Default `false`, matching the
530 /// upstream library default (`ZSTD_c_checksumFlag = 0`). Must be called
531 /// before the frame's first [`write`](Self::write); once the frame header
532 /// is emitted the flag is fixed, so a late change returns an error rather
533 /// than producing a header/trailer mismatch. Without the `hash` feature
534 /// no checksum is emitted regardless.
535 pub fn set_content_checksum(&mut self, emit: bool) -> Result<(), Error> {
536 self.ensure_settable("content checksum must be set before the first write")?;
537 self.content_checksum = emit;
538 Ok(())
539 }
540
541 /// Enable or disable magicless frame format (`ZSTD_f_zstd1_magicless`).
542 ///
543 /// When set to `true`, the frame header omits the 4-byte magic number
544 /// prefix. Must be called BEFORE the frame's first [`write`](Self::write)
545 /// call; calling it after the frame header has already been emitted
546 /// returns an error so the caller can't be misled into thinking they
547 /// produced a magicless stream.
548 pub fn set_magicless(&mut self, magicless: bool) -> Result<(), Error> {
549 self.ensure_settable("magicless format must be set before the first write")?;
550 self.magicless = magicless;
551 Ok(())
552 }
553
554 /// Pledge the total uncompressed content size of the next frame.
555 ///
556 /// When set, the frame header will include a `Frame_Content_Size` field.
557 /// This enables decoders to pre-allocate output buffers.
558 /// The pledged size is also forwarded as a source-size hint to the
559 /// matcher so small inputs can use smaller matching tables.
560 ///
561 /// Must be called **before** the frame's first [`write`](Self::write);
562 /// calling it after the frame header has already been emitted returns an
563 /// error. The pledge ends with the frame.
564 pub fn set_pledged_content_size(&mut self, size: u64) -> Result<(), Error> {
565 self.ensure_settable("pledged content size must be set before the first write")?;
566 // The matcher is handed it as its size hint when the frame starts
567 // (`ensure_frame_started`), not here: a pledge that goes unused must
568 // not stay behind in the matcher for a later frame.
569 self.pledged_content_size = Some(size);
570 Ok(())
571 }
572
573 /// Control whether the pledged size is written into the header's
574 /// `Frame_Content_Size` field (upstream `ZSTD_c_contentSizeFlag`,
575 /// default on). With the flag off the header omits the field, but a
576 /// pledge set via [`set_pledged_content_size`](Self::set_pledged_content_size)
577 /// is still enforced against the bytes actually written. Must be
578 /// called before the frame's first [`write`](Self::write).
579 pub fn set_content_size_flag(&mut self, emit: bool) -> Result<(), Error> {
580 self.ensure_settable("content size flag must be set before the first write")?;
581 self.content_size_flag = emit;
582 Ok(())
583 }
584
585 /// Provide a hint about the total uncompressed size of each frame.
586 ///
587 /// Unlike [`set_pledged_content_size`](Self::set_pledged_content_size),
588 /// this does **not** enforce that exactly `size` bytes are written; it
589 /// may reduce matcher tables, advertised frame window, and block sizing
590 /// for small inputs. A parameter, like upstream `ZSTD_c_srcSizeHint`: it
591 /// applies to every frame until replaced. Must be called before the
592 /// frame's first [`write`](Self::write).
593 pub fn set_source_size_hint(&mut self, size: u64) -> Result<(), Error> {
594 self.ensure_settable("source size hint must be set before the first write")?;
595 // Read at each frame start as `pledged_content_size.or(source_size_hint)`,
596 // by the matcher's sizing and the Fast HUF fast-path gate alike
597 // (`ensure_frame_started`), so a small advisory size also lifts Fast
598 // streams off the expensive optimal-HUF search.
599 self.source_size_hint = Some(size);
600 Ok(())
601 }
602
603 /// Drop the hint [`set_source_size_hint`](Self::set_source_size_hint)
604 /// set, so later frames are sized by their pledge alone. Must be called
605 /// before the frame's first [`write`](Self::write).
606 pub fn clear_source_size_hint(&mut self) -> Result<(), Error> {
607 self.ensure_settable("source size hint must be cleared before the first write")?;
608 self.source_size_hint = None;
609 Ok(())
610 }
611
612 /// Attach a dictionary blob to each frame (upstream zstd
613 /// `ZSTD_CCtx_loadDictionary` on a streaming context, which loads in
614 /// `ZSTD_dct_auto` mode): a
615 /// [serialized](crate::decoding::Dictionary::is_serialized) dictionary,
616 /// or raw content. The dictionary primes the
617 /// match-finder and seeds the first block's entropy tables + repeat
618 /// offsets; a serialized one's ID is written into the frame header, while
619 /// raw content has none to write, so the decoder must be given the same
620 /// bytes explicitly. Must be called before the frame's first
621 /// [`write`](Self::write); repeat offsets must be non-zero.
622 pub fn set_dictionary_from_bytes(&mut self, raw_dictionary: &[u8]) -> Result<(), Error> {
623 if raw_dictionary.is_empty() {
624 // An empty buffer is how the same upstream entry point is told
625 // there is no dictionary: it clears and succeeds. Still refused
626 // once the frame is open, like any other attach.
627 self.ensure_settable("dictionary must be attached before the first write")?;
628 // What the match finder kept of it across frames (its primed
629 // snapshot, the copy left resident) goes with it, entropy included.
630 self.state.matcher.invalidate_primed_dictionary();
631 self.dictionary = None;
632 return Ok(());
633 }
634 let dict = EncoderDictionary::from_serialized_or_raw_content(raw_dictionary)
635 .map_err(|err| invalid_input_error(&alloc::format!("invalid dictionary: {err:?}")))?;
636 self.set_encoder_dictionary(dict)
637 }
638
639 /// Whether the frame header records the attached dictionary's ID
640 /// (upstream `ZSTD_c_dictIDFlag` semantics; default `true`).
641 /// Mirrors [`FrameCompressor::set_dictionary_id_flag`](crate::encoding::FrameCompressor::set_dictionary_id_flag).
642 /// Decoders can still decode such frames by supplying the dictionary
643 /// explicitly.
644 pub fn set_dictionary_id_flag(&mut self, emit: bool) -> Result<(), Error> {
645 self.ensure_settable("dictionary ID flag must be set before the first write")?;
646 self.dictionary_id_flag = emit;
647 Ok(())
648 }
649
650 /// Attach an already-parsed [`EncoderDictionary`] to each frame. See
651 /// [`set_dictionary_from_bytes`](Self::set_dictionary_from_bytes); must be
652 /// called before the frame's first write. The entropy tables it seeds were
653 /// built when it was prepared, so attaching builds nothing.
654 pub fn set_encoder_dictionary(&mut self, dict: EncoderDictionary) -> Result<(), Error> {
655 self.ensure_settable("dictionary must be attached before the first write")?;
656 // A zero id marks a raw-content dictionary, which carries no header to
657 // hold one; the frame then records no dictionary ID and the decoder
658 // must be given the same bytes explicitly.
659 if dict.inner.offset_hist.contains(&0) {
660 return Err(invalid_input_error(
661 "dictionary carries a zero repeat offset",
662 ));
663 }
664 // The match finder's primed snapshot and resident copy belong to the
665 // dictionary being replaced; the next frame primes the new one.
666 self.state.matcher.invalidate_primed_dictionary();
667 // One under 8 bytes is attached as none; see
668 // `EncoderDictionary::is_used_for_compression`.
669 self.dictionary = dict.is_used_for_compression().then_some(dict);
670 Ok(())
671 }
672
673 /// The dictionary each frame is compressed with, if one is attached.
674 ///
675 /// # Examples
676 /// ```
677 /// use structured_zstd::encoding::{CompressionContext, CompressionLevel, EncoderDictionary};
678 ///
679 /// let dictionary = EncoderDictionary::from_serialized_or_raw_content(b"some shared history").unwrap();
680 /// let mut context = CompressionContext::new(CompressionLevel::Default);
681 /// assert!(context.dictionary().is_none());
682 /// context.set_encoder_dictionary(dictionary).unwrap();
683 /// assert!(context.dictionary().is_some());
684 /// ```
685 pub fn dictionary(&self) -> Option<&EncoderDictionary> {
686 self.dictionary.as_ref()
687 }
688
689 /// Total heap bytes this context's allocations hold, excluding the inline
690 /// struct: match-finder tables / history, retained Huffman tables, the
691 /// `encoded_scratch` block buffer, the retained dictionary content, and
692 /// its entropy tables. Mirrors
693 /// `FrameCompressor::heap_size` so a context can report its true
694 /// footprint through `ZSTD_sizeof_CCtx`.
695 pub fn heap_size(&self) -> usize {
696 let mut total = self.state.matcher.heap_size();
697 total += self
698 .state
699 .last_huff_table
700 .as_ref()
701 .map_or(0, |table| table.heap_size());
702 total += self
703 .state
704 .huff_table_spare
705 .as_ref()
706 .map_or(0, |table| table.heap_size());
707 // Kept between blocks and frames; see `FrameCompressor::heap_size`.
708 total += self.state.huff_weights.heap_size();
709 total += self.state.retained_scratch_heap_size();
710 total += self.state.seen_content.heap_size();
711 total += self.state.captured_parse_heap_size();
712 total += self.encoded_scratch.capacity();
713 total += self
714 .dictionary
715 .as_ref()
716 .map_or(0, EncoderDictionary::heap_size);
717 total
718 }
719
720 /// Compress `buf` into the frame in progress, starting one (and writing
721 /// its header to `drain`) if none is. Full blocks are compressed as they
722 /// fill and written to `drain`; the rest stays buffered for the next call,
723 /// [`flush`](Self::flush) or [`finish_frame`](Self::finish_frame).
724 ///
725 /// Returns how much of `buf` was taken, which is all of it unless a pledge
726 /// set with [`set_pledged_content_size`](Self::set_pledged_content_size)
727 /// allows less.
728 pub fn write<D: Write + ?Sized>(&mut self, drain: &mut D, buf: &[u8]) -> Result<usize, Error> {
729 // The call's one choice of kernel tier: the frame start and every block
730 // it emits run compiled for it.
731 crate::encoding::fastpath::with_encode_tier!(
732 self.state.copy_kernel,
733 K => self.write_with::<D, K>(drain, buf)
734 )
735 }
736
737 /// [`Self::write`] compiled for the kernel tier `K`.
738 fn write_with<D: Write + ?Sized, K: EncodeTier>(
739 &mut self,
740 drain: &mut D,
741 buf: &[u8],
742 ) -> Result<usize, Error> {
743 self.ensure_open()?;
744 if buf.is_empty() {
745 return Ok(0);
746 }
747
748 // Check pledge before emitting the frame header so that a misuse
749 // like set_pledged_content_size(0) + write(non_empty) doesn't leave
750 // a partially-written header in the drain.
751 if let Some(pledged) = self.pledged_content_size
752 && self.bytes_consumed >= pledged
753 {
754 return Err(invalid_input_error(
755 "write would exceed pledged content size",
756 ));
757 }
758
759 self.ensure_frame_started::<D, K>(drain)?;
760
761 // Enforce pledged upper bound: truncate the accepted slice to the
762 // remaining allowance so that partial-write semantics are honored
763 // (return Ok(n) with n < buf.len()) instead of failing the full call.
764 // The check above leaves at least one byte of it.
765 let buf = if let Some(pledged) = self.pledged_content_size {
766 debug_assert!(self.bytes_consumed < pledged);
767 let remaining_allowed = pledged - self.bytes_consumed;
768 let accepted = core::cmp::min(
769 buf.len(),
770 usize::try_from(remaining_allowed).unwrap_or(usize::MAX),
771 );
772 &buf[..accepted]
773 } else {
774 buf
775 };
776
777 let block_capacity = self.block_capacity();
778 let mut remaining = buf;
779 let mut consumed = 0usize;
780
781 while !remaining.is_empty() {
782 // A block is emitted the moment it fills, below, and the
783 // capacity is fixed for the frame, so there is always room here.
784 let pending = self.pending_len();
785 debug_assert!(pending < block_capacity);
786 let to_take = core::cmp::min(remaining.len(), block_capacity - pending);
787 self.buffer_input(&remaining[..to_take]);
788 remaining = &remaining[to_take..];
789 consumed += to_take;
790
791 if self
792 .emit_full_pending_block::<D, K>(drain, block_capacity)
793 .is_err()
794 {
795 // The bytes this call took are reported; the failure, sticky
796 // on the context now, is what the next call returns.
797 break;
798 }
799 }
800 self.bytes_consumed += consumed as u64;
801 Ok(consumed)
802 }
803
804 /// Emit the buffered partial block as a non-last block and flush `drain`.
805 pub fn flush<D: Write + ?Sized>(&mut self, drain: &mut D) -> Result<(), Error> {
806 self.ensure_open()?;
807 if self.pending_len() == 0 {
808 return drain.flush().map_err(|err| self.fail(err));
809 }
810 crate::encoding::fastpath::with_encode_tier!(
811 self.state.copy_kernel,
812 K => {
813 self.ensure_frame_started::<D, K>(drain)?;
814 self.emit_pending_block::<D, K>(drain, false)
815 }
816 )?;
817 drain.flush().map_err(|err| self.fail(err))
818 }
819
820 /// Close the frame in progress into `drain`: its last block, then its
821 /// checksum when enabled. A frame nothing was written to is still a valid
822 /// empty frame. The context is then ready for the next frame, with every
823 /// setting and the dictionary as they were and the pledge cleared.
824 pub fn finish_frame<D: Write + ?Sized>(&mut self, drain: &mut D) -> Result<(), Error> {
825 self.ensure_open()?;
826
827 // Validate the pledge before finalizing the frame. If this is called
828 // before any writes, this also avoids emitting a header with an
829 // incorrect FCS into the drain on mismatch.
830 if let Some(pledged) = self.pledged_content_size
831 && self.bytes_consumed != pledged
832 {
833 return Err(invalid_input_error(
834 "pledged content size does not match bytes consumed",
835 ));
836 }
837
838 // An empty frame, or one whose input ended on a block boundary, closes
839 // with an empty last block.
840 crate::encoding::fastpath::with_encode_tier!(
841 self.state.copy_kernel,
842 K => {
843 self.ensure_frame_started::<D, K>(drain)?;
844 self.emit_pending_block::<D, K>(drain, true)
845 }
846 )?;
847
848 #[cfg(feature = "hash")]
849 if self.content_checksum {
850 let checksum = self.hasher.finish() as u32;
851 drain
852 .write_all(&checksum.to_le_bytes())
853 .map_err(|err| self.fail(err))?;
854 }
855
856 drain.flush().map_err(|err| self.fail(err))?;
857 self.abandon_frame();
858 Ok(())
859 }
860
861 /// Drop the frame in progress without closing it (upstream
862 /// `ZSTD_CCtx_reset(ZSTD_reset_session_only)`): its buffered input and its
863 /// pledge go, and the next [`write`](Self::write) starts a new frame. The
864 /// settings, the dictionary and every allocation stay. What the frame
865 /// already wrote to its drain stays there, an unfinished frame.
866 ///
867 /// This is the way on from a [`finish_frame`](Self::finish_frame) refused
868 /// for a pledge the frame did not meet, when the rest of the input is not
869 /// coming; writing it and finishing again completes the frame instead.
870 ///
871 /// # Examples
872 /// ```
873 /// use structured_zstd::encoding::{CompressionContext, CompressionLevel};
874 ///
875 /// let mut context = CompressionContext::new(CompressionLevel::Default);
876 /// let mut unfinished = Vec::new();
877 /// context.set_pledged_content_size(100).unwrap();
878 /// context.write(&mut unfinished, b"only part of it").unwrap();
879 /// assert!(context.finish_frame(&mut unfinished).is_err());
880 /// context.abandon_frame();
881 ///
882 /// let mut frame = Vec::new();
883 /// context.write(&mut frame, b"a frame of its own").unwrap();
884 /// context.finish_frame(&mut frame).unwrap();
885 /// ```
886 pub fn abandon_frame(&mut self) {
887 // What belongs to the frame goes with it; the settings, the dictionary
888 // and every allocation stay for the next one.
889 self.frame_started = false;
890 self.bytes_consumed = 0;
891 self.pledged_content_size = None;
892 // A compressed frame's buffered input is uncommitted in the match
893 // finder, which the next frame's reset drops.
894 self.encoded_scratch.clear();
895 }
896
897 /// Start the next frame now, writing its header to `drain`, rather than at
898 /// its first [`write`](Self::write). Settings are fixed from here on, as
899 /// after a write; a frame already open is left as it is.
900 ///
901 /// ```
902 /// use structured_zstd::encoding::{CompressionContext, CompressionLevel};
903 ///
904 /// let mut context = CompressionContext::new(CompressionLevel::Default);
905 /// let mut frame = Vec::new();
906 /// context.start_frame(&mut frame).unwrap();
907 /// assert!(!frame.is_empty());
908 /// assert!(context.frame_block_size().is_some());
909 /// ```
910 pub fn start_frame<D: Write + ?Sized>(&mut self, drain: &mut D) -> Result<(), Error> {
911 self.ensure_open()?;
912 crate::encoding::fastpath::with_encode_tier!(
913 self.state.copy_kernel,
914 K => self.ensure_frame_started::<D, K>(drain)
915 )
916 }
917
918 /// The largest block the open frame cuts: the block size limits held to
919 /// the window (upstream `ZSTD_getBlockSize`). `None` between frames.
920 pub fn frame_block_size(&self) -> Option<usize> {
921 self.frame_started.then(|| self.block_capacity())
922 }
923
924 /// Compress `block` as one block of the open frame and append its content,
925 /// without the block header, to `payload` (upstream `ZSTD_compressBlock`).
926 /// Returns whether it compressed: `false` appends nothing, and the caller
927 /// stores the block raw, as a decoder then takes it into its history.
928 /// Either way the block joins the history later blocks match against.
929 ///
930 /// The frame must be open ([`start_frame`](Self::start_frame)), hold no
931 /// input still waiting for its block, and `block` must fit
932 /// [`frame_block_size`](Self::frame_block_size).
933 ///
934 /// ```
935 /// use structured_zstd::encoding::{CompressionContext, CompressionLevel};
936 ///
937 /// let mut context = CompressionContext::new(CompressionLevel::Default);
938 /// context.start_frame(&mut Vec::new()).unwrap();
939 /// let mut payload = Vec::new();
940 /// let compressed = context.compress_block(&b"blocks".repeat(100), &mut payload).unwrap();
941 /// assert!(compressed && payload.len() < 600);
942 /// ```
943 pub fn compress_block(&mut self, block: &[u8], payload: &mut Vec<u8>) -> Result<bool, Error> {
944 self.ensure_open()?;
945 if !self.frame_started {
946 return Err(invalid_input_error(
947 "the frame must be started before a block of it is compressed",
948 ));
949 }
950 if self.pending_len() != 0 {
951 return Err(invalid_input_error(
952 "a block cannot follow input still waiting for its block",
953 ));
954 }
955 if block.len() > self.block_capacity() {
956 return Err(invalid_input_error(
957 "the block is larger than the frame's block size",
958 ));
959 }
960 if block.is_empty() {
961 return Ok(false);
962 }
963 self.buffer_input(block);
964 let start = payload.len();
965 // One block in, one block out: the caller writes its single header,
966 // so the post-split, which would emit several physical blocks, stays
967 // off (upstream `ZSTD_compressBlock` never splits either).
968 let post_split = core::mem::replace(&mut self.state.post_split_allowed, false);
969 // The context's one choice of kernel tier for this block.
970 let encoded = crate::encoding::fastpath::with_encode_tier!(
971 self.state.copy_kernel,
972 K => self.encode_block::<_, K>(payload, block.len(), false)
973 );
974 self.state.post_split_allowed = post_split;
975 encoded.map_err(|err| self.fail(err))?;
976 self.bytes_consumed += block.len() as u64;
977 // RFC 8878 3.1.1.2.2: Block_Type 2 is Compressed. A Raw or RLE block
978 // changes no entropy table and no repeat offset, exactly as a block
979 // the caller stores raw does on the decoding side.
980 let compressed = (payload[start] >> 1) & 3 == 2;
981 if compressed {
982 payload.drain(start..start + BLOCK_HEADER_LEN);
983 } else {
984 payload.truncate(start);
985 }
986 Ok(compressed)
987 }
988
989 /// Input written to the open frame that no block has compressed yet: it
990 /// waits for its block to fill, or for a flush. Zero between frames.
991 /// Backs the progress the C `ZSTD_getFrameProgression` reports.
992 ///
993 /// ```
994 /// use structured_zstd::encoding::{CompressionContext, CompressionLevel};
995 ///
996 /// let mut context = CompressionContext::new(CompressionLevel::Default);
997 /// let mut frame = Vec::new();
998 /// context.write(&mut frame, b"less than a block").unwrap();
999 /// assert_eq!(context.buffered_input_len(), 17);
1000 /// context.flush(&mut frame).unwrap();
1001 /// assert_eq!(context.buffered_input_len(), 0);
1002 /// ```
1003 pub fn buffered_input_len(&self) -> usize {
1004 self.pending_len()
1005 }
1006
1007 /// How much more input completes the block being filled, or a whole block
1008 /// when none is partly filled: the read size that lets the next write
1009 /// compress a block straight away (upstream zstd `ZSTD_nextInputSizeHint`).
1010 ///
1011 /// ```
1012 /// use structured_zstd::encoding::{CompressionContext, CompressionLevel};
1013 ///
1014 /// let mut context = CompressionContext::new(CompressionLevel::Default);
1015 /// let mut frame = Vec::new();
1016 /// context.write(&mut frame, &[7u8; 100]).unwrap();
1017 /// assert_eq!(context.next_input_size_hint(), 128 * 1024 - 100);
1018 /// ```
1019 pub fn next_input_size_hint(&self) -> usize {
1020 let block = self.block_capacity();
1021 block - self.pending_len() % block
1022 }
1023
1024 fn ensure_open(&self) -> Result<(), Error> {
1025 if self.errored {
1026 return Err(self.sticky_error());
1027 }
1028 Ok(())
1029 }
1030
1031 /// Refuse a setting once the frame it would change is under way.
1032 fn ensure_settable(&self, too_late: &str) -> Result<(), Error> {
1033 self.ensure_open()?;
1034 if self.frame_started {
1035 return Err(invalid_input_error(too_late));
1036 }
1037 Ok(())
1038 }
1039
1040 // Cold path (only reached after poisoning). The format!() calls still allocate
1041 // in no_std even though error_with_kind_message/other_error_owned drop the
1042 // message; this is acceptable on an error recovery path to keep match arms simple.
1043 fn sticky_error(&self) -> Error {
1044 match (self.last_error_kind, self.last_error_message.as_deref()) {
1045 (Some(kind), Some(message)) => error_with_kind_message(
1046 kind,
1047 format!(
1048 "streaming encoder is in an errored state due to previous {kind:?} failure: {message}"
1049 ),
1050 ),
1051 (Some(kind), None) => error_from_kind(kind),
1052 (None, Some(message)) => other_error_owned(format!(
1053 "streaming encoder is in an errored state: {message}"
1054 )),
1055 (None, None) => other_error("streaming encoder is in an errored state"),
1056 }
1057 }
1058
1059 fn ensure_frame_started<D: Write + ?Sized, K: EncodeTier>(
1060 &mut self,
1061 drain: &mut D,
1062 ) -> Result<(), Error> {
1063 if self.frame_started {
1064 return Ok(());
1065 }
1066
1067 // Frames are independent, so the raw-skip's memory of emitted content
1068 // starts empty; the allocation is kept across frames.
1069 self.state.seen_content.reset_for_frame();
1070 // Same reason as the frame compressor's start: what the last frame
1071 // ended on is about to be replaced, and it is exactly the buffer this
1072 // frame wants to build into.
1073 self.state.fse_tables.park_previous_before_frame();
1074 self.ensure_level_supported()?;
1075 // A dictionary is only active when it can actually be primed: the level
1076 // compresses (not `Uncompressed`) AND the matcher supports priming AND a
1077 // dictionary is attached. Mirrors `FrameCompressor`'s `use_dictionary_state`
1078 // so a streaming frame never advertises a `Dictionary_ID` or seeds dict
1079 // entropy/offsets unless the dictionary is genuinely in play (otherwise
1080 // it would emit frames that needlessly require a dictionary at decode
1081 // time).
1082 let use_dictionary_state =
1083 !matches!(self.compression_level, CompressionLevel::Uncompressed)
1084 && self.state.matcher.supports_dictionary_priming()
1085 && self.dictionary.is_some();
1086 // The dictionary sizes select the CDict cParams tier (consumed inside
1087 // `reset`), so hand them over BEFORE reset.
1088 if use_dictionary_state && let Some(dict) = self.dictionary.as_ref() {
1089 self.state.matcher.set_dictionary_size_hint(dict.sizes());
1090 }
1091 // The matcher resolves the frame from the LAST hint it was handed, so
1092 // re-forward the authoritative size (`pledge.or(advisory)`, the same
1093 // value the gates below read) right before the reset: without this a
1094 // pledge followed by a different advisory hint (or vice versa) left
1095 // the matcher and the frame gates on different size tiers.
1096 if let Some(size) = self.pledged_content_size.or(self.source_size_hint) {
1097 self.state.matcher.set_source_size_hint(size);
1098 }
1099 let block_target = crate::encoding::frame_compressor::block_ceiling(
1100 self.target_block_size,
1101 self.max_block_size,
1102 );
1103 // Raw frames keep no history; everything else is read block by block,
1104 // and a pledge makes its length exact. One past what the address space
1105 // holds is laid out as a stream, which sizes to the window anyway.
1106 let ingest = if matches!(self.compression_level, CompressionLevel::Uncompressed) {
1107 crate::encoding::workspace::IngestPlan::Raw
1108 } else if let Some(pledged) = self
1109 .pledged_content_size
1110 .and_then(|pledged| usize::try_from(pledged).ok())
1111 {
1112 crate::encoding::workspace::IngestPlan::PledgedStream(pledged)
1113 } else {
1114 crate::encoding::workspace::IngestPlan::Stream
1115 };
1116 self.state
1117 .reset_for_frame(self.compression_level, block_target, ingest);
1118 let block_capacity = self.block_capacity();
1119 self.state.finish_layout(block_capacity);
1120 // Sync `state.strategy_tag` / `state.pre_split` to the strategy the
1121 // matcher's reset resolved (size- and dictionary-adaptive; a public
1122 // strategy override wins) so the literal-compression gates, the block
1123 // pre-splitter and the dictionary load below agree with the parse.
1124 // Mirrors `FrameCompressor::compress` and keeps both entry points
1125 // byte-equivalent.
1126 let hint = self.pledged_content_size.or(self.source_size_hint);
1127 let params = crate::encoding::frame_compressor::resolve_frame_params(
1128 self.compression_level,
1129 hint,
1130 self.dictionary.as_ref().filter(|_| use_dictionary_state),
1131 );
1132 crate::encoding::frame_compressor::sync_effective_strategy(
1133 &mut self.state,
1134 self.compression_level,
1135 ¶ms,
1136 self.tuning.strategy,
1137 );
1138 self.state.huf_optimal_search =
1139 crate::encoding::frame_compressor::huf_search_enabled(self.state.strategy_tag, hint);
1140 self.state.literal_compression_disabled =
1141 crate::encoding::frame_compressor::literal_compression_disabled(
1142 self.state.strategy_tag,
1143 self.compression_level,
1144 crate::encoding::frame_compressor::gate_target_length(
1145 self.compression_level,
1146 &self.tuning,
1147 self.dictionary.as_ref().filter(|_| use_dictionary_state),
1148 ),
1149 self.tuning.literal_compression,
1150 );
1151 // Seed the repeat-offset history from the dictionary (upstream zstd
1152 // `ZSTD_compress_insertDictionary`), or the default rep codes
1153 // otherwise, and load the dictionary into the match finder: primed,
1154 // restored from a snapshot, or, on a reused context whose reset kept
1155 // it resident, left in place with its offsets reapplied.
1156 // `dict` borrows `self.dictionary`; `self.state` is a disjoint field.
1157 self.state.offset_hist = [1, 4, 8];
1158 if use_dictionary_state && let Some(dict) = self.dictionary.as_ref() {
1159 self.state.offset_hist = dict.inner.offset_hist;
1160 crate::encoding::frame_compressor::load_frame_dictionary::<M, K>(
1161 &mut self.state,
1162 self.compression_level,
1163 dict,
1164 hint,
1165 );
1166 }
1167 // Seed the first block's entropy from the dictionary's encoder tables
1168 // (upstream zstd `cdict->cBlockState`), or clear to defaults.
1169 if use_dictionary_state && let Some(dict) = self.dictionary.as_ref() {
1170 let cache = &dict.inner.entropy;
1171 self.state.last_huff_table.clone_from(&cache.huff);
1172 self.state
1173 .fse_tables
1174 .ll_previous
1175 .clone_from(&cache.ll_previous);
1176 self.state
1177 .fse_tables
1178 .ml_previous
1179 .clone_from(&cache.ml_previous);
1180 self.state
1181 .fse_tables
1182 .of_previous
1183 .clone_from(&cache.of_previous);
1184 let ll_entropy = match cache.ll_previous.as_ref() {
1185 Some(PreviousFseTable::Custom(table)) => Some(table.as_ref()),
1186 _ => None,
1187 };
1188 let ml_entropy = match cache.ml_previous.as_ref() {
1189 Some(PreviousFseTable::Custom(table)) => Some(table.as_ref()),
1190 _ => None,
1191 };
1192 let of_entropy = match cache.of_previous.as_ref() {
1193 Some(PreviousFseTable::Custom(table)) => Some(table.as_ref()),
1194 _ => None,
1195 };
1196 self.state.matcher.seed_dictionary_entropy(
1197 self.state.last_huff_table.as_ref(),
1198 ll_entropy,
1199 ml_entropy,
1200 of_entropy,
1201 );
1202 } else {
1203 self.state.last_huff_table = None;
1204 self.state.fse_tables.ll_previous = None;
1205 self.state.fse_tables.ml_previous = None;
1206 self.state.fse_tables.of_previous = None;
1207 }
1208 self.savings = 0;
1209 #[cfg(feature = "hash")]
1210 {
1211 self.hasher = XxHash64::with_seed(0);
1212 }
1213
1214 let window_size = self.state.matcher.window_size();
1215 if window_size == 0 {
1216 return Err(invalid_input_error(
1217 "matcher reported window_size == 0, which is invalid",
1218 ));
1219 }
1220
1221 let frame_content_size = self.pledged_content_size.filter(|_| self.content_size_flag);
1222 let single_segment =
1223 crate::encoding::frame_header::single_segment(frame_content_size, window_size);
1224
1225 let header = FrameHeader {
1226 frame_content_size,
1227 single_segment,
1228 content_checksum: cfg!(feature = "hash") && self.content_checksum,
1229 dictionary_id: if use_dictionary_state && self.dictionary_id_flag {
1230 // Id 0 is a raw-content dictionary: RFC 8878 spells "no
1231 // dictionary ID" as an absent field, not as a stored zero.
1232 self.dictionary
1233 .as_ref()
1234 .map(|dict| dict.inner.id)
1235 .filter(|id| *id != 0)
1236 .map(u64::from)
1237 } else {
1238 None
1239 },
1240 window_size: if single_segment {
1241 None
1242 } else {
1243 Some(window_size)
1244 },
1245 magicless: self.magicless,
1246 };
1247 let mut encoded_header = Vec::new();
1248 header.serialize(&mut encoded_header);
1249 drain
1250 .write_all(&encoded_header)
1251 .map_err(|err| self.fail(err))?;
1252
1253 if matches!(self.compression_level, CompressionLevel::Uncompressed) {
1254 // The raw block's header goes in front of its input once the
1255 // block's length is known.
1256 self.encoded_scratch.clear();
1257 self.encoded_scratch
1258 .reserve(BLOCK_HEADER_LEN + self.block_capacity());
1259 self.encoded_scratch
1260 .extend_from_slice(&[0; BLOCK_HEADER_LEN]);
1261 }
1262 self.frame_started = true;
1263 Ok(())
1264 }
1265
1266 fn block_capacity(&self) -> usize {
1267 let matcher_window = self.state.matcher.window_size() as usize;
1268 let ceiling = crate::encoding::frame_compressor::block_ceiling(
1269 self.target_block_size,
1270 self.max_block_size,
1271 );
1272 core::cmp::max(1, core::cmp::min(matcher_window, ceiling))
1273 }
1274
1275 /// Input taken for the frame that no block has carried out yet.
1276 fn pending_len(&self) -> usize {
1277 if !self.frame_started {
1278 return 0;
1279 }
1280 match self.compression_level {
1281 CompressionLevel::Uncompressed => self.encoded_scratch.len() - BLOCK_HEADER_LEN,
1282 _ => self.state.matcher.uncommitted_input().len(),
1283 }
1284 }
1285
1286 /// Buffers `bytes` for the block being filled: behind the raw block's
1287 /// header, or uncommitted in the match finder's history, where the block
1288 /// is later matched without another copy.
1289 fn buffer_input(&mut self, bytes: &[u8]) {
1290 match self.compression_level {
1291 CompressionLevel::Uncompressed => self.encoded_scratch.extend_from_slice(bytes),
1292 _ => {
1293 self.state
1294 .matcher
1295 .fill_in_place(bytes.len(), &mut |history| {
1296 history.extend_from_slice(bytes);
1297 (bytes.len(), false)
1298 });
1299 }
1300 }
1301 }
1302
1303 /// Where the full pending block is cut (upstream `ZSTD_compress_frameChunk`
1304 /// sizing every block with `ZSTD_optimalBlockSize`): the pre-splitter's
1305 /// boundary once the frame has saved enough, else the whole block. Writes
1306 /// may continue, so the splitter is told a full block is still to come.
1307 fn pre_split_len<K: EncodeTier>(&mut self, block_capacity: usize) -> usize {
1308 if matches!(self.compression_level, CompressionLevel::Uncompressed) {
1309 return self.pending_len();
1310 }
1311 let pending = self.state.matcher.uncommitted_input();
1312 let window = pending.len();
1313 let cut = crate::encoding::frame_compressor::optimal_block_size_with(
1314 self.state.pre_split.map(usize::from),
1315 pending,
1316 block_capacity,
1317 block_capacity,
1318 self.savings,
1319 )
1320 .min(window);
1321 // The splitter sees byte statistics only; a cut it proposes inside a
1322 // match the matcher already knows of is moved out.
1323 if cut < window {
1324 match self.state.matcher.builtin_driver() {
1325 Some(driver) => driver.settle_block_cut_with::<K>(window, cut),
1326 None => self.state.matcher.settle_block_cut(window, cut),
1327 }
1328 } else {
1329 cut
1330 }
1331 }
1332
1333 /// Emit the pending block once it is full, cut where the pre-splitter
1334 /// says; the suffix stays pending and heads the next block. A failure
1335 /// leaves the context failed.
1336 fn emit_full_pending_block<D: Write + ?Sized, K: EncodeTier>(
1337 &mut self,
1338 drain: &mut D,
1339 block_capacity: usize,
1340 ) -> Result<(), Error> {
1341 if self.pending_len() != block_capacity {
1342 return Ok(());
1343 }
1344 let block_len = self.pre_split_len::<K>(block_capacity);
1345 self.encode_block::<D, K>(drain, block_len, false)
1346 .map_err(|err| self.fail(err))
1347 }
1348
1349 fn emit_pending_block<D: Write + ?Sized, K: EncodeTier>(
1350 &mut self,
1351 drain: &mut D,
1352 last_block: bool,
1353 ) -> Result<(), Error> {
1354 // `write` emits a block the moment it fills, so what is left here is
1355 // always short of one and goes out whole; a write that failed on the
1356 // full block left the context failed, and `ensure_open` stops here.
1357 let block_len = self.pending_len();
1358 debug_assert!(block_len < self.block_capacity());
1359 self.encode_block::<D, K>(drain, block_len, last_block)
1360 .map_err(|err| self.fail(err))
1361 }
1362
1363 // Exhaustive match kept intentionally: adding a new CompressionLevel
1364 // variant will produce a compile error here, forcing the developer to
1365 // decide whether the streaming encoder supports it before shipping.
1366 fn ensure_level_supported(&self) -> Result<(), Error> {
1367 match self.compression_level {
1368 CompressionLevel::Uncompressed
1369 | CompressionLevel::Fastest
1370 | CompressionLevel::Default
1371 | CompressionLevel::Better
1372 | CompressionLevel::Best
1373 | CompressionLevel::Level(_) => Ok(()),
1374 }
1375 }
1376
1377 /// Writes the first `block_len` pending bytes to `drain` as one block,
1378 /// `last_block` if it closes the frame.
1379 fn encode_block<D: Write + ?Sized, K: EncodeTier>(
1380 &mut self,
1381 drain: &mut D,
1382 block_len: usize,
1383 last_block: bool,
1384 ) -> Result<(), Error> {
1385 if matches!(self.compression_level, CompressionLevel::Uncompressed) {
1386 // The block was assembled behind its header as it was written.
1387 debug_assert_eq!(block_len, self.pending_len());
1388 let header = BlockHeader {
1389 last_block,
1390 block_type: crate::blocks::block::BlockType::Raw,
1391 // At most one block, which a block header's 21-bit size holds.
1392 block_size: block_len as u32,
1393 };
1394 self.encoded_scratch[..BLOCK_HEADER_LEN].copy_from_slice(&header.to_le_bytes());
1395 drain.write_all(&self.encoded_scratch)?;
1396 #[cfg(feature = "hash")]
1397 if self.content_checksum {
1398 self.hasher.write(&self.encoded_scratch[BLOCK_HEADER_LEN..]);
1399 }
1400 self.encoded_scratch.truncate(BLOCK_HEADER_LEN);
1401 return Ok(());
1402 }
1403 self.encoded_scratch.clear();
1404 if block_len == 0 {
1405 // Only the last block of a frame is ever empty.
1406 let header = BlockHeader {
1407 last_block,
1408 block_type: crate::blocks::block::BlockType::Raw,
1409 block_size: 0,
1410 };
1411 header.serialize(&mut self.encoded_scratch);
1412 return drain.write_all(&self.encoded_scratch);
1413 }
1414 let dict_active =
1415 self.dictionary.is_some() && self.state.matcher.supports_dictionary_priming();
1416 compress_block_encoded::<M, K>(
1417 &mut self.state,
1418 self.compression_level,
1419 last_block,
1420 block_len,
1421 &mut self.encoded_scratch,
1422 dict_active,
1423 // No FrameEmitInfo on the streaming encoder path: it does not
1424 // surface per-block layout, so no sidecar.
1425 #[cfg(feature = "lsm")]
1426 None,
1427 #[cfg(all(feature = "lsm", feature = "hash"))]
1428 None,
1429 );
1430 // The block is committed to the match finder whether or not the drain
1431 // takes it; a drain that fails leaves the context failed.
1432 drain.write_all(&self.encoded_scratch)?;
1433 // `savings` counts the block header too, as upstream's
1434 // `ZSTD_compress_frameChunk` does (`cSize` includes it).
1435 self.savings += block_len as i64 - self.encoded_scratch.len() as i64;
1436 #[cfg(feature = "hash")]
1437 if self.content_checksum {
1438 self.hasher.write(self.state.matcher.get_last_space());
1439 }
1440 Ok(())
1441 }
1442
1443 fn fail(&mut self, err: Error) -> Error {
1444 self.errored = true;
1445 if self.last_error_kind.is_none() {
1446 self.last_error_kind = Some(err.kind());
1447 }
1448 if self.last_error_message.is_none() {
1449 self.last_error_message = Some(err.to_string());
1450 }
1451 err
1452 }
1453}
1454
1455/// Bytes of a block header (RFC 8878 3.1.1.2).
1456const BLOCK_HEADER_LEN: usize = 3;
1457
1458fn error_from_kind(kind: ErrorKind) -> Error {
1459 Error::from(kind)
1460}
1461
1462fn error_with_kind_message(kind: ErrorKind, message: String) -> Error {
1463 #[cfg(feature = "std")]
1464 {
1465 Error::new(kind, message)
1466 }
1467 #[cfg(not(feature = "std"))]
1468 {
1469 Error::new(kind, alloc::boxed::Box::new(message))
1470 }
1471}
1472
1473fn invalid_input_error(message: &str) -> Error {
1474 #[cfg(feature = "std")]
1475 {
1476 Error::new(ErrorKind::InvalidInput, message)
1477 }
1478 #[cfg(not(feature = "std"))]
1479 {
1480 Error::new(
1481 ErrorKind::Other,
1482 alloc::boxed::Box::new(alloc::string::String::from(message)),
1483 )
1484 }
1485}
1486
1487fn other_error_owned(message: String) -> Error {
1488 #[cfg(feature = "std")]
1489 {
1490 Error::other(message)
1491 }
1492 #[cfg(not(feature = "std"))]
1493 {
1494 Error::new(ErrorKind::Other, alloc::boxed::Box::new(message))
1495 }
1496}
1497
1498fn other_error(message: &str) -> Error {
1499 #[cfg(feature = "std")]
1500 {
1501 Error::other(message)
1502 }
1503 #[cfg(not(feature = "std"))]
1504 {
1505 Error::new(
1506 ErrorKind::Other,
1507 alloc::boxed::Box::new(alloc::string::String::from(message)),
1508 )
1509 }
1510}
1511
1512#[cfg(test)]
1513mod tests;