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