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