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