Skip to main content

structured_zstd/encoding/
streaming_encoder.rs

1use alloc::format;
2use alloc::string::{String, ToString};
3use alloc::vec::Vec;
4use core::mem;
5
6use crate::common::MAX_BLOCK_SIZE;
7#[cfg(feature = "hash")]
8use core::hash::Hasher;
9#[cfg(feature = "hash")]
10use twox_hash::XxHash64;
11
12use crate::encoding::levels::compress_block_encoded;
13use crate::encoding::{
14    CompressionLevel, EncoderDictionary, MatchGeneratorDriver, Matcher, block_header::BlockHeader,
15    frame_compressor::CachedDictionaryEntropy, frame_compressor::CompressState,
16    frame_compressor::FseTables, frame_compressor::PreviousFseTable, frame_header::FrameHeader,
17};
18use crate::io::{Error, ErrorKind, Write};
19
20/// Incremental frame encoder that implements [`Write`].
21///
22/// Data can be provided with multiple `write()` calls. Full blocks are compressed
23/// automatically, `flush()` emits the currently buffered partial block as non-last,
24/// and `finish()` closes the frame and returns the wrapped writer.
25pub struct StreamingEncoder<W: Write, M: Matcher = MatchGeneratorDriver> {
26    drain: Option<W>,
27    compression_level: CompressionLevel,
28    state: CompressState<M>,
29    pending: Vec<u8>,
30    encoded_scratch: Vec<u8>,
31    errored: bool,
32    last_error_kind: Option<ErrorKind>,
33    last_error_message: Option<String>,
34    frame_started: bool,
35    /// Upper bound on emitted block sizes (upstream `ZSTD_c_targetCBlockSize`
36    /// semantics; see `FrameCompressor::set_target_block_size`). `None` =
37    /// the format's 128 KiB ceiling.
38    target_block_size: Option<u32>,
39    pledged_content_size: Option<u64>,
40    /// Advisory source-size hint from [`set_source_size_hint`](Self::set_source_size_hint).
41    /// Unlike `pledged_content_size` it carries no end-of-frame enforcement, but
42    /// it still feeds the small-input gates (matcher sizing AND the Fast HUF
43    /// fast-path gate) so `set_source_size_hint(small)` reduces work the same way
44    /// a pledge does. The HUF gate reads `pledged_content_size.or(source_size_hint)`.
45    source_size_hint: Option<u64>,
46    /// Whether a pledged size is written into the header's
47    /// `Frame_Content_Size` field (upstream `ZSTD_c_contentSizeFlag`).
48    /// Pledge *enforcement* is independent of this flag — upstream
49    /// validates consumed bytes against the pledge at frame end even
50    /// when the header omits the field. Default `true`.
51    content_size_flag: bool,
52    bytes_consumed: u64,
53    /// Upstream `ZSTD_compress_frameChunk` `savings`: bytes consumed minus
54    /// bytes produced so far in this frame; the block pre-splitter only cuts
55    /// full blocks once the frame has saved enough.
56    savings: i64,
57    /// Effective strategy tag (and lazy depth) when a public-parameter
58    /// [`Strategy`](crate::encoding::Strategy) override (#27) is active, mirroring
59    /// [`FrameCompressor`](crate::encoding::FrameCompressor)'s field. `Some`
60    /// survives frame start so the literal-compression gates and the block
61    /// splitter run the same strategy the matcher does; `None` keeps the
62    /// level-derived tag.
63    strategy_override: Option<(crate::encoding::strategy::StrategyTag, u8)>,
64    /// Public `target_length` override (#27), kept so the raw-literals gate
65    /// resolved at frame start reads the value the matcher runs (dropped on a
66    /// dictionary frame, where the CDict's targetLength applies).
67    target_length_override: Option<u32>,
68    /// `ZSTD_f_zstd1_magicless` — omit the 4-byte magic number prefix.
69    /// Default false. See [`Self::set_magicless`].
70    magicless: bool,
71    /// Whether to emit a trailing XXH64 content checksum and set the frame
72    /// header's `Content_Checksum_flag` (upstream `ZSTD_c_checksumFlag`).
73    /// Default `false`, matching the upstream library default; combined with
74    /// the `hash` feature, so without `hash` no checksum is emitted
75    /// regardless. See [`Self::set_content_checksum`].
76    content_checksum: bool,
77    /// Dictionary applied to the frame (upstream zstd `ZSTD_CCtx_loadDictionary` on a
78    /// streaming context). `None` = no dictionary. Set before the first write.
79    dictionary: Option<EncoderDictionary>,
80    /// Whether the frame header records the attached dictionary's ID
81    /// (upstream `ZSTD_c_dictIDFlag`). Default `true`. Raw-content
82    /// dictionaries (upstream `ZSTD_CCtx_refPrefix`) carry a synthetic
83    /// non-zero ID that must not reach the wire, so their attach path
84    /// turns this off. See [`Self::set_dictionary_id_flag`].
85    dictionary_id_flag: bool,
86    /// Encoder entropy tables (literals Huffman + LL/ML/OF FSE "previous"
87    /// tables) the dictionary seeds into the first block, derived once when the
88    /// dictionary is attached so each frame start is a cheap clone.
89    dictionary_entropy_cache: Option<CachedDictionaryEntropy>,
90    #[cfg(feature = "hash")]
91    hasher: XxHash64,
92}
93
94impl<W: Write> StreamingEncoder<W, MatchGeneratorDriver> {
95    /// Creates a streaming encoder backed by the default match generator.
96    ///
97    /// The encoder writes compressed bytes into `drain` and applies `compression_level`
98    /// to all subsequently written blocks.
99    pub fn new(drain: W, compression_level: CompressionLevel) -> Self {
100        Self::new_with_matcher(
101            MatchGeneratorDriver::new(MAX_BLOCK_SIZE as usize, 1),
102            drain,
103            compression_level,
104        )
105    }
106
107    /// Configure fine-grained compression parameters (#27): resets the level to
108    /// the parameters' level and installs the per-knob overrides (window / hash
109    /// / chain / search logs, strategy, long-distance matching) applied at the
110    /// next frame. Mirrors [`FrameCompressor::set_parameters`]. Must be called
111    /// before the first [`write`](Write::write). Only the built-in
112    /// `MatchGeneratorDriver` exposes the override knobs, so this lives on the
113    /// default-matcher impl.
114    pub fn set_parameters(
115        &mut self,
116        params: &crate::encoding::CompressionParameters,
117    ) -> Result<(), Error> {
118        self.ensure_open()?;
119        if self.frame_started {
120            return Err(invalid_input_error(
121                "compression parameters must be set before the first write",
122            ));
123        }
124        self.compression_level = params.level();
125        let overrides = params.overrides();
126        // Persist the strategy override so `ensure_frame_started`'s level-based
127        // resync does not discard it (matching `FrameCompressor::set_parameters`).
128        self.strategy_override = overrides.strategy.map(|s| (s.tag(), s.lazy_depth()));
129        self.target_length_override = overrides.target_length;
130        self.state.strategy_tag = self.strategy_override.map_or_else(
131            || {
132                crate::encoding::strategy::StrategyTag::for_compression_level(
133                    self.compression_level,
134                )
135            },
136            |(tag, _)| tag,
137        );
138        self.state.huf_optimal_search = crate::encoding::frame_compressor::huf_search_enabled(
139            self.state.strategy_tag,
140            self.pledged_content_size.or(self.source_size_hint),
141        );
142        self.state.matcher.set_param_overrides(Some(overrides));
143        Ok(())
144    }
145}
146
147impl<W: Write, M: Matcher> StreamingEncoder<W, M> {
148    /// Creates a streaming encoder with an explicitly provided matcher implementation.
149    ///
150    /// This constructor is primarily intended for tests and advanced callers that need
151    /// custom match-window behavior.
152    pub fn new_with_matcher(matcher: M, drain: W, compression_level: CompressionLevel) -> Self {
153        Self {
154            drain: Some(drain),
155            compression_level,
156            state: CompressState {
157                matcher,
158                copy_tier: crate::decoding::simd_copy::ExactCopyTier::resolve(),
159                last_huff_table: None,
160                huff_table_spare: None,
161                huff_weights: Default::default(),
162                fse_tables: FseTables::new(),
163                block_scratch: crate::encoding::blocks::CompressedBlockScratch::new(),
164                offset_hist: [1, 4, 8],
165                strategy_tag: crate::encoding::strategy::StrategyTag::for_compression_level(
166                    compression_level,
167                ),
168                pre_split: crate::encoding::levels::config::level_pre_split(compression_level)
169                    .map(|tier| tier as u8),
170                huf_optimal_search: true,
171                literal_compression_disabled: matches!(
172                    compression_level,
173                    CompressionLevel::Level(n) if n < 0
174                ),
175            },
176            pending: Vec::new(),
177            encoded_scratch: Vec::new(),
178            errored: false,
179            last_error_kind: None,
180            last_error_message: None,
181            frame_started: false,
182            target_block_size: None,
183            pledged_content_size: None,
184            source_size_hint: None,
185            content_size_flag: true,
186            bytes_consumed: 0,
187            savings: 0,
188            strategy_override: None,
189            target_length_override: None,
190            magicless: false,
191            content_checksum: false,
192            dictionary: None,
193            dictionary_id_flag: true,
194            dictionary_entropy_cache: None,
195            #[cfg(feature = "hash")]
196            hasher: XxHash64::with_seed(0),
197        }
198    }
199
200    /// Set an upper bound on each physical block's payload (semantics of
201    /// upstream `ZSTD_c_targetCBlockSize`): every block carries at most
202    /// `target` payload bytes, +3-byte block header on the wire — the
203    /// upstream knob is likewise a convergence target for block sizing,
204    /// not a cap on header-inclusive wire bytes. Clamped to
205    /// `[MIN_TARGET_BLOCK_SIZE, MAX_BLOCK_SIZE]`; mirrors
206    /// `FrameCompressor::set_target_block_size`. Must be set before the
207    /// first write.
208    pub fn set_target_block_size(&mut self, target: Option<u32>) -> Result<(), Error> {
209        self.ensure_open()?;
210        if self.frame_started {
211            return Err(invalid_input_error(
212                "the block-size target must be set before the first write",
213            ));
214        }
215        self.target_block_size = target.map(|t| {
216            t.clamp(
217                crate::common::MIN_TARGET_BLOCK_SIZE,
218                crate::common::MAX_BLOCK_SIZE,
219            )
220        });
221        Ok(())
222    }
223
224    /// Enable or disable the trailing XXH64 content checksum
225    /// (upstream `ZSTD_c_checksumFlag`). Default `false`, matching the
226    /// upstream library default (`ZSTD_c_checksumFlag = 0`). Must be called
227    /// before the first [`write`](Write::write); once the frame header is
228    /// emitted the flag is fixed, so a late change returns an error rather
229    /// than producing a header/trailer mismatch. Without the `hash` feature
230    /// no checksum is emitted regardless.
231    pub fn set_content_checksum(&mut self, emit: bool) -> Result<(), Error> {
232        self.ensure_open()?;
233        if self.frame_started {
234            return Err(invalid_input_error(
235                "content checksum must be set before the first write",
236            ));
237        }
238        self.content_checksum = emit;
239        Ok(())
240    }
241
242    /// Enable or disable magicless frame format (`ZSTD_f_zstd1_magicless`).
243    ///
244    /// When set to `true`, the frame header serialized by this encoder
245    /// omits the 4-byte magic number prefix. Must be called BEFORE the
246    /// first [`write`](Write::write) call; calling it after the frame
247    /// header has already been emitted returns an error so the caller
248    /// can't be misled into thinking they produced a magicless stream.
249    pub fn set_magicless(&mut self, magicless: bool) -> Result<(), Error> {
250        self.ensure_open()?;
251        if self.frame_started {
252            return Err(invalid_input_error(
253                "magicless format must be set before the first write",
254            ));
255        }
256        self.magicless = magicless;
257        Ok(())
258    }
259
260    /// Pledge the total uncompressed content size for this frame.
261    ///
262    /// When set, the frame header will include a `Frame_Content_Size` field.
263    /// This enables decoders to pre-allocate output buffers.
264    /// The pledged size is also forwarded as a source-size hint to the
265    /// matcher so small inputs can use smaller matching tables.
266    ///
267    /// Must be called **before** the first [`write`](Write::write) call;
268    /// calling it after the frame header has already been emitted returns an
269    /// error.
270    pub fn set_pledged_content_size(&mut self, size: u64) -> Result<(), Error> {
271        self.ensure_open()?;
272        if self.frame_started {
273            return Err(invalid_input_error(
274                "pledged content size must be set before the first write",
275            ));
276        }
277        self.pledged_content_size = Some(size);
278        // Also use pledged size as source-size hint so the matcher
279        // can select smaller tables for small inputs.
280        self.state.matcher.set_source_size_hint(size);
281        Ok(())
282    }
283
284    /// Control whether the pledged size is written into the header's
285    /// `Frame_Content_Size` field (upstream `ZSTD_c_contentSizeFlag`,
286    /// default on). With the flag off the header omits the field, but a
287    /// pledge set via [`set_pledged_content_size`](Self::set_pledged_content_size)
288    /// is still enforced against the bytes actually written. Must be
289    /// called before the first [`write`](Write::write).
290    pub fn set_content_size_flag(&mut self, emit: bool) -> Result<(), Error> {
291        self.ensure_open()?;
292        if self.frame_started {
293            return Err(invalid_input_error(
294                "content size flag must be set before the first write",
295            ));
296        }
297        self.content_size_flag = emit;
298        Ok(())
299    }
300
301    /// Provide a hint about the total uncompressed size for the next frame.
302    ///
303    /// Unlike [`set_pledged_content_size`](Self::set_pledged_content_size),
304    /// this does **not** enforce that exactly `size` bytes are written; it
305    /// may reduce matcher tables, advertised frame window, and block sizing
306    /// for small inputs. Must be called before the first
307    /// [`write`](Write::write).
308    pub fn set_source_size_hint(&mut self, size: u64) -> Result<(), Error> {
309        self.ensure_open()?;
310        if self.frame_started {
311            return Err(invalid_input_error(
312                "source size hint must be set before the first write",
313            ));
314        }
315        self.state.matcher.set_source_size_hint(size);
316        // Feed the same hint to the Fast HUF fast-path gate (resolved in
317        // `set_parameters` / `ensure_frame_started` via
318        // `pledged_content_size.or(source_size_hint)`), so a small advisory size
319        // also lifts Fast streams off the expensive optimal-HUF search.
320        self.source_size_hint = Some(size);
321        Ok(())
322    }
323
324    /// Attach a serialized dictionary blob to the frame (upstream zstd
325    /// `ZSTD_CCtx_loadDictionary` on a streaming context). The dictionary primes
326    /// the match-finder and seeds the first block's entropy tables + repeat
327    /// offsets, and its ID is written into the frame header. Must be called
328    /// before the first [`write`](Write::write); the parsed dictionary must have
329    /// a non-zero ID and non-zero repeat offsets.
330    pub fn set_dictionary_from_bytes(&mut self, raw_dictionary: &[u8]) -> Result<(), Error> {
331        let dict = EncoderDictionary::from_bytes(raw_dictionary)
332            .map_err(|err| invalid_input_error(&alloc::format!("invalid dictionary: {err:?}")))?;
333        self.set_encoder_dictionary(dict)
334    }
335
336    /// Whether the frame header records the dictionary ID when a dictionary
337    /// is attached (upstream `ZSTD_c_dictIDFlag` semantics; default `true`).
338    /// Mirrors [`FrameCompressor::set_dictionary_id_flag`]. Decoders can still
339    /// decode such frames by supplying the dictionary explicitly.
340    pub fn set_dictionary_id_flag(&mut self, emit: bool) -> Result<(), Error> {
341        self.ensure_open()?;
342        if self.frame_started {
343            return Err(invalid_input_error(
344                "dictionary ID flag must be set before the first write",
345            ));
346        }
347        self.dictionary_id_flag = emit;
348        Ok(())
349    }
350
351    /// Attach an already-parsed [`EncoderDictionary`] to the frame. See
352    /// [`set_dictionary_from_bytes`](Self::set_dictionary_from_bytes); must be
353    /// called before the first write.
354    pub fn set_encoder_dictionary(&mut self, dict: EncoderDictionary) -> Result<(), Error> {
355        self.ensure_open()?;
356        if self.frame_started {
357            return Err(invalid_input_error(
358                "dictionary must be attached before the first write",
359            ));
360        }
361        // A zero id marks a raw-content dictionary, which carries no header to
362        // hold one; the frame then records no dictionary ID and the decoder
363        // must be given the same bytes explicitly.
364        let inner = &dict.inner;
365        if inner.offset_hist.contains(&0) {
366            return Err(invalid_input_error(
367                "dictionary carries a zero repeat offset",
368            ));
369        }
370        self.dictionary_entropy_cache = Some(CachedDictionaryEntropy::from_dictionary(inner));
371        self.dictionary = Some(dict);
372        Ok(())
373    }
374
375    /// Returns an immutable reference to the wrapped output drain.
376    ///
377    /// The drain remains available for the encoder lifetime; [`finish`](Self::finish)
378    /// consumes the encoder and returns ownership of the drain.
379    pub fn get_ref(&self) -> &W {
380        self.drain
381            .as_ref()
382            .expect("streaming encoder drain is present until finish consumes self")
383    }
384
385    /// Total heap bytes this encoder's allocations hold, excluding the
386    /// inline struct and the drain `W` (whose footprint the owner can
387    /// measure through [`get_ref`](Self::get_ref)): match-finder tables /
388    /// history / recycled buffers, retained Huffman tables, the staging
389    /// `pending` / `encoded_scratch` buffers, the retained dictionary
390    /// content, and the cached dictionary entropy tables. Mirrors
391    /// `FrameCompressor::heap_size` so a context can report its true
392    /// footprint through `ZSTD_sizeof_CCtx`.
393    pub fn heap_size(&self) -> usize {
394        let mut total = self.state.matcher.heap_size();
395        total += self
396            .state
397            .last_huff_table
398            .as_ref()
399            .map_or(0, |table| table.heap_size());
400        total += self
401            .state
402            .huff_table_spare
403            .as_ref()
404            .map_or(0, |table| table.heap_size());
405        // Kept between blocks and frames; see `FrameCompressor::heap_size`.
406        total += self.state.huff_weights.heap_size();
407        total += self.pending.capacity();
408        total += self.encoded_scratch.capacity();
409        total += self
410            .dictionary
411            .as_ref()
412            .map_or(0, |d| d.inner.dict_content.capacity());
413        total += self
414            .dictionary_entropy_cache
415            .as_ref()
416            .map_or(0, CachedDictionaryEntropy::heap_size);
417        total
418    }
419
420    /// Returns a mutable reference to the wrapped output drain.
421    ///
422    /// It is inadvisable to directly write to the underlying writer, as doing
423    /// so would corrupt the zstd frame being assembled by the encoder.
424    ///
425    /// The drain remains available for the encoder lifetime; [`finish`](Self::finish)
426    /// consumes the encoder and returns ownership of the drain.
427    pub fn get_mut(&mut self) -> &mut W {
428        self.drain
429            .as_mut()
430            .expect("streaming encoder drain is present until finish consumes self")
431    }
432
433    /// Finalizes the current zstd frame and returns the wrapped output drain.
434    ///
435    /// If no payload was written yet, this still emits a valid empty frame.
436    /// Calling this method consumes the encoder.
437    pub fn finish(mut self) -> Result<W, Error> {
438        self.ensure_open()?;
439
440        // Validate the pledge before finalizing the frame. If finish() is
441        // called before any writes, this also avoids emitting a header with
442        // an incorrect FCS into the drain on mismatch.
443        if let Some(pledged) = self.pledged_content_size
444            && self.bytes_consumed != pledged
445        {
446            return Err(invalid_input_error(
447                "pledged content size does not match bytes consumed",
448            ));
449        }
450
451        self.ensure_frame_started()?;
452
453        if self.pending.is_empty() {
454            self.write_empty_last_block()
455                .map_err(|err| self.fail(err))?;
456        } else {
457            self.emit_pending_block(true)?;
458        }
459
460        let mut drain = self
461            .drain
462            .take()
463            .expect("streaming encoder drain must be present when finishing");
464
465        #[cfg(feature = "hash")]
466        if self.content_checksum {
467            let checksum = self.hasher.finish() as u32;
468            drain
469                .write_all(&checksum.to_le_bytes())
470                .map_err(|err| self.fail(err))?;
471        }
472
473        drain.flush().map_err(|err| self.fail(err))?;
474        Ok(drain)
475    }
476
477    fn ensure_open(&self) -> Result<(), Error> {
478        if self.errored {
479            return Err(self.sticky_error());
480        }
481        Ok(())
482    }
483
484    // Cold path (only reached after poisoning). The format!() calls still allocate
485    // in no_std even though error_with_kind_message/other_error_owned drop the
486    // message; this is acceptable on an error recovery path to keep match arms simple.
487    fn sticky_error(&self) -> Error {
488        match (self.last_error_kind, self.last_error_message.as_deref()) {
489            (Some(kind), Some(message)) => error_with_kind_message(
490                kind,
491                format!(
492                    "streaming encoder is in an errored state due to previous {kind:?} failure: {message}"
493                ),
494            ),
495            (Some(kind), None) => error_from_kind(kind),
496            (None, Some(message)) => other_error_owned(format!(
497                "streaming encoder is in an errored state: {message}"
498            )),
499            (None, None) => other_error("streaming encoder is in an errored state"),
500        }
501    }
502
503    fn drain_mut(&mut self) -> Result<&mut W, Error> {
504        self.drain
505            .as_mut()
506            .ok_or_else(|| other_error("streaming encoder has no active drain"))
507    }
508
509    fn ensure_frame_started(&mut self) -> Result<(), Error> {
510        if self.frame_started {
511            return Ok(());
512        }
513
514        self.ensure_level_supported()?;
515        // A dictionary is only active when it can actually be primed: the level
516        // compresses (not `Uncompressed`) AND the matcher supports priming AND a
517        // dictionary is attached. Mirrors `FrameCompressor`'s `use_dictionary_state`
518        // so a streaming frame never advertises a `Dictionary_ID`, disables
519        // single-segment, or seeds dict entropy/offsets unless the dictionary is
520        // genuinely in play (otherwise it would emit frames that needlessly
521        // require a dictionary at decode time).
522        let use_dictionary_state =
523            !matches!(self.compression_level, CompressionLevel::Uncompressed)
524                && self.state.matcher.supports_dictionary_priming()
525                && self.dictionary.is_some();
526        // The dictionary sizes select the CDict cParams tier (consumed inside
527        // `reset`), so hand them over BEFORE reset.
528        if use_dictionary_state && let Some(dict) = self.dictionary.as_ref() {
529            self.state.matcher.set_dictionary_size_hint(dict.sizes());
530        }
531        // The matcher resolves the frame from the LAST hint it was handed, so
532        // re-forward the authoritative size (`pledge.or(advisory)`, the same
533        // value the gates below read) right before the reset: without this a
534        // pledge followed by a different advisory hint (or vice versa) left
535        // the matcher and the frame gates on different size tiers.
536        if let Some(size) = self.pledged_content_size.or(self.source_size_hint) {
537            self.state.matcher.set_source_size_hint(size);
538        }
539        self.state.matcher.reset(self.compression_level);
540        // Seed the repeat-offset history from the dictionary (upstream zstd
541        // `ZSTD_compress_insertDictionary`), or the default rep codes otherwise.
542        self.state.offset_hist = if use_dictionary_state {
543            self.dictionary
544                .as_ref()
545                .map(|dict| dict.inner.offset_hist)
546                .unwrap_or([1, 4, 8])
547        } else {
548            [1, 4, 8]
549        };
550        // Prime the match-finder with the dictionary content + offsets.
551        // `dict` borrows `self.dictionary`; `self.state.matcher` is a disjoint
552        // field, so the immutable dict borrow and the mutable matcher borrow
553        // coexist (field-level borrow splitting) with no conflict.
554        if use_dictionary_state && let Some(dict) = self.dictionary.as_ref() {
555            let offset_hist = dict.inner.offset_hist;
556            self.state
557                .matcher
558                .prime_with_dictionary(dict.inner.dict_content.as_slice(), offset_hist);
559        }
560        // Seed the first block's entropy from the dictionary's cached encoder
561        // tables (upstream zstd `cdict->cBlockState`), or clear to defaults.
562        if use_dictionary_state && let Some(cache) = self.dictionary_entropy_cache.as_ref() {
563            self.state.last_huff_table.clone_from(&cache.huff);
564            self.state
565                .fse_tables
566                .ll_previous
567                .clone_from(&cache.ll_previous);
568            self.state
569                .fse_tables
570                .ml_previous
571                .clone_from(&cache.ml_previous);
572            self.state
573                .fse_tables
574                .of_previous
575                .clone_from(&cache.of_previous);
576            let ll_entropy = match cache.ll_previous.as_ref() {
577                Some(PreviousFseTable::Custom(table)) => Some(table.as_ref()),
578                _ => None,
579            };
580            let ml_entropy = match cache.ml_previous.as_ref() {
581                Some(PreviousFseTable::Custom(table)) => Some(table.as_ref()),
582                _ => None,
583            };
584            let of_entropy = match cache.of_previous.as_ref() {
585                Some(PreviousFseTable::Custom(table)) => Some(table.as_ref()),
586                _ => None,
587            };
588            self.state.matcher.seed_dictionary_entropy(
589                self.state.last_huff_table.as_ref(),
590                ll_entropy,
591                ml_entropy,
592                of_entropy,
593            );
594        } else {
595            self.state.last_huff_table = None;
596            self.state.fse_tables.ll_previous = None;
597            self.state.fse_tables.ml_previous = None;
598            self.state.fse_tables.of_previous = None;
599        }
600        // Sync `state.strategy_tag` / `state.pre_split` to the strategy the
601        // matcher's reset resolved (size- and dictionary-adaptive; a public
602        // strategy override wins on a plain frame, a dictionary frame runs the
603        // CDict's strategy) so the literal-compression gates and the block
604        // pre-splitter agree with the parse. Mirrors `FrameCompressor::compress`
605        // and keeps both entry points byte-equivalent.
606        let hint = self.pledged_content_size.or(self.source_size_hint);
607        let (params, dict_frame) = crate::encoding::frame_compressor::resolve_frame_params(
608            self.compression_level,
609            hint,
610            self.dictionary.as_ref().filter(|_| use_dictionary_state),
611        );
612        crate::encoding::frame_compressor::sync_effective_strategy(
613            &mut self.state,
614            self.compression_level,
615            &params,
616            self.strategy_override.filter(|_| !dict_frame),
617        );
618        self.state.huf_optimal_search =
619            crate::encoding::frame_compressor::huf_search_enabled(self.state.strategy_tag, hint);
620        self.state.literal_compression_disabled =
621            crate::encoding::frame_compressor::literal_compression_disabled(
622                self.state.strategy_tag,
623                self.compression_level,
624                self.target_length_override.filter(|_| !dict_frame),
625            );
626        self.savings = 0;
627        #[cfg(feature = "hash")]
628        {
629            self.hasher = XxHash64::with_seed(0);
630        }
631
632        let window_size = self.state.matcher.window_size();
633        if window_size == 0 {
634            return Err(invalid_input_error(
635                "matcher reported window_size == 0, which is invalid",
636            ));
637        }
638
639        // Single-segment is incompatible with a dictionary (the dictionary
640        // pushes referenceable history before the content, so the frame needs
641        // an explicit window descriptor); gate it off when a dict is attached,
642        // mirroring `FrameCompressor`'s `!use_dictionary_state` guard.
643        // Single-segment also requires the FCS field to be present
644        // (`content_size_flag`): the layout drops the window descriptor,
645        // so the header must carry the content size for decoders to size
646        // their window.
647        let single_segment = self.content_size_flag
648            && !use_dictionary_state
649            && self
650                .pledged_content_size
651                .map(|size| (512..=(1 << 14)).contains(&size) && size <= window_size)
652                .unwrap_or(false);
653
654        let header = FrameHeader {
655            frame_content_size: if self.content_size_flag {
656                self.pledged_content_size
657            } else {
658                None
659            },
660            single_segment,
661            content_checksum: cfg!(feature = "hash") && self.content_checksum,
662            dictionary_id: if use_dictionary_state && self.dictionary_id_flag {
663                // Id 0 is a raw-content dictionary: RFC 8878 spells "no
664                // dictionary ID" as an absent field, not as a stored zero.
665                self.dictionary
666                    .as_ref()
667                    .map(|dict| dict.inner.id)
668                    .filter(|id| *id != 0)
669                    .map(u64::from)
670            } else {
671                None
672            },
673            window_size: if single_segment {
674                None
675            } else {
676                Some(window_size)
677            },
678            magicless: self.magicless,
679        };
680        let mut encoded_header = Vec::new();
681        header.serialize(&mut encoded_header);
682        self.drain_mut()
683            .and_then(|drain| drain.write_all(&encoded_header))
684            .map_err(|err| self.fail(err))?;
685
686        self.frame_started = true;
687        Ok(())
688    }
689
690    fn block_capacity(&self) -> usize {
691        let matcher_window = self.state.matcher.window_size() as usize;
692        let ceiling = self
693            .target_block_size
694            .map_or(MAX_BLOCK_SIZE as usize, |t| t as usize);
695        core::cmp::max(1, core::cmp::min(matcher_window, ceiling))
696    }
697
698    fn allocate_pending_space(&mut self, block_capacity: usize) -> Vec<u8> {
699        let mut space = match self.compression_level {
700            CompressionLevel::Fastest
701            | CompressionLevel::Default
702            | CompressionLevel::Better
703            | CompressionLevel::Best
704            | CompressionLevel::Level(_) => self.state.matcher.get_next_space(),
705            CompressionLevel::Uncompressed => Vec::new(),
706        };
707        space.clear();
708        if space.capacity() > block_capacity {
709            space.shrink_to(block_capacity);
710        }
711        if space.capacity() < block_capacity {
712            space.reserve(block_capacity - space.capacity());
713        }
714        space
715    }
716
717    /// Where the full pending block is cut (upstream `ZSTD_compress_frameChunk`
718    /// sizing every block with `ZSTD_optimalBlockSize`): the pre-splitter's
719    /// boundary once the frame has saved enough, else the whole block.
720    /// `remaining` is the input still to come as far as the splitter knows:
721    /// a full block again while writes continue, the buffered bytes at the
722    /// end of the frame.
723    fn pre_split_len(&self, block_capacity: usize, remaining: usize) -> usize {
724        if matches!(self.compression_level, CompressionLevel::Uncompressed) {
725            return self.pending.len();
726        }
727        crate::encoding::frame_compressor::optimal_block_size_with(
728            self.state.pre_split.map(usize::from),
729            &self.pending,
730            remaining,
731            block_capacity,
732            self.savings,
733        )
734        .min(self.pending.len())
735    }
736
737    /// Emit the first `block_len` pending bytes as a non-last block; the
738    /// suffix stays pending (the next block starts with it, as the frame
739    /// compressor's reader path carries a pre-split suffix). On a drain
740    /// error the whole pending buffer is restored so no input is lost.
741    fn emit_pending_prefix(
742        &mut self,
743        block_len: usize,
744        block_capacity: usize,
745    ) -> Result<(), Error> {
746        let mut suffix = self.allocate_pending_space(block_capacity);
747        suffix.extend_from_slice(&self.pending[block_len..]);
748        let mut block = mem::replace(&mut self.pending, suffix);
749        block.truncate(block_len);
750        if let Err((err, mut restored_block)) = self.encode_block(block, false) {
751            restored_block.extend_from_slice(&self.pending);
752            self.pending = restored_block;
753            return Err(err);
754        }
755        Ok(())
756    }
757
758    fn emit_full_pending_block(
759        &mut self,
760        block_capacity: usize,
761        consumed: usize,
762    ) -> Option<Result<usize, Error>> {
763        if self.pending.len() != block_capacity {
764            return None;
765        }
766        let block_len = self.pre_split_len(block_capacity, block_capacity);
767        if let Err(err) = self.emit_pending_prefix(block_len, block_capacity) {
768            let err = self.fail(err);
769            if consumed > 0 {
770                return Some(Ok(consumed));
771            }
772            return Some(Err(err));
773        }
774        None
775    }
776
777    fn emit_pending_block(&mut self, last_block: bool) -> Result<(), Error> {
778        let block_capacity = self.block_capacity();
779        if last_block {
780            // A full final buffer is cut like any other block (the reader
781            // path splits it with `remaining = len`); each cut prefix goes
782            // out as a non-last block and the suffix is re-examined.
783            while self.pending.len() == block_capacity {
784                let block_len = self.pre_split_len(block_capacity, self.pending.len());
785                if block_len == self.pending.len() {
786                    break;
787                }
788                self.emit_pending_prefix(block_len, block_capacity)
789                    .map_err(|err| self.fail(err))?;
790            }
791        }
792        let block = mem::take(&mut self.pending);
793        if let Err((err, restored_block)) = self.encode_block(block, last_block) {
794            self.pending = restored_block;
795            return Err(self.fail(err));
796        }
797        if !last_block {
798            self.pending = self.allocate_pending_space(block_capacity);
799        }
800        Ok(())
801    }
802
803    // Exhaustive match kept intentionally: adding a new CompressionLevel
804    // variant will produce a compile error here, forcing the developer to
805    // decide whether the streaming encoder supports it before shipping.
806    fn ensure_level_supported(&self) -> Result<(), Error> {
807        match self.compression_level {
808            CompressionLevel::Uncompressed
809            | CompressionLevel::Fastest
810            | CompressionLevel::Default
811            | CompressionLevel::Better
812            | CompressionLevel::Best
813            | CompressionLevel::Level(_) => Ok(()),
814        }
815    }
816
817    fn encode_block(
818        &mut self,
819        uncompressed_data: Vec<u8>,
820        last_block: bool,
821    ) -> Result<(), (Error, Vec<u8>)> {
822        let mut raw_block = Some(uncompressed_data);
823        let mut encoded = Vec::new();
824        mem::swap(&mut encoded, &mut self.encoded_scratch);
825        encoded.clear();
826        let needed_capacity = self.block_capacity() + 3;
827        if encoded.capacity() < needed_capacity {
828            encoded.reserve(needed_capacity.saturating_sub(encoded.len()));
829        }
830        let mut moved_into_matcher = false;
831        let raw_len = raw_block.as_ref().map_or(0, Vec::len);
832        if raw_block.as_ref().is_some_and(|block| block.is_empty()) {
833            let header = BlockHeader {
834                last_block,
835                block_type: crate::blocks::block::BlockType::Raw,
836                block_size: 0,
837            };
838            header.serialize(&mut encoded);
839        } else {
840            match self.compression_level {
841                CompressionLevel::Uncompressed => {
842                    let block = raw_block.as_ref().expect("raw block missing");
843                    let header = BlockHeader {
844                        last_block,
845                        block_type: crate::blocks::block::BlockType::Raw,
846                        block_size: block.len() as u32,
847                    };
848                    header.serialize(&mut encoded);
849                    encoded.extend_from_slice(block);
850                }
851                CompressionLevel::Fastest
852                | CompressionLevel::Default
853                | CompressionLevel::Better
854                | CompressionLevel::Best
855                | CompressionLevel::Level(_) => {
856                    let block = raw_block.take().expect("raw block missing");
857                    debug_assert!(!block.is_empty(), "empty blocks handled above");
858                    compress_block_encoded(
859                        &mut self.state,
860                        self.compression_level,
861                        last_block,
862                        crate::encoding::levels::BlockInput::Staged(block),
863                        &mut encoded,
864                        // No FrameEmitInfo on the streaming encoder path — it
865                        // does not surface per-block layout, so no sidecar.
866                        #[cfg(feature = "lsm")]
867                        None,
868                        #[cfg(all(feature = "lsm", feature = "hash"))]
869                        None,
870                    );
871                    moved_into_matcher = true;
872                }
873            }
874        }
875
876        if let Err(err) = self.drain_mut().and_then(|drain| drain.write_all(&encoded)) {
877            encoded.clear();
878            mem::swap(&mut encoded, &mut self.encoded_scratch);
879            let restored = if moved_into_matcher {
880                self.state.matcher.get_last_space().to_vec()
881            } else {
882                raw_block.unwrap_or_default()
883            };
884            return Err((err, restored));
885        }
886        // `savings` counts the block header too, as upstream's
887        // `ZSTD_compress_frameChunk` does (`cSize` includes it).
888        self.savings += raw_len as i64 - encoded.len() as i64;
889
890        if moved_into_matcher {
891            #[cfg(feature = "hash")]
892            if self.content_checksum {
893                self.hasher.write(self.state.matcher.get_last_space());
894            }
895        } else {
896            self.hash_block(raw_block.as_deref().unwrap_or(&[]));
897        }
898        encoded.clear();
899        mem::swap(&mut encoded, &mut self.encoded_scratch);
900        Ok(())
901    }
902
903    fn write_empty_last_block(&mut self) -> Result<(), Error> {
904        self.encode_block(Vec::new(), true).map_err(|(err, _)| err)
905    }
906
907    fn fail(&mut self, err: Error) -> Error {
908        self.errored = true;
909        if self.last_error_kind.is_none() {
910            self.last_error_kind = Some(err.kind());
911        }
912        if self.last_error_message.is_none() {
913            self.last_error_message = Some(err.to_string());
914        }
915        err
916    }
917
918    #[cfg(feature = "hash")]
919    fn hash_block(&mut self, uncompressed_data: &[u8]) {
920        if self.content_checksum {
921            self.hasher.write(uncompressed_data);
922        }
923    }
924
925    #[cfg(not(feature = "hash"))]
926    fn hash_block(&mut self, _uncompressed_data: &[u8]) {}
927}
928
929impl<W: Write, M: Matcher> Write for StreamingEncoder<W, M> {
930    fn write(&mut self, buf: &[u8]) -> Result<usize, Error> {
931        self.ensure_open()?;
932        if buf.is_empty() {
933            return Ok(0);
934        }
935
936        // Check pledge before emitting the frame header so that a misuse
937        // like set_pledged_content_size(0) + write(non_empty) doesn't leave
938        // a partially-written header in the drain.
939        if let Some(pledged) = self.pledged_content_size
940            && self.bytes_consumed >= pledged
941        {
942            return Err(invalid_input_error(
943                "write would exceed pledged content size",
944            ));
945        }
946
947        self.ensure_frame_started()?;
948
949        // Enforce pledged upper bound: truncate the accepted slice to the
950        // remaining allowance so that partial-write semantics are honored
951        // (return Ok(n) with n < buf.len()) instead of failing the full call.
952        let buf = if let Some(pledged) = self.pledged_content_size {
953            let remaining_allowed = pledged
954                .checked_sub(self.bytes_consumed)
955                .ok_or_else(|| invalid_input_error("bytes consumed exceed pledged content size"))?;
956            if remaining_allowed == 0 {
957                return Err(invalid_input_error(
958                    "write would exceed pledged content size",
959                ));
960            }
961            let accepted = core::cmp::min(
962                buf.len(),
963                usize::try_from(remaining_allowed).unwrap_or(usize::MAX),
964            );
965            &buf[..accepted]
966        } else {
967            buf
968        };
969
970        let block_capacity = self.block_capacity();
971        if self.pending.capacity() == 0 {
972            self.pending = self.allocate_pending_space(block_capacity);
973        }
974        let mut remaining = buf;
975        let mut consumed = 0usize;
976
977        while !remaining.is_empty() {
978            if let Some(result) = self.emit_full_pending_block(block_capacity, consumed) {
979                return result;
980            }
981
982            let available = block_capacity - self.pending.len();
983            let to_take = core::cmp::min(remaining.len(), available);
984            if to_take == 0 {
985                break;
986            }
987            self.pending.extend_from_slice(&remaining[..to_take]);
988            remaining = &remaining[to_take..];
989            consumed += to_take;
990
991            if let Some(result) = self.emit_full_pending_block(block_capacity, consumed) {
992                if let Ok(n) = &result {
993                    self.bytes_consumed += *n as u64;
994                }
995                return result;
996            }
997        }
998        self.bytes_consumed += consumed as u64;
999        Ok(consumed)
1000    }
1001
1002    fn flush(&mut self) -> Result<(), Error> {
1003        self.ensure_open()?;
1004        if self.pending.is_empty() {
1005            return self
1006                .drain_mut()
1007                .and_then(|drain| drain.flush())
1008                .map_err(|err| self.fail(err));
1009        }
1010        self.ensure_frame_started()?;
1011        self.emit_pending_block(false)?;
1012        self.drain_mut()
1013            .and_then(|drain| drain.flush())
1014            .map_err(|err| self.fail(err))
1015    }
1016}
1017
1018fn error_from_kind(kind: ErrorKind) -> Error {
1019    Error::from(kind)
1020}
1021
1022fn error_with_kind_message(kind: ErrorKind, message: String) -> Error {
1023    #[cfg(feature = "std")]
1024    {
1025        Error::new(kind, message)
1026    }
1027    #[cfg(not(feature = "std"))]
1028    {
1029        Error::new(kind, alloc::boxed::Box::new(message))
1030    }
1031}
1032
1033fn invalid_input_error(message: &str) -> Error {
1034    #[cfg(feature = "std")]
1035    {
1036        Error::new(ErrorKind::InvalidInput, message)
1037    }
1038    #[cfg(not(feature = "std"))]
1039    {
1040        Error::new(
1041            ErrorKind::Other,
1042            alloc::boxed::Box::new(alloc::string::String::from(message)),
1043        )
1044    }
1045}
1046
1047fn other_error_owned(message: String) -> Error {
1048    #[cfg(feature = "std")]
1049    {
1050        Error::other(message)
1051    }
1052    #[cfg(not(feature = "std"))]
1053    {
1054        Error::new(ErrorKind::Other, alloc::boxed::Box::new(message))
1055    }
1056}
1057
1058fn other_error(message: &str) -> Error {
1059    #[cfg(feature = "std")]
1060    {
1061        Error::other(message)
1062    }
1063    #[cfg(not(feature = "std"))]
1064    {
1065        Error::new(
1066            ErrorKind::Other,
1067            alloc::boxed::Box::new(alloc::string::String::from(message)),
1068        )
1069    }
1070}
1071
1072#[cfg(test)]
1073mod tests;