Skip to main content

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][..]);
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_kernel: crate::encoding::fastpath::select_kernel(),
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                post_split_allowed: true,
417            },
418            encoded_scratch: Vec::new(),
419            errored: false,
420            last_error_kind: None,
421            last_error_message: None,
422            frame_started: false,
423            target_block_size: None,
424            pledged_content_size: None,
425            source_size_hint: None,
426            content_size_flag: true,
427            bytes_consumed: 0,
428            savings: 0,
429            tuning: FrameTuning::default(),
430            magicless: false,
431            content_checksum: false,
432            dictionary: None,
433            dictionary_id_flag: true,
434            #[cfg(feature = "hash")]
435            hasher: XxHash64::with_seed(0),
436        }
437    }
438
439    /// Compress the next frames at `level`, with the level's own tuning: any
440    /// parameter override installed by
441    /// [`set_parameters`](CompressionContext::set_parameters) is dropped, as
442    /// [`FrameCompressor::set_compression_level`](crate::encoding::FrameCompressor::set_compression_level)
443    /// drops it. Must be called before the frame's first [`write`](Self::write).
444    ///
445    /// # Examples
446    /// ```
447    /// use structured_zstd::encoding::{CompressionContext, CompressionLevel};
448    ///
449    /// let mut context = CompressionContext::new(CompressionLevel::Fastest);
450    /// context.set_compression_level(CompressionLevel::Better).unwrap();
451    /// let mut frame = Vec::new();
452    /// context.write(&mut frame, b"compressed at the new level").unwrap();
453    /// context.finish_frame(&mut frame).unwrap();
454    /// ```
455    pub fn set_compression_level(&mut self, level: CompressionLevel) -> Result<(), Error> {
456        self.ensure_settable("the compression level must be set before the first write")?;
457        self.compression_level = level;
458        self.tuning = FrameTuning::default();
459        self.state.matcher.clear_param_overrides();
460        Ok(())
461    }
462
463    /// Set an upper bound on each physical block's payload (semantics of
464    /// upstream `ZSTD_c_targetCBlockSize`): every block carries at most
465    /// `target` payload bytes, +3-byte block header on the wire — the
466    /// upstream knob is likewise a convergence target for block sizing,
467    /// not a cap on header-inclusive wire bytes. Clamped to
468    /// `[MIN_TARGET_BLOCK_SIZE, MAX_BLOCK_SIZE]`; mirrors
469    /// `FrameCompressor::set_target_block_size`. Must be set before the
470    /// frame's first write.
471    pub fn set_target_block_size(&mut self, target: Option<u32>) -> Result<(), Error> {
472        self.ensure_settable("the block-size target must be set before the first write")?;
473        self.target_block_size = target.map(|t| {
474            t.clamp(
475                crate::common::MIN_TARGET_BLOCK_SIZE,
476                crate::common::MAX_BLOCK_SIZE,
477            )
478        });
479        Ok(())
480    }
481
482    /// Enable or disable the trailing XXH64 content checksum
483    /// (upstream `ZSTD_c_checksumFlag`). Default `false`, matching the
484    /// upstream library default (`ZSTD_c_checksumFlag = 0`). Must be called
485    /// before the frame's first [`write`](Self::write); once the frame header
486    /// is emitted the flag is fixed, so a late change returns an error rather
487    /// than producing a header/trailer mismatch. Without the `hash` feature
488    /// no checksum is emitted regardless.
489    pub fn set_content_checksum(&mut self, emit: bool) -> Result<(), Error> {
490        self.ensure_settable("content checksum must be set before the first write")?;
491        self.content_checksum = emit;
492        Ok(())
493    }
494
495    /// Enable or disable magicless frame format (`ZSTD_f_zstd1_magicless`).
496    ///
497    /// When set to `true`, the frame header omits the 4-byte magic number
498    /// prefix. Must be called BEFORE the frame's first [`write`](Self::write)
499    /// call; calling it after the frame header has already been emitted
500    /// returns an error so the caller can't be misled into thinking they
501    /// produced a magicless stream.
502    pub fn set_magicless(&mut self, magicless: bool) -> Result<(), Error> {
503        self.ensure_settable("magicless format must be set before the first write")?;
504        self.magicless = magicless;
505        Ok(())
506    }
507
508    /// Pledge the total uncompressed content size of the next frame.
509    ///
510    /// When set, the frame header will include a `Frame_Content_Size` field.
511    /// This enables decoders to pre-allocate output buffers.
512    /// The pledged size is also forwarded as a source-size hint to the
513    /// matcher so small inputs can use smaller matching tables.
514    ///
515    /// Must be called **before** the frame's first [`write`](Self::write);
516    /// calling it after the frame header has already been emitted returns an
517    /// error. The pledge ends with the frame.
518    pub fn set_pledged_content_size(&mut self, size: u64) -> Result<(), Error> {
519        self.ensure_settable("pledged content size must be set before the first write")?;
520        // The matcher is handed it as its size hint when the frame starts
521        // (`ensure_frame_started`), not here: a pledge that goes unused must
522        // not stay behind in the matcher for a later frame.
523        self.pledged_content_size = Some(size);
524        Ok(())
525    }
526
527    /// Control whether the pledged size is written into the header's
528    /// `Frame_Content_Size` field (upstream `ZSTD_c_contentSizeFlag`,
529    /// default on). With the flag off the header omits the field, but a
530    /// pledge set via [`set_pledged_content_size`](Self::set_pledged_content_size)
531    /// is still enforced against the bytes actually written. Must be
532    /// called before the frame's first [`write`](Self::write).
533    pub fn set_content_size_flag(&mut self, emit: bool) -> Result<(), Error> {
534        self.ensure_settable("content size flag must be set before the first write")?;
535        self.content_size_flag = emit;
536        Ok(())
537    }
538
539    /// Provide a hint about the total uncompressed size of each frame.
540    ///
541    /// Unlike [`set_pledged_content_size`](Self::set_pledged_content_size),
542    /// this does **not** enforce that exactly `size` bytes are written; it
543    /// may reduce matcher tables, advertised frame window, and block sizing
544    /// for small inputs. A parameter, like upstream `ZSTD_c_srcSizeHint`: it
545    /// applies to every frame until replaced. Must be called before the
546    /// frame's first [`write`](Self::write).
547    pub fn set_source_size_hint(&mut self, size: u64) -> Result<(), Error> {
548        self.ensure_settable("source size hint must be set before the first write")?;
549        // Read at each frame start as `pledged_content_size.or(source_size_hint)`,
550        // by the matcher's sizing and the Fast HUF fast-path gate alike
551        // (`ensure_frame_started`), so a small advisory size also lifts Fast
552        // streams off the expensive optimal-HUF search.
553        self.source_size_hint = Some(size);
554        Ok(())
555    }
556
557    /// Attach a dictionary blob to each frame (upstream zstd
558    /// `ZSTD_CCtx_loadDictionary` on a streaming context, which loads in
559    /// `ZSTD_dct_auto` mode): a blob prefixed with
560    /// [`DICTIONARY_MAGIC`](crate::decoding::DICTIONARY_MAGIC) is a serialized
561    /// dictionary, anything else is raw content. The dictionary primes the
562    /// match-finder and seeds the first block's entropy tables + repeat
563    /// offsets; a serialized one's ID is written into the frame header, while
564    /// raw content has none to write, so the decoder must be given the same
565    /// bytes explicitly. Must be called before the frame's first
566    /// [`write`](Self::write); repeat offsets must be non-zero.
567    pub fn set_dictionary_from_bytes(&mut self, raw_dictionary: &[u8]) -> Result<(), Error> {
568        if raw_dictionary.is_empty() {
569            // An empty buffer is how the same upstream entry point is told
570            // there is no dictionary: it clears and succeeds. Still refused
571            // once the frame is open, like any other attach.
572            self.ensure_settable("dictionary must be attached before the first write")?;
573            // What the match finder kept of it across frames (its primed
574            // snapshot, the copy left resident) goes with it, entropy included.
575            self.state.matcher.invalidate_primed_dictionary();
576            self.dictionary = None;
577            return Ok(());
578        }
579        let dict = EncoderDictionary::from_serialized_or_raw_content(raw_dictionary)
580            .map_err(|err| invalid_input_error(&alloc::format!("invalid dictionary: {err:?}")))?;
581        self.set_encoder_dictionary(dict)
582    }
583
584    /// Whether the frame header records the attached dictionary's ID
585    /// (upstream `ZSTD_c_dictIDFlag` semantics; default `true`).
586    /// Mirrors [`FrameCompressor::set_dictionary_id_flag`](crate::encoding::FrameCompressor::set_dictionary_id_flag).
587    /// Decoders can still decode such frames by supplying the dictionary
588    /// explicitly.
589    pub fn set_dictionary_id_flag(&mut self, emit: bool) -> Result<(), Error> {
590        self.ensure_settable("dictionary ID flag must be set before the first write")?;
591        self.dictionary_id_flag = emit;
592        Ok(())
593    }
594
595    /// Attach an already-parsed [`EncoderDictionary`] to each frame. See
596    /// [`set_dictionary_from_bytes`](Self::set_dictionary_from_bytes); must be
597    /// called before the frame's first write. The entropy tables it seeds were
598    /// built when it was prepared, so attaching builds nothing.
599    pub fn set_encoder_dictionary(&mut self, dict: EncoderDictionary) -> Result<(), Error> {
600        self.ensure_settable("dictionary must be attached before the first write")?;
601        // A zero id marks a raw-content dictionary, which carries no header to
602        // hold one; the frame then records no dictionary ID and the decoder
603        // must be given the same bytes explicitly.
604        if dict.inner.offset_hist.contains(&0) {
605            return Err(invalid_input_error(
606                "dictionary carries a zero repeat offset",
607            ));
608        }
609        // The match finder's primed snapshot and resident copy belong to the
610        // dictionary being replaced; the next frame primes the new one.
611        self.state.matcher.invalidate_primed_dictionary();
612        self.dictionary = Some(dict);
613        Ok(())
614    }
615
616    /// The dictionary each frame is compressed with, if one is attached.
617    ///
618    /// # Examples
619    /// ```
620    /// use structured_zstd::encoding::{CompressionContext, CompressionLevel, EncoderDictionary};
621    ///
622    /// let dictionary = EncoderDictionary::from_serialized_or_raw_content(b"some shared history").unwrap();
623    /// let mut context = CompressionContext::new(CompressionLevel::Default);
624    /// assert!(context.dictionary().is_none());
625    /// context.set_encoder_dictionary(dictionary).unwrap();
626    /// assert!(context.dictionary().is_some());
627    /// ```
628    pub fn dictionary(&self) -> Option<&EncoderDictionary> {
629        self.dictionary.as_ref()
630    }
631
632    /// Total heap bytes this context's allocations hold, excluding the inline
633    /// struct: match-finder tables / history, retained Huffman tables, the
634    /// `encoded_scratch` block buffer, the retained dictionary content, and
635    /// its entropy tables. Mirrors
636    /// `FrameCompressor::heap_size` so a context can report its true
637    /// footprint through `ZSTD_sizeof_CCtx`.
638    pub fn heap_size(&self) -> usize {
639        let mut total = self.state.matcher.heap_size();
640        total += self
641            .state
642            .last_huff_table
643            .as_ref()
644            .map_or(0, |table| table.heap_size());
645        total += self
646            .state
647            .huff_table_spare
648            .as_ref()
649            .map_or(0, |table| table.heap_size());
650        // Kept between blocks and frames; see `FrameCompressor::heap_size`.
651        total += self.state.huff_weights.heap_size();
652        total += self.state.retained_scratch_heap_size();
653        total += self.state.seen_content.heap_size();
654        total += self.encoded_scratch.capacity();
655        total += self
656            .dictionary
657            .as_ref()
658            .map_or(0, EncoderDictionary::heap_size);
659        total
660    }
661
662    /// Compress `buf` into the frame in progress, starting one (and writing
663    /// its header to `drain`) if none is. Full blocks are compressed as they
664    /// fill and written to `drain`; the rest stays buffered for the next call,
665    /// [`flush`](Self::flush) or [`finish_frame`](Self::finish_frame).
666    ///
667    /// Returns how much of `buf` was taken, which is all of it unless a pledge
668    /// set with [`set_pledged_content_size`](Self::set_pledged_content_size)
669    /// allows less.
670    pub fn write<D: Write + ?Sized>(&mut self, drain: &mut D, buf: &[u8]) -> Result<usize, Error> {
671        self.ensure_open()?;
672        if buf.is_empty() {
673            return Ok(0);
674        }
675
676        // Check pledge before emitting the frame header so that a misuse
677        // like set_pledged_content_size(0) + write(non_empty) doesn't leave
678        // a partially-written header in the drain.
679        if let Some(pledged) = self.pledged_content_size
680            && self.bytes_consumed >= pledged
681        {
682            return Err(invalid_input_error(
683                "write would exceed pledged content size",
684            ));
685        }
686
687        self.ensure_frame_started(drain)?;
688
689        // Enforce pledged upper bound: truncate the accepted slice to the
690        // remaining allowance so that partial-write semantics are honored
691        // (return Ok(n) with n < buf.len()) instead of failing the full call.
692        // The check above leaves at least one byte of it.
693        let buf = if let Some(pledged) = self.pledged_content_size {
694            debug_assert!(self.bytes_consumed < pledged);
695            let remaining_allowed = pledged - self.bytes_consumed;
696            let accepted = core::cmp::min(
697                buf.len(),
698                usize::try_from(remaining_allowed).unwrap_or(usize::MAX),
699            );
700            &buf[..accepted]
701        } else {
702            buf
703        };
704
705        let block_capacity = self.block_capacity();
706        let mut remaining = buf;
707        let mut consumed = 0usize;
708
709        while !remaining.is_empty() {
710            // A block is emitted the moment it fills, below, and the
711            // capacity is fixed for the frame, so there is always room here.
712            let pending = self.pending_len();
713            debug_assert!(pending < block_capacity);
714            let to_take = core::cmp::min(remaining.len(), block_capacity - pending);
715            self.buffer_input(&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_len() == 0 {
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        // An empty frame, or one whose input ended on a block boundary, closes
760        // with an empty last block.
761        self.emit_pending_block(drain, true)?;
762
763        #[cfg(feature = "hash")]
764        if self.content_checksum {
765            let checksum = self.hasher.finish() as u32;
766            drain
767                .write_all(&checksum.to_le_bytes())
768                .map_err(|err| self.fail(err))?;
769        }
770
771        drain.flush().map_err(|err| self.fail(err))?;
772        self.abandon_frame();
773        Ok(())
774    }
775
776    /// Drop the frame in progress without closing it (upstream
777    /// `ZSTD_CCtx_reset(ZSTD_reset_session_only)`): its buffered input and its
778    /// pledge go, and the next [`write`](Self::write) starts a new frame. The
779    /// settings, the dictionary and every allocation stay. What the frame
780    /// already wrote to its drain stays there, an unfinished frame.
781    ///
782    /// This is the way on from a [`finish_frame`](Self::finish_frame) refused
783    /// for a pledge the frame did not meet, when the rest of the input is not
784    /// coming; writing it and finishing again completes the frame instead.
785    ///
786    /// # Examples
787    /// ```
788    /// use structured_zstd::encoding::{CompressionContext, CompressionLevel};
789    ///
790    /// let mut context = CompressionContext::new(CompressionLevel::Default);
791    /// let mut unfinished = Vec::new();
792    /// context.set_pledged_content_size(100).unwrap();
793    /// context.write(&mut unfinished, b"only part of it").unwrap();
794    /// assert!(context.finish_frame(&mut unfinished).is_err());
795    /// context.abandon_frame();
796    ///
797    /// let mut frame = Vec::new();
798    /// context.write(&mut frame, b"a frame of its own").unwrap();
799    /// context.finish_frame(&mut frame).unwrap();
800    /// ```
801    pub fn abandon_frame(&mut self) {
802        // What belongs to the frame goes with it; the settings, the dictionary
803        // and every allocation stay for the next one.
804        self.frame_started = false;
805        self.bytes_consumed = 0;
806        self.pledged_content_size = None;
807        // A compressed frame's buffered input is uncommitted in the match
808        // finder, which the next frame's reset drops.
809        self.encoded_scratch.clear();
810    }
811
812    fn ensure_open(&self) -> Result<(), Error> {
813        if self.errored {
814            return Err(self.sticky_error());
815        }
816        Ok(())
817    }
818
819    /// Refuse a setting once the frame it would change is under way.
820    fn ensure_settable(&self, too_late: &str) -> Result<(), Error> {
821        self.ensure_open()?;
822        if self.frame_started {
823            return Err(invalid_input_error(too_late));
824        }
825        Ok(())
826    }
827
828    // Cold path (only reached after poisoning). The format!() calls still allocate
829    // in no_std even though error_with_kind_message/other_error_owned drop the
830    // message; this is acceptable on an error recovery path to keep match arms simple.
831    fn sticky_error(&self) -> Error {
832        match (self.last_error_kind, self.last_error_message.as_deref()) {
833            (Some(kind), Some(message)) => error_with_kind_message(
834                kind,
835                format!(
836                    "streaming encoder is in an errored state due to previous {kind:?} failure: {message}"
837                ),
838            ),
839            (Some(kind), None) => error_from_kind(kind),
840            (None, Some(message)) => other_error_owned(format!(
841                "streaming encoder is in an errored state: {message}"
842            )),
843            (None, None) => other_error("streaming encoder is in an errored state"),
844        }
845    }
846
847    fn ensure_frame_started<D: Write + ?Sized>(&mut self, drain: &mut D) -> Result<(), Error> {
848        if self.frame_started {
849            return Ok(());
850        }
851
852        // Frames are independent, so the raw-skip's memory of emitted content
853        // starts empty; the allocation is kept across frames.
854        self.state.seen_content.reset_for_frame();
855        // Same reason as the frame compressor's start: what the last frame
856        // ended on is about to be replaced, and it is exactly the buffer this
857        // frame wants to build into.
858        self.state.fse_tables.park_previous_before_frame();
859        self.ensure_level_supported()?;
860        // A dictionary is only active when it can actually be primed: the level
861        // compresses (not `Uncompressed`) AND the matcher supports priming AND a
862        // dictionary is attached. Mirrors `FrameCompressor`'s `use_dictionary_state`
863        // so a streaming frame never advertises a `Dictionary_ID`, disables
864        // single-segment, or seeds dict entropy/offsets unless the dictionary is
865        // genuinely in play (otherwise it would emit frames that needlessly
866        // require a dictionary at decode time).
867        let use_dictionary_state =
868            !matches!(self.compression_level, CompressionLevel::Uncompressed)
869                && self.state.matcher.supports_dictionary_priming()
870                && self.dictionary.is_some();
871        // The dictionary sizes select the CDict cParams tier (consumed inside
872        // `reset`), so hand them over BEFORE reset.
873        if use_dictionary_state && let Some(dict) = self.dictionary.as_ref() {
874            self.state.matcher.set_dictionary_size_hint(dict.sizes());
875        }
876        // The matcher resolves the frame from the LAST hint it was handed, so
877        // re-forward the authoritative size (`pledge.or(advisory)`, the same
878        // value the gates below read) right before the reset: without this a
879        // pledge followed by a different advisory hint (or vice versa) left
880        // the matcher and the frame gates on different size tiers.
881        if let Some(size) = self.pledged_content_size.or(self.source_size_hint) {
882            self.state.matcher.set_source_size_hint(size);
883        }
884        let block_target = self
885            .target_block_size
886            .map_or(MAX_BLOCK_SIZE as usize, |t| t as usize);
887        // Raw frames keep no history; everything else is read block by block,
888        // and a pledge makes its length exact. One past what the address space
889        // holds is laid out as a stream, which sizes to the window anyway.
890        let ingest = if matches!(self.compression_level, CompressionLevel::Uncompressed) {
891            crate::encoding::workspace::IngestPlan::Raw
892        } else if let Some(pledged) = self
893            .pledged_content_size
894            .and_then(|pledged| usize::try_from(pledged).ok())
895        {
896            crate::encoding::workspace::IngestPlan::PledgedStream(pledged)
897        } else {
898            crate::encoding::workspace::IngestPlan::Stream
899        };
900        self.state
901            .reset_for_frame(self.compression_level, block_target, ingest);
902        let block_capacity = self.block_capacity();
903        self.state.finish_layout(block_capacity);
904        // Sync `state.strategy_tag` / `state.pre_split` to the strategy the
905        // matcher's reset resolved (size- and dictionary-adaptive; a public
906        // strategy override wins) so the literal-compression gates, the block
907        // pre-splitter and the dictionary load below agree with the parse.
908        // Mirrors `FrameCompressor::compress` and keeps both entry points
909        // byte-equivalent.
910        let hint = self.pledged_content_size.or(self.source_size_hint);
911        let params = crate::encoding::frame_compressor::resolve_frame_params(
912            self.compression_level,
913            hint,
914            self.dictionary.as_ref().filter(|_| use_dictionary_state),
915        );
916        crate::encoding::frame_compressor::sync_effective_strategy(
917            &mut self.state,
918            self.compression_level,
919            &params,
920            self.tuning.strategy,
921        );
922        self.state.huf_optimal_search =
923            crate::encoding::frame_compressor::huf_search_enabled(self.state.strategy_tag, hint);
924        self.state.literal_compression_disabled =
925            crate::encoding::frame_compressor::literal_compression_disabled(
926                self.state.strategy_tag,
927                self.compression_level,
928                crate::encoding::frame_compressor::gate_target_length(
929                    self.compression_level,
930                    &self.tuning,
931                    self.dictionary.as_ref().filter(|_| use_dictionary_state),
932                ),
933                self.tuning.literal_compression,
934            );
935        // Seed the repeat-offset history from the dictionary (upstream zstd
936        // `ZSTD_compress_insertDictionary`), or the default rep codes
937        // otherwise, and load the dictionary into the match finder: primed,
938        // restored from a snapshot, or, on a reused context whose reset kept
939        // it resident, left in place with its offsets reapplied.
940        // `dict` borrows `self.dictionary`; `self.state` is a disjoint field.
941        self.state.offset_hist = [1, 4, 8];
942        if use_dictionary_state && let Some(dict) = self.dictionary.as_ref() {
943            self.state.offset_hist = dict.inner.offset_hist;
944            crate::encoding::frame_compressor::load_frame_dictionary(
945                &mut self.state,
946                self.compression_level,
947                dict,
948                hint,
949            );
950        }
951        // Seed the first block's entropy from the dictionary's encoder tables
952        // (upstream zstd `cdict->cBlockState`), or clear to defaults.
953        if use_dictionary_state && let Some(dict) = self.dictionary.as_ref() {
954            let cache = &dict.inner.entropy;
955            self.state.last_huff_table.clone_from(&cache.huff);
956            self.state
957                .fse_tables
958                .ll_previous
959                .clone_from(&cache.ll_previous);
960            self.state
961                .fse_tables
962                .ml_previous
963                .clone_from(&cache.ml_previous);
964            self.state
965                .fse_tables
966                .of_previous
967                .clone_from(&cache.of_previous);
968            let ll_entropy = match cache.ll_previous.as_ref() {
969                Some(PreviousFseTable::Custom(table)) => Some(table.as_ref()),
970                _ => None,
971            };
972            let ml_entropy = match cache.ml_previous.as_ref() {
973                Some(PreviousFseTable::Custom(table)) => Some(table.as_ref()),
974                _ => None,
975            };
976            let of_entropy = match cache.of_previous.as_ref() {
977                Some(PreviousFseTable::Custom(table)) => Some(table.as_ref()),
978                _ => None,
979            };
980            self.state.matcher.seed_dictionary_entropy(
981                self.state.last_huff_table.as_ref(),
982                ll_entropy,
983                ml_entropy,
984                of_entropy,
985            );
986        } else {
987            self.state.last_huff_table = None;
988            self.state.fse_tables.ll_previous = None;
989            self.state.fse_tables.ml_previous = None;
990            self.state.fse_tables.of_previous = None;
991        }
992        self.savings = 0;
993        #[cfg(feature = "hash")]
994        {
995            self.hasher = XxHash64::with_seed(0);
996        }
997
998        let window_size = self.state.matcher.window_size();
999        if window_size == 0 {
1000            return Err(invalid_input_error(
1001                "matcher reported window_size == 0, which is invalid",
1002            ));
1003        }
1004
1005        // Single-segment is incompatible with a dictionary (the dictionary
1006        // pushes referenceable history before the content, so the frame needs
1007        // an explicit window descriptor); gate it off when a dict is attached,
1008        // mirroring `FrameCompressor`'s `!use_dictionary_state` guard.
1009        // Single-segment also requires the FCS field to be present
1010        // (`content_size_flag`): the layout drops the window descriptor,
1011        // so the header must carry the content size for decoders to size
1012        // their window.
1013        let single_segment = self.content_size_flag
1014            && !use_dictionary_state
1015            && self
1016                .pledged_content_size
1017                .map(|size| (512..=(1 << 14)).contains(&size) && size <= window_size)
1018                .unwrap_or(false);
1019
1020        let header = FrameHeader {
1021            frame_content_size: if self.content_size_flag {
1022                self.pledged_content_size
1023            } else {
1024                None
1025            },
1026            single_segment,
1027            content_checksum: cfg!(feature = "hash") && self.content_checksum,
1028            dictionary_id: if use_dictionary_state && self.dictionary_id_flag {
1029                // Id 0 is a raw-content dictionary: RFC 8878 spells "no
1030                // dictionary ID" as an absent field, not as a stored zero.
1031                self.dictionary
1032                    .as_ref()
1033                    .map(|dict| dict.inner.id)
1034                    .filter(|id| *id != 0)
1035                    .map(u64::from)
1036            } else {
1037                None
1038            },
1039            window_size: if single_segment {
1040                None
1041            } else {
1042                Some(window_size)
1043            },
1044            magicless: self.magicless,
1045        };
1046        let mut encoded_header = Vec::new();
1047        header.serialize(&mut encoded_header);
1048        drain
1049            .write_all(&encoded_header)
1050            .map_err(|err| self.fail(err))?;
1051
1052        if matches!(self.compression_level, CompressionLevel::Uncompressed) {
1053            // The raw block's header goes in front of its input once the
1054            // block's length is known.
1055            self.encoded_scratch.clear();
1056            self.encoded_scratch
1057                .reserve(BLOCK_HEADER_LEN + self.block_capacity());
1058            self.encoded_scratch
1059                .extend_from_slice(&[0; BLOCK_HEADER_LEN]);
1060        }
1061        self.frame_started = true;
1062        Ok(())
1063    }
1064
1065    fn block_capacity(&self) -> usize {
1066        let matcher_window = self.state.matcher.window_size() as usize;
1067        let ceiling = self
1068            .target_block_size
1069            .map_or(MAX_BLOCK_SIZE as usize, |t| t as usize);
1070        core::cmp::max(1, core::cmp::min(matcher_window, ceiling))
1071    }
1072
1073    /// Input taken for the frame that no block has carried out yet.
1074    fn pending_len(&self) -> usize {
1075        if !self.frame_started {
1076            return 0;
1077        }
1078        match self.compression_level {
1079            CompressionLevel::Uncompressed => self.encoded_scratch.len() - BLOCK_HEADER_LEN,
1080            _ => self.state.matcher.uncommitted_input().len(),
1081        }
1082    }
1083
1084    /// Buffers `bytes` for the block being filled: behind the raw block's
1085    /// header, or uncommitted in the match finder's history, where the block
1086    /// is later matched without another copy.
1087    fn buffer_input(&mut self, bytes: &[u8]) {
1088        match self.compression_level {
1089            CompressionLevel::Uncompressed => self.encoded_scratch.extend_from_slice(bytes),
1090            _ => {
1091                self.state
1092                    .matcher
1093                    .fill_in_place(bytes.len(), &mut |history| {
1094                        history.extend_from_slice(bytes);
1095                        (bytes.len(), false)
1096                    });
1097            }
1098        }
1099    }
1100
1101    /// Where the full pending block is cut (upstream `ZSTD_compress_frameChunk`
1102    /// sizing every block with `ZSTD_optimalBlockSize`): the pre-splitter's
1103    /// boundary once the frame has saved enough, else the whole block.
1104    /// `remaining` is the input still to come as far as the splitter knows:
1105    /// a full block again while writes continue, the buffered bytes at the
1106    /// end of the frame.
1107    fn pre_split_len(&self, block_capacity: usize, remaining: usize) -> usize {
1108        if matches!(self.compression_level, CompressionLevel::Uncompressed) {
1109            return self.pending_len();
1110        }
1111        let pending = self.state.matcher.uncommitted_input();
1112        crate::encoding::frame_compressor::optimal_block_size_with(
1113            self.state.pre_split.map(usize::from),
1114            pending,
1115            remaining,
1116            block_capacity,
1117            self.savings,
1118        )
1119        .min(pending.len())
1120    }
1121
1122    /// Emit the pending block once it is full, cut where the pre-splitter
1123    /// says; the suffix stays pending and heads the next block. A failure
1124    /// leaves the context failed.
1125    fn emit_full_pending_block<D: Write + ?Sized>(
1126        &mut self,
1127        drain: &mut D,
1128        block_capacity: usize,
1129    ) -> Result<(), Error> {
1130        if self.pending_len() != block_capacity {
1131            return Ok(());
1132        }
1133        let block_len = self.pre_split_len(block_capacity, block_capacity);
1134        self.encode_block(drain, block_len, false)
1135            .map_err(|err| self.fail(err))
1136    }
1137
1138    fn emit_pending_block<D: Write + ?Sized>(
1139        &mut self,
1140        drain: &mut D,
1141        last_block: bool,
1142    ) -> Result<(), Error> {
1143        let block_capacity = self.block_capacity();
1144        if last_block {
1145            // A full final buffer is cut like any other block (the reader
1146            // path splits it with `remaining = len`); each cut prefix goes
1147            // out as a non-last block and the suffix is re-examined.
1148            while self.pending_len() == block_capacity {
1149                let block_len = self.pre_split_len(block_capacity, block_capacity);
1150                if block_len == block_capacity {
1151                    break;
1152                }
1153                self.encode_block(drain, block_len, false)
1154                    .map_err(|err| self.fail(err))?;
1155            }
1156        }
1157        let block_len = self.pending_len();
1158        self.encode_block(drain, block_len, last_block)
1159            .map_err(|err| self.fail(err))
1160    }
1161
1162    // Exhaustive match kept intentionally: adding a new CompressionLevel
1163    // variant will produce a compile error here, forcing the developer to
1164    // decide whether the streaming encoder supports it before shipping.
1165    fn ensure_level_supported(&self) -> Result<(), Error> {
1166        match self.compression_level {
1167            CompressionLevel::Uncompressed
1168            | CompressionLevel::Fastest
1169            | CompressionLevel::Default
1170            | CompressionLevel::Better
1171            | CompressionLevel::Best
1172            | CompressionLevel::Level(_) => Ok(()),
1173        }
1174    }
1175
1176    /// Writes the first `block_len` pending bytes to `drain` as one block,
1177    /// `last_block` if it closes the frame.
1178    fn encode_block<D: Write + ?Sized>(
1179        &mut self,
1180        drain: &mut D,
1181        block_len: usize,
1182        last_block: bool,
1183    ) -> Result<(), Error> {
1184        if matches!(self.compression_level, CompressionLevel::Uncompressed) {
1185            // The block was assembled behind its header as it was written.
1186            debug_assert_eq!(block_len, self.pending_len());
1187            let header = BlockHeader {
1188                last_block,
1189                block_type: crate::blocks::block::BlockType::Raw,
1190                // At most one block, which a block header's 21-bit size holds.
1191                block_size: block_len as u32,
1192            };
1193            self.encoded_scratch[..BLOCK_HEADER_LEN].copy_from_slice(&header.to_le_bytes());
1194            drain.write_all(&self.encoded_scratch)?;
1195            #[cfg(feature = "hash")]
1196            if self.content_checksum {
1197                self.hasher.write(&self.encoded_scratch[BLOCK_HEADER_LEN..]);
1198            }
1199            self.encoded_scratch.truncate(BLOCK_HEADER_LEN);
1200            return Ok(());
1201        }
1202        self.encoded_scratch.clear();
1203        if block_len == 0 {
1204            // Only the last block of a frame is ever empty.
1205            let header = BlockHeader {
1206                last_block,
1207                block_type: crate::blocks::block::BlockType::Raw,
1208                block_size: 0,
1209            };
1210            header.serialize(&mut self.encoded_scratch);
1211            return drain.write_all(&self.encoded_scratch);
1212        }
1213        let dict_active =
1214            self.dictionary.is_some() && self.state.matcher.supports_dictionary_priming();
1215        compress_block_encoded(
1216            &mut self.state,
1217            self.compression_level,
1218            last_block,
1219            block_len,
1220            &mut self.encoded_scratch,
1221            dict_active,
1222            // No FrameEmitInfo on the streaming encoder path: it does not
1223            // surface per-block layout, so no sidecar.
1224            #[cfg(feature = "lsm")]
1225            None,
1226            #[cfg(all(feature = "lsm", feature = "hash"))]
1227            None,
1228        );
1229        // The block is committed to the match finder whether or not the drain
1230        // takes it; a drain that fails leaves the context failed.
1231        drain.write_all(&self.encoded_scratch)?;
1232        // `savings` counts the block header too, as upstream's
1233        // `ZSTD_compress_frameChunk` does (`cSize` includes it).
1234        self.savings += block_len as i64 - self.encoded_scratch.len() as i64;
1235        #[cfg(feature = "hash")]
1236        if self.content_checksum {
1237            self.hasher.write(self.state.matcher.get_last_space());
1238        }
1239        Ok(())
1240    }
1241
1242    fn fail(&mut self, err: Error) -> Error {
1243        self.errored = true;
1244        if self.last_error_kind.is_none() {
1245            self.last_error_kind = Some(err.kind());
1246        }
1247        if self.last_error_message.is_none() {
1248            self.last_error_message = Some(err.to_string());
1249        }
1250        err
1251    }
1252}
1253
1254/// Bytes of a block header (RFC 8878 3.1.1.2).
1255const BLOCK_HEADER_LEN: usize = 3;
1256
1257fn error_from_kind(kind: ErrorKind) -> Error {
1258    Error::from(kind)
1259}
1260
1261fn error_with_kind_message(kind: ErrorKind, message: String) -> Error {
1262    #[cfg(feature = "std")]
1263    {
1264        Error::new(kind, message)
1265    }
1266    #[cfg(not(feature = "std"))]
1267    {
1268        Error::new(kind, alloc::boxed::Box::new(message))
1269    }
1270}
1271
1272fn invalid_input_error(message: &str) -> Error {
1273    #[cfg(feature = "std")]
1274    {
1275        Error::new(ErrorKind::InvalidInput, message)
1276    }
1277    #[cfg(not(feature = "std"))]
1278    {
1279        Error::new(
1280            ErrorKind::Other,
1281            alloc::boxed::Box::new(alloc::string::String::from(message)),
1282        )
1283    }
1284}
1285
1286fn other_error_owned(message: String) -> Error {
1287    #[cfg(feature = "std")]
1288    {
1289        Error::other(message)
1290    }
1291    #[cfg(not(feature = "std"))]
1292    {
1293        Error::new(ErrorKind::Other, alloc::boxed::Box::new(message))
1294    }
1295}
1296
1297fn other_error(message: &str) -> Error {
1298    #[cfg(feature = "std")]
1299    {
1300        Error::other(message)
1301    }
1302    #[cfg(not(feature = "std"))]
1303    {
1304        Error::new(
1305            ErrorKind::Other,
1306            alloc::boxed::Box::new(alloc::string::String::from(message)),
1307        )
1308    }
1309}
1310
1311#[cfg(test)]
1312mod tests;