structured_zstd/decoding/

streaming_decoder.rs

//! The [StreamingDecoder] wraps a [FrameDecoder] and provides a Read impl that decodes data when necessary

use core::borrow::BorrowMut;

use crate::common::MAX_BLOCK_SIZE;
use crate::decoding::errors::FrameDecoderError;
use crate::decoding::{BlockDecodingStrategy, DictionaryHandle, FrameDecoder};
#[cfg(not(feature = "std"))]
use crate::io::ErrorKind;
use crate::io::{Error, Read};

/// High level Zstandard frame decoder that can be used to decompress a given Zstandard frame.
///
/// This decoder implements `io::Read`, so you can interact with it by calling
/// `io::Read::read_to_end` / `io::Read::read_exact` or passing this to another library / module as a source for the decoded content
///
/// If you need more control over how decompression takes place, you can use
/// the lower level [FrameDecoder], which allows for greater control over how
/// decompression takes place but the implementor must call
/// [FrameDecoder::decode_blocks] repeatedly to decode the entire frame.
///
/// ## Caveat
/// Plain `read` / `read_exact` operate on the single frame this decoder was
/// initialised with: they do not advance into following frames. `read_to_end`,
/// by contrast, is specialised to consume a finite source to EOF, decoding
/// concatenated frames and skipping skippable frames along the way.
///
/// To recover the bytes that follow one frame WITHOUT consuming the rest of the
/// source, recreate the decoder manually and handle
/// [crate::decoding::errors::ReadFrameHeaderError::SkipFrame]
/// errors by skipping forward the `length` amount of bytes, see <https://github.com/KillingSpark/zstd-rs/issues/57>
///
/// ```no_run
/// // `File` is std-only; `read_to_end` itself is available under no_std too.
/// #[cfg(feature = "std")]
/// {
///     use std::fs::File;
///     use std::io::Read;
///     use structured_zstd::decoding::StreamingDecoder;
///
///     // Read a Zstandard archive from the filesystem then decompress it into a vec.
///     let mut f: File = todo!("Read a .zstd archive from somewhere");
///     let mut decoder = StreamingDecoder::new(f).unwrap();
///     let mut result = Vec::new();
///     Read::read_to_end(&mut decoder, &mut result).unwrap();
/// }
/// ```
pub struct StreamingDecoder<READ: Read, DEC: BorrowMut<FrameDecoder>> {
    pub decoder: DEC,
    source: READ,
    /// Whether the decoder was constructed with a dictionary it applies to
    /// every frame. The `read_to_end` paths re-initialise FOLLOWING
    /// concatenated frames with it (a plain re-init resolves dictionaries by
    /// frame id only and would lose it for frames omitting the id); the
    /// decoder already holds its handle, so this one keeps none of its own.
    forced_dictionary: bool,
}

impl<READ: Read, DEC: BorrowMut<FrameDecoder>> StreamingDecoder<READ, DEC> {
    pub fn new_with_decoder(
        mut source: READ,
        mut decoder: DEC,
    ) -> Result<StreamingDecoder<READ, DEC>, FrameDecoderError> {
        decoder.borrow_mut().init(&mut source)?;
        Ok(StreamingDecoder {
            decoder,
            source,
            forced_dictionary: false,
        })
    }

    /// [`new_with_decoder`](Self::new_with_decoder) with `dict` applied to
    /// the frame even when its header omits the dictionary ID, as
    /// [`new_with_dictionary_handle`](StreamingDecoder::new_with_dictionary_handle)
    /// applies it (same warning). A decoder reused frame after frame keeps its
    /// buffers and, for the same dictionary, the handle it already holds, so
    /// after the first frame it allocates nothing and touches no reference
    /// count.
    ///
    /// # Examples
    /// ```
    /// use std::io::Read;
    /// use structured_zstd::decoding::{Dictionary, DictionaryHandle, FrameDecoder, StreamingDecoder};
    /// use structured_zstd::encoding::{FrameCompressor, CompressionLevel};
    ///
    /// let content = b"a dictionary of words the frames reuse".to_vec();
    /// let dictionary = DictionaryHandle::from_dictionary(
    ///     Dictionary::from_raw_content(1, content.clone()).unwrap(),
    /// );
    /// let mut compressor: FrameCompressor = FrameCompressor::new(CompressionLevel::Default);
    /// compressor.set_dictionary(Dictionary::from_raw_content(1, content).unwrap()).unwrap();
    /// let frame = compressor.compress_independent_frame(b"the frames reuse words");
    ///
    /// let mut decoder = FrameDecoder::new();
    /// for _ in 0..2 {
    ///     let mut stream =
    ///         StreamingDecoder::new_with_decoder_and_dictionary_handle(&frame[..], &mut decoder, &dictionary)
    ///             .unwrap();
    ///     let mut decoded = Vec::new();
    ///     stream.read_to_end(&mut decoded).unwrap();
    ///     assert_eq!(decoded, b"the frames reuse words");
    /// }
    /// ```
    pub fn new_with_decoder_and_dictionary_handle(
        mut source: READ,
        mut decoder: DEC,
        dict: &DictionaryHandle,
    ) -> Result<StreamingDecoder<READ, DEC>, FrameDecoderError> {
        decoder
            .borrow_mut()
            .init_with_dict_handle(&mut source, dict)?;
        Ok(StreamingDecoder {
            decoder,
            source,
            forced_dictionary: true,
        })
    }
}

impl<READ: Read> StreamingDecoder<READ, FrameDecoder> {
    pub fn new(
        mut source: READ,
    ) -> Result<StreamingDecoder<READ, FrameDecoder>, FrameDecoderError> {
        let mut decoder = FrameDecoder::new();
        decoder.init(&mut source)?;
        Ok(StreamingDecoder {
            decoder,
            source,
            forced_dictionary: false,
        })
    }

    /// Create a streaming decoder using a pre-parsed dictionary handle.
    ///
    /// # Warning
    ///
    /// This constructor initializes the underlying [`FrameDecoder`] with
    /// `dict`, even if a frame header omits the optional dictionary ID.
    /// Callers must only use it when they already know the stream was encoded
    /// with this dictionary; otherwise decoded output can be silently
    /// corrupted.
    pub fn new_with_dictionary_handle(
        source: READ,
        dict: &DictionaryHandle,
    ) -> Result<StreamingDecoder<READ, FrameDecoder>, FrameDecoderError> {
        Self::new_with_decoder_and_dictionary_handle(source, FrameDecoder::new(), dict)
    }

    /// Create a streaming decoder using a serialized dictionary blob.
    ///
    /// # Warning
    ///
    /// This API forwards to [`StreamingDecoder::new_with_dictionary_handle`]
    /// and therefore applies the decoded dictionary to frames whose headers may
    /// omit the optional dictionary ID. Only use it when the stream is known to
    /// be encoded with that dictionary.
    pub fn new_with_dictionary_bytes(
        source: READ,
        raw_dictionary: &[u8],
    ) -> Result<StreamingDecoder<READ, FrameDecoder>, FrameDecoderError> {
        let dict = DictionaryHandle::decode_dict(raw_dictionary)?;
        Self::new_with_dictionary_handle(source, &dict)
    }
}

impl<READ: Read, DEC: BorrowMut<FrameDecoder>> StreamingDecoder<READ, DEC> {
    /// Gets a reference to the underlying reader.
    pub fn get_ref(&self) -> &READ {
        &self.source
    }

    /// Gets a mutable reference to the underlying reader.
    ///
    /// It is inadvisable to directly read from the underlying reader.
    pub fn get_mut(&mut self) -> &mut READ {
        &mut self.source
    }

    /// Gets a mutable reference to the frame decoder driving this stream.
    ///
    /// Exposed for settings that are read as decoding proceeds rather than at
    /// construction — [`FrameDecoder::set_content_checksum`] above all, which a
    /// caller that wants mismatches to fail (rather than merely be computed)
    /// has to reach after the constructor has chosen and initialised the
    /// decoder, including on the dictionary paths.
    pub fn decoder_mut(&mut self) -> &mut FrameDecoder {
        self.decoder.borrow_mut()
    }

    /// Destructures this object into the inner reader.
    pub fn into_inner(self) -> READ
    where
        READ: Sized,
    {
        self.source
    }

    /// Destructures this object into both the inner reader and [FrameDecoder].
    pub fn into_parts(self) -> (READ, DEC)
    where
        READ: Sized,
    {
        (self.source, self.decoder)
    }

    /// Destructures this object into the inner [FrameDecoder].
    pub fn into_frame_decoder(self) -> DEC {
        self.decoder
    }
}

impl<READ: Read, DEC: BorrowMut<FrameDecoder>> Read for StreamingDecoder<READ, DEC> {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize, Error> {
        let decoder = self.decoder.borrow_mut();
        if decoder.is_finished() && decoder.can_collect() == 0 {
            // Frame fully decoded and fully drained: the running XXH64 digest
            // is final, so a `Verify`-mode decoder validates the content
            // checksum at this finish point. No-op in other modes.
            #[cfg(feature = "hash")]
            if let Err(e) = decoder.verify_content_checksum() {
                #[cfg(feature = "std")]
                return Err(Error::other(e));
                #[cfg(not(feature = "std"))]
                return Err(Error::new(ErrorKind::Other, alloc::boxed::Box::new(e)));
            }
            //No more bytes can ever be decoded
            return Ok(0);
        }

        // Interleave bounded decode with draining so the decode window
        // (`RingBuffer`) stays near `window_size` instead of accumulating the
        // whole request before a single end-of-call drain. `read_to_end` hands
        // ever-larger buffers; decoding `buf.len()` worth into the ring up
        // front grew it far past the window (repeated `reserve_amortized`
        // alloc+copy). Decode at most one block worth per step, then drain
        // what is now collectable into `buf`, mirroring upstream zstd's
        // window-bounded flush loop.
        let mut written = 0;
        while written < buf.len() {
            // Drain whatever is collectable now (retaining `window_size` until
            // the frame finishes). Reclaims the ring promptly so the next
            // decode step reuses the same capacity.
            written += decoder.read(&mut buf[written..])?;
            if written == buf.len() || decoder.is_finished() {
                break;
            }
            // Decode one bounded chunk. `UptoBytes` may overshoot a little but
            // is capped to one block, so the ring's live region stays within
            // `window_size + MAX_BLOCK_SIZE`.
            let step = (buf.len() - written).min(MAX_BLOCK_SIZE as usize);
            if let Err(e) =
                decoder.decode_blocks(&mut self.source, BlockDecodingStrategy::UptoBytes(step))
            {
                #[cfg(feature = "std")]
                {
                    return Err(Error::other(e));
                }
                #[cfg(not(feature = "std"))]
                {
                    return Err(Error::new(ErrorKind::Other, alloc::boxed::Box::new(e)));
                }
            }
        }

        // The loop can finish AND fully drain a frame within this same call
        // (decode last block, then drain it into `buf`). Validate here too when
        // the frame is finished and nothing is left to collect, but ONLY when
        // this call wrote no bytes: the `Read` contract forbids returning `Err`
        // after bytes were delivered, so when `written > 0` the verify is
        // deferred to the next call, where the top early-return runs it and
        // returns `Err` on the zero-byte path. Idempotent with that top check.
        #[cfg(feature = "hash")]
        if written == 0
            && decoder.is_finished()
            && decoder.can_collect() == 0
            && let Err(e) = decoder.verify_content_checksum()
        {
            #[cfg(feature = "std")]
            return Err(Error::other(e));
            #[cfg(not(feature = "std"))]
            return Err(Error::new(ErrorKind::Other, alloc::boxed::Box::new(e)));
        }

        Ok(written)
    }

