zstd_complete/encoding/
streaming_encoder.rs

1use alloc::vec::Vec;
2use core::fmt;
3#[cfg(feature = "hash")]
4use core::hash::Hasher;
5use std::io::{self, Read, Write};
6
7use super::{levels::c_port, CompressionLevel};
8use crate::{
9    blocks::block::BlockType,
10    common::MAX_BLOCK_SIZE,
11    encoding::{block_header::BlockHeader, frame_header::FrameHeader},
12};
13
14/// Default maximum estimated resident working set: 96 MiB.
15pub const DEFAULT_MEMORY_LIMIT: usize = 96 * 1024 * 1024;
16/// Default amount of input stored in each independent Zstandard frame: 8 MiB.
17pub const DEFAULT_FRAME_CHUNK_SIZE: usize = 8 * 1024 * 1024;
18
19/// Configuration for the bounded-memory streaming encoder.
20#[derive(Clone)]
21pub struct EncoderOptions {
22    level: CompressionLevel,
23    memory_limit: usize,
24    frame_chunk_size: usize,
25    dictionary: Option<EncoderDictionary>,
26    checksum: bool,
27}
28
29impl EncoderOptions {
30    /// Creates options using the 96 MiB memory budget and 8 MiB frame chunks.
31    pub const fn new(level: CompressionLevel) -> Self {
32        Self {
33            level,
34            memory_limit: DEFAULT_MEMORY_LIMIT,
35            frame_chunk_size: DEFAULT_FRAME_CHUNK_SIZE,
36            dictionary: None,
37            checksum: false,
38        }
39    }
40
41    pub const fn level(&self) -> CompressionLevel {
42        self.level
43    }
44
45    pub const fn memory_limit(&self) -> usize {
46        self.memory_limit
47    }
48
49    pub const fn frame_chunk_size(&self) -> usize {
50        self.frame_chunk_size
51    }
52
53    pub fn dictionary(&self) -> Option<&EncoderDictionary> {
54        self.dictionary.as_ref()
55    }
56
57    pub const fn checksum(&self) -> bool {
58        self.checksum
59    }
60
61    /// Returns the conservative memory estimate checked by [`Encoder::new`].
62    pub fn estimated_memory_usage(&self) -> usize {
63        let required = c_port::estimated_frame_memory(self.level.c_level(), self.frame_chunk_size);
64        self.dictionary.as_ref().map_or(required, |dictionary| {
65            let retained_copies = if cfg!(feature = "multithreading") {
66                3
67            } else {
68                2
69            };
70            required.saturating_add(dictionary.raw_size().saturating_mul(retained_copies))
71        })
72    }
73
74    /// Sets the maximum conservative working-memory estimate.
75    pub const fn with_memory_limit(mut self, bytes: usize) -> Self {
76        self.memory_limit = bytes;
77        self
78    }
79
80    /// Sets the input bytes collected per independent Zstandard frame.
81    ///
82    /// Smaller chunks reduce latency and memory use but can reduce compression
83    /// ratio because matches do not cross frame boundaries.
84    pub const fn with_frame_chunk_size(mut self, bytes: usize) -> Self {
85        self.frame_chunk_size = bytes;
86        self
87    }
88
89    /// Uses a prepared dictionary for every frame in the output archive.
90    pub fn with_dictionary(mut self, dictionary: EncoderDictionary) -> Self {
91        self.dictionary = Some(dictionary);
92        self
93    }
94
95    /// Enables or disables the standard 32-bit frame content checksum.
96    pub const fn with_checksum(mut self, enabled: bool) -> Self {
97        self.checksum = enabled;
98        self
99    }
100
101    pub(super) fn validate(&self) -> Result<(), EncodeError> {
102        if self.frame_chunk_size == 0 {
103            return Err(EncodeError::InvalidOptions(
104                "frame chunk size must be greater than zero",
105            ));
106        }
107        if self.level.is_uncompressed() && self.dictionary.is_some() {
108            return Err(EncodeError::InvalidOptions(
109                "an uncompressed encoder cannot use a dictionary",
110            ));
111        }
112        if self.checksum && !cfg!(feature = "hash") {
113            return Err(EncodeError::InvalidOptions(
114                "checksums require the zstd-complete `hash` feature",
115            ));
116        }
117        let required = self.estimated_memory_usage();
118        if required > self.memory_limit {
119            return Err(EncodeError::MemoryLimitExceeded {
120                limit: self.memory_limit,
121                required,
122            });
123        }
124        Ok(())
125    }
126}
127
128impl fmt::Debug for EncoderOptions {
129    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
130        formatter
131            .debug_struct("EncoderOptions")
132            .field("level", &self.level)
133            .field("memory_limit", &self.memory_limit)
134            .field("frame_chunk_size", &self.frame_chunk_size)
135            .field("dictionary", &self.dictionary)
136            .field("checksum", &self.checksum)
137            .finish()
138    }
139}
140
141/// An owned, parsed dictionary reusable across independent frames.
142#[derive(Clone)]
143pub struct EncoderDictionary {
144    prepared: c_port::PreparedDictionary,
145    raw_size: usize,
146    #[cfg(feature = "multithreading")]
147    raw: Vec<u8>,
148}
149
150impl EncoderDictionary {
151    /// Copies and validates raw-content or formatted Zstandard dictionary data.
152    pub fn copy(dictionary: &[u8]) -> Result<Self, DictionaryError> {
153        let prepared = c_port::PreparedDictionary::from_bytes(dictionary)
154            .map_err(|error| match error {
155                c_port::DictionaryParseError::WrongDictionary => DictionaryError::WrongDictionary,
156                c_port::DictionaryParseError::CorruptedDictionary => {
157                    DictionaryError::CorruptedDictionary
158                }
159            })?
160            .ok_or(DictionaryError::WrongDictionary)?;
161        Ok(Self {
162            prepared,
163            raw_size: dictionary.len(),
164            #[cfg(feature = "multithreading")]
165            raw: dictionary.to_vec(),
166        })
167    }
168
169    pub const fn raw_size(&self) -> usize {
170        self.raw_size
171    }
172
173    #[cfg(feature = "multithreading")]
174    pub(super) fn raw(&self) -> &[u8] {
175        &self.raw
176    }
177}
178
179impl fmt::Debug for EncoderDictionary {
180    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
181        formatter
182            .debug_struct("EncoderDictionary")
183            .field("raw_size", &self.raw_size)
184            .finish_non_exhaustive()
185    }
186}
187
188/// Error returned when dictionary bytes are not a usable Zstandard dictionary.
189#[derive(Clone, Copy, Debug, Eq, PartialEq)]
190pub enum DictionaryError {
191    WrongDictionary,
192    CorruptedDictionary,
193}
194
195impl fmt::Display for DictionaryError {
196    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
197        formatter.write_str(match self {
198            Self::WrongDictionary => "input is not a usable Zstandard dictionary",
199            Self::CorruptedDictionary => "Zstandard dictionary metadata is corrupted",
200        })
201    }
202}
203
204impl std::error::Error for DictionaryError {}
205
206impl Default for EncoderOptions {
207    fn default() -> Self {
208        Self::new(CompressionLevel::DEFAULT)
209    }
210}
211
212/// Errors produced by the bounded streaming API.
213#[derive(Debug)]
214pub enum EncodeError {
215    Io(io::Error),
216    InvalidOptions(&'static str),
217    MemoryLimitExceeded {
218        limit: usize,
219        required: usize,
220    },
221    #[cfg(feature = "multithreading")]
222    WorkerFailed,
223}
224
225impl fmt::Display for EncodeError {
226    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
227        match self {
228            Self::Io(error) => error.fmt(formatter),
229            Self::InvalidOptions(message) => formatter.write_str(message),
230            Self::MemoryLimitExceeded { limit, required } => write!(
231                formatter,
232                "encoder needs an estimated {required} bytes, exceeding the {limit}-byte memory limit"
233            ),
234            #[cfg(feature = "multithreading")]
235            Self::WorkerFailed => formatter.write_str("a compression worker failed"),
236        }
237    }
238}
239
240impl std::error::Error for EncodeError {
241    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
242        match self {
243            Self::Io(error) => Some(error),
244            _ => None,
245        }
246    }
247}
248
249impl From<io::Error> for EncodeError {
250    fn from(error: io::Error) -> Self {
251        Self::Io(error)
252    }
253}
254
255/// A bounded-memory Zstandard encoder implementing [`Write`].
256///
257/// Each full input chunk becomes an independent frame in the output archive.
258/// Call [`Encoder::finish`] to emit the final frame and recover the writer.
259pub struct Encoder<W: Write> {
260    inner: W,
261    options: EncoderOptions,
262    input: Vec<u8>,
263    emitted_frame: bool,
264}
265
266impl<W: Write> Encoder<W> {
267    pub fn new(inner: W, options: EncoderOptions) -> Result<Self, EncodeError> {
268        options.validate()?;
269        let input = Vec::with_capacity(options.frame_chunk_size);
270        Ok(Self {
271            inner,
272            options,
273            input,
274            emitted_frame: false,
275        })
276    }
277
278    pub fn get_ref(&self) -> &W {
279        &self.inner
280    }
281
282    pub fn get_mut(&mut self) -> &mut W {
283        &mut self.inner
284    }
285
286    /// Emits the final frame and returns the wrapped writer.
287    pub fn finish(mut self) -> Result<W, EncodeError> {
288        if !self.input.is_empty() || !self.emitted_frame {
289            self.emit_frame()?;
290        }
291        self.inner.flush()?;
292        Ok(self.inner)
293    }
294
295    fn emit_frame(&mut self) -> io::Result<()> {
296        let mut compressed = if self.options.level.is_uncompressed() {
297            encode_uncompressed_frame(&self.input)
298        } else if let Some(dictionary) = self.options.dictionary.as_ref() {
299            c_port::encode_frame_with_prepared_dictionary(
300                &self.input,
301                self.options.level.c_level(),
302                &dictionary.prepared,
303            )
304        } else {
305            c_port::encode_frame_no_dict(&self.input, self.options.level.c_level())
306        };
307        if self.options.checksum {
308            append_checksum(&mut compressed, &self.input);
309        }
310        self.inner.write_all(&compressed)?;
311        self.input.clear();
312        self.emitted_frame = true;
313        Ok(())
314    }
315}
316
317fn append_checksum(frame: &mut Vec<u8>, source: &[u8]) {
318    #[cfg(feature = "hash")]
319    {
320        debug_assert!(frame.len() >= 5);
321        frame[4] |= 1 << 2;
322        let mut hasher = twox_hash::XxHash64::with_seed(0);
323        hasher.write(source);
324        frame.extend_from_slice(&(hasher.finish() as u32).to_le_bytes());
325    }
326    #[cfg(not(feature = "hash"))]
327    {
328        let _ = (frame, source);
329        unreachable!("checksum options are rejected without the hash feature");
330    }
331}
332
333#[cfg(feature = "multithreading")]
334pub(super) fn encode_frame_for_options(source: &[u8], options: &EncoderOptions) -> Vec<u8> {
335    let mut compressed = if options.level.is_uncompressed() {
336        encode_uncompressed_frame(source)
337    } else if let Some(dictionary) = options.dictionary.as_ref() {
338        c_port::encode_frame_with_prepared_dictionary(
339            source,
340            options.level.c_level(),
341            &dictionary.prepared,
342        )
343    } else {
344        c_port::encode_frame_no_dict(source, options.level.c_level())
345    };
346    if options.checksum {
347        append_checksum(&mut compressed, source);
348    }
349    compressed
350}
351
352fn encode_uncompressed_frame(source: &[u8]) -> Vec<u8> {
353    let mut output = Vec::with_capacity(source.len().saturating_add(16));
354    FrameHeader {
355        frame_content_size: Some(source.len() as u64),
356        single_segment: true,
357        content_checksum: false,
358        dictionary_id: None,
359        window_size: None,
360    }
361    .serialize(&mut output);
362
363    if source.is_empty() {
364        BlockHeader {
365            last_block: true,
366            block_type: BlockType::Raw,
367            block_size: 0,
368        }
369        .serialize(&mut output);
370        return output;
371    }
372
373    for (index, block) in source.chunks(MAX_BLOCK_SIZE as usize).enumerate() {
374        BlockHeader {
375            last_block: (index + 1) * MAX_BLOCK_SIZE as usize >= source.len(),
376            block_type: BlockType::Raw,
377            block_size: block.len() as u32,
378        }
379        .serialize(&mut output);
380        output.extend_from_slice(block);
381    }
382    output
383}
384
385impl<W: Write> Write for Encoder<W> {
386    fn write(&mut self, source: &[u8]) -> io::Result<usize> {
387        if source.is_empty() {
388            return Ok(0);
389        }
390        if self.input.len() == self.options.frame_chunk_size {
391            self.emit_frame()?;
392        }
393        let available = self.options.frame_chunk_size - self.input.len();
394        let consumed = available.min(source.len());
395        self.input.extend_from_slice(&source[..consumed]);
396        Ok(consumed)
397    }
398
399    /// Closes the current frame before flushing the wrapped writer.
400    fn flush(&mut self) -> io::Result<()> {
401        if !self.input.is_empty() {
402            self.emit_frame()?;
403        }
404        self.inner.flush()
405    }
406}
407
408/// Streams all input into a bounded-memory encoder.
409pub fn encode<R: Read, W: Write>(
410    mut source: R,
411    target: W,
412    options: EncoderOptions,
413) -> Result<(), EncodeError> {
414    let mut encoder = Encoder::new(target, options)?;
415    io::copy(&mut source, &mut encoder)?;
416    encoder.finish()?;
417    Ok(())
418}
419
420/// Compresses all input and returns the resulting Zstandard archive.
421pub fn encode_all<R: Read>(source: R, options: EncoderOptions) -> Result<Vec<u8>, EncodeError> {
422    let mut output = Vec::new();
423    encode(source, &mut output, options)?;
424    Ok(output)
425}
426
427#[cfg(test)]
428mod tests {
429    use super::*;
430    use crate::decoding::FrameDecoder;
431    use alloc::vec;
432
433    #[test]
434    fn multi_frame_stream_roundtrips() {
435        let input = b"abcdefgh".repeat(4096);
436        let options = EncoderOptions::new(CompressionLevel::FASTEST)
437            .with_frame_chunk_size(1024)
438            .with_memory_limit(16 * 1024 * 1024);
439        let compressed = encode_all(input.as_slice(), options).unwrap();
440        let mut decoded = Vec::with_capacity(input.len());
441        FrameDecoder::new()
442            .decode_all_to_vec(&compressed, &mut decoded)
443            .unwrap();
444        assert_eq!(decoded, input);
445    }
446
447    #[test]
448    fn empty_input_is_a_valid_frame() {
449        let compressed = encode_all(&[][..], EncoderOptions::default()).unwrap();
450        let mut decoded = Vec::with_capacity(1);
451        FrameDecoder::new()
452            .decode_all_to_vec(&compressed, &mut decoded)
453            .unwrap();
454        assert!(decoded.is_empty());
455    }
456
457    #[test]
458    fn uncompressed_stream_roundtrips() {
459        let input = vec![42; 300_000];
460        let options =
461            EncoderOptions::new(CompressionLevel::UNCOMPRESSED).with_frame_chunk_size(150_000);
462        let compressed = encode_all(input.as_slice(), options).unwrap();
463        assert_eq!(zstd::decode_all(compressed.as_slice()).unwrap(), input);
464    }
465
466    #[test]
467    fn archive_is_interoperable_with_the_c_decoder() {
468        let input = b"interoperable streaming input".repeat(10_000);
469        let options = EncoderOptions::new(CompressionLevel::DEFAULT)
470            .with_frame_chunk_size(32 * 1024)
471            .with_memory_limit(32 * 1024 * 1024);
472        let compressed = encode_all(input.as_slice(), options).unwrap();
473        assert_eq!(zstd::decode_all(compressed.as_slice()).unwrap(), input);
474    }
475
476    #[cfg(feature = "hash")]
477    #[test]
478    fn checksummed_frames_roundtrip() {
479        let input = b"checksummed content".repeat(10_000);
480        let options = EncoderOptions::new(CompressionLevel::DEFAULT)
481            .with_frame_chunk_size(32 * 1024)
482            .with_memory_limit(32 * 1024 * 1024)
483            .with_checksum(true);
484        let compressed = encode_all(input.as_slice(), options).unwrap();
485        assert_eq!(zstd::decode_all(compressed.as_slice()).unwrap(), input);
486    }
487
488    #[test]
489    fn prepared_dictionary_stream_roundtrips() {
490        let dictionary_bytes = b"dictionary words and repeated prefixes for the test".repeat(8);
491        let input = dictionary_bytes.repeat(200);
492        let dictionary = EncoderDictionary::copy(&dictionary_bytes).unwrap();
493        let options = EncoderOptions::new(CompressionLevel::DEFAULT)
494            .with_frame_chunk_size(16 * 1024)
495            .with_memory_limit(32 * 1024 * 1024)
496            .with_dictionary(dictionary);
497        let compressed = encode_all(input.as_slice(), options).unwrap();
498
499        let mut decoded = Vec::with_capacity(input.len());
500        let mut decoder =
501            zstd::stream::read::Decoder::with_dictionary(compressed.as_slice(), &dictionary_bytes)
502                .unwrap();
503        decoder.read_to_end(&mut decoded).unwrap();
504        assert_eq!(decoded, input);
505    }
506
507    #[test]
508    fn too_short_dictionary_is_rejected_instead_of_silently_ignored() {
509        assert_eq!(
510            EncoderDictionary::copy(b"short").unwrap_err(),
511            DictionaryError::WrongDictionary
512        );
513    }
514
515    #[test]
516    fn invalid_budget_is_reported_before_writing() {
517        let options = EncoderOptions::new(CompressionLevel::MAXIMUM);
518        assert!(matches!(
519            Encoder::new(io::sink(), options),
520            Err(EncodeError::MemoryLimitExceeded { .. })
521        ));
522    }
523
524    #[test]
525    fn target_io_errors_are_preserved() {
526        struct FailingWriter;
527
528        impl Write for FailingWriter {
529            fn write(&mut self, _source: &[u8]) -> io::Result<usize> {
530                Err(io::Error::new(io::ErrorKind::BrokenPipe, "test failure"))
531            }
532
533            fn flush(&mut self) -> io::Result<()> {
534                Ok(())
535            }
536        }
537
538        let error = encode(
539            b"accepted input".as_slice(),
540            FailingWriter,
541            EncoderOptions::default(),
542        )
543        .unwrap_err();
544        assert!(matches!(
545            error,
546            EncodeError::Io(ref error) if error.kind() == io::ErrorKind::BrokenPipe
547        ));
548    }
549
550    #[test]
551    fn input_buffer_never_exceeds_the_configured_chunk() {
552        let options = EncoderOptions::new(CompressionLevel::FASTEST)
553            .with_frame_chunk_size(1024)
554            .with_memory_limit(16 * 1024 * 1024);
555        let mut encoder = Encoder::new(io::sink(), options).unwrap();
556        for _ in 0..1024 {
557            encoder.write_all(&[7; 4096]).unwrap();
558            assert!(encoder.input.len() <= 1024);
559            assert_eq!(encoder.input.capacity(), 1024);
560        }
561        encoder.finish().unwrap();
562    }
563
564    #[test]
565    fn large_logical_source_does_not_require_a_large_input_allocation() {
566        struct RepeatingReader {
567            remaining: usize,
568        }
569
570        impl Read for RepeatingReader {
571            fn read(&mut self, target: &mut [u8]) -> io::Result<usize> {
572                let count = self.remaining.min(target.len());
573                target[..count].fill(19);
574                self.remaining -= count;
575                Ok(count)
576            }
577        }
578
579        let options = EncoderOptions::new(CompressionLevel::UNCOMPRESSED)
580            .with_frame_chunk_size(1024 * 1024)
581            .with_memory_limit(4 * 1024 * 1024);
582        encode(
583            RepeatingReader {
584                remaining: 64 * 1024 * 1024,
585            },
586            io::sink(),
587            options,
588        )
589        .unwrap();
590    }
591}
zstd_complete/encoding/streaming_encoder.rs

zstd_complete/encoding/
streaming_encoder.rs