    /// Decode-in-place fast path for whole-frame consumption. Instead of the
    /// generic `read` loop (decode block -> `RingBuffer` -> copy into the
    /// caller buffer), buffer the (compressed, hence small) source and decode
    /// STRAIGHT into `output`'s spare capacity via the single-copy direct path,
    /// pre-sized from the frame's declared content size. Only taken when the
    /// decoder is at a frame boundary (nothing partially decoded / undrained);
    /// otherwise it falls back to the generic grow-and-`read` loop so a caller
    /// that mixed `read` with `read_to_end` still gets correct output.
    ///
    /// Per the `Read::read_to_end` contract this consumes the source to EOF: if
    /// the stream holds several concatenated frames they are ALL decoded (and
    /// skippable frames skipped). To recover bytes that follow a single frame,
    /// use `read` plus the
    /// [`SkipFrame`](crate::decoding::errors::ReadFrameHeaderError::SkipFrame)
    /// recreate-the-decoder pattern instead.
    #[cfg(feature = "std")]
    fn read_to_end(&mut self, output: &mut alloc::vec::Vec<u8>) -> Result<usize, Error> {
        let start_total = output.len();
        // `new()` already read the frame header, so the fast path applies when
        // the decoder sits at the start of that frame with nothing decoded yet.
        let at_start = {
            let d = self.decoder.borrow_mut();
            d.is_at_frame_start() && d.can_collect() == 0
        };
        // A forced dictionary is the one the decoder already holds; following
        // frames are re-initialised with it in place, touching no reference
        // count.
        let keep_dictionary = self.forced_dictionary;
        if at_start {
            let mut compressed = alloc::vec::Vec::new();
            self.source.read_to_end(&mut compressed)?;
            self.decoder
                .borrow_mut()
                .decode_current_frame_to_vec(&compressed, output, keep_dictionary)
                .map_err(Error::other)?;
            return Ok(output.len() - start_total);
        }
        // Mid-frame fallback: drain the partially-read CURRENT frame through the
        // generic path, then decode any FOLLOWING concatenated frames so
        // read_to_end still consumes the source to true EOF.
        loop {
            let start = output.len();
            output.resize(start + MAX_BLOCK_SIZE as usize, 0);
            // On error, drop the just-grown (zeroed) tail before propagating so
            // the caller never observes bytes that were never decoded.
            let n = match self.read(&mut output[start..]) {
                Ok(n) => n,
                Err(e) => {
                    output.truncate(start);
                    return Err(e);
                }
            };
            output.truncate(start + n);
            if n == 0 {
                break;
            }
        }
        // Current frame fully drained; `source` is positioned at the next frame.
        let mut rest = alloc::vec::Vec::new();
        self.source.read_to_end(&mut rest)?;
        if !rest.is_empty() {
            let mut input = rest.as_slice();
            self.decoder
                .borrow_mut()
                .decode_concatenated_frames_to_vec(&mut input, output, keep_dictionary)
                .map_err(Error::other)?;
        }
        Ok(output.len() - start_total)
    }

    /// no_std counterpart of the decode-in-place `read_to_end` fast path above
    /// (the no_std `Read::read_to_end` returns `()` instead of the byte count).
    #[cfg(not(feature = "std"))]
    fn read_to_end(&mut self, output: &mut alloc::vec::Vec<u8>) -> Result<(), Error> {
        let at_start = {
            let d = self.decoder.borrow_mut();
            d.is_at_frame_start() && d.can_collect() == 0
        };
        // As in the std path: the decoder's own dictionary, reused in place.
        let keep_dictionary = self.forced_dictionary;
        if at_start {
            let mut compressed = alloc::vec::Vec::new();
            self.source.read_to_end(&mut compressed)?;
            self.decoder
                .borrow_mut()
                .decode_current_frame_to_vec(&compressed, output, keep_dictionary)
                .map_err(|e| Error::new(ErrorKind::Other, alloc::boxed::Box::new(e)))?;
            return Ok(());
        }
        // Mid-frame fallback: drain the partial CURRENT frame, then decode the
        // FOLLOWING concatenated frames so the source is consumed to true EOF.
        loop {
            let start = output.len();
            output.resize(start + MAX_BLOCK_SIZE as usize, 0);
            // On error, drop the just-grown (zeroed) tail before propagating so
            // the caller never observes bytes that were never decoded.
            let n = match self.read(&mut output[start..]) {
                Ok(n) => n,
                Err(e) => {
                    output.truncate(start);
                    return Err(e);
                }
            };
            output.truncate(start + n);
            if n == 0 {
                break;
            }
        }
        let mut rest = alloc::vec::Vec::new();
        self.source.read_to_end(&mut rest)?;
        if !rest.is_empty() {
            let mut input = rest.as_slice();
            self.decoder
                .borrow_mut()
                .decode_concatenated_frames_to_vec(&mut input, output, keep_dictionary)
                .map_err(|e| Error::new(ErrorKind::Other, alloc::boxed::Box::new(e)))?;
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests;