mzstd 0.2.1

use super::{DecodeError, Decompressor};

/// Whether more input may follow a call.
///
/// Use [`Self::Process`] while more input may arrive. Once the input is
/// known to end, use [`Self::Finish`] on every remaining call, each time
/// passing exactly the input not yet consumed. See
/// [`DecoderSession::process`] for a complete loop.
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub enum DecodeOperation {
    /// More input may follow; an empty slice does not end the input.
    #[default]
    Process,
    /// This is all the remaining input.
    Finish,
}

/// Why an incremental call stopped.
///
/// Handle the counts in [`DecodeProgress`] before acting on the status.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum DecoderStatus {
    /// All offered input was consumed and more is needed.
    NeedsInput,
    /// Decoded output is waiting for output space.
    NeedsOutput,
    /// The operation is complete: every frame was decoded and validated,
    /// and all output delivered.
    Finished,
}

/// Bytes accepted and delivered by one successful call.
///
/// Advance the input by `consumed` and use `output[..produced]`. The counts
/// are per call, not cumulative.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct DecodeProgress {
    /// Length of the accepted input prefix.
    pub consumed: usize,
    /// Length of the written output prefix.
    pub produced: usize,
    /// Why the call stopped.
    pub status: DecoderStatus,
}

/// A failed incremental call, with that call's exact progress.
///
/// Output counted by `produced` was written before the failure and is not
/// rolled back; it is a partial, unvalidated result. The session stays
/// failed: later calls report [`DecodeError::InvalidState`] until the
/// decoder starts a new operation.
///
/// # Examples
///
/// ```
/// use mzstd::{DecodeError, DecodeOperation, Decompressor};
/// let mut decoder = Decompressor::default();
/// let mut session = decoder.start();
/// // The input ends before any frame.
/// let failure = session.process(&[], &mut [], DecodeOperation::Finish).unwrap_err();
/// assert_eq!((failure.consumed, failure.produced), (0, 0));
/// assert!(matches!(failure.into_error(), DecodeError::UnexpectedEndOfInput));
/// ```
#[derive(Debug, thiserror::Error)]
#[error("{error}")]
pub struct DecodeFailure {
    /// The underlying error.
    #[source]
    pub error: DecodeError,
    /// Length of the input prefix the call accepted.
    pub consumed: usize,
    /// Length of the output prefix the call wrote.
    pub produced: usize,
}

impl DecodeFailure {
    /// The error, once the progress has been handled.
    #[must_use]
    pub fn into_error(self) -> DecodeError {
        self.error
    }
}

impl From<DecodeFailure> for DecodeError {
    fn from(failure: DecodeFailure) -> Self {
        failure.error
    }
}

/// An incremental decoding operation borrowing its [`Decompressor`].
///
/// Created by [`Decompressor::start`]. Input and output move through
/// caller-owned slices of any size; the session keeps its own buffers for
/// the headers and block being assembled and for regenerated output not yet
/// delivered. The decoder's [`RetentionPolicy`](crate::RetentionPolicy)
/// is applied when the operation finishes or fails, and again when the next
/// operation starts. Dropping the session abandons the operation.
#[derive(Debug)]
pub struct DecoderSession<'d> {
    decoder: &'d mut Decompressor,
}

impl<'d> DecoderSession<'d> {
    pub(super) fn new(decoder: &'d mut Decompressor) -> Self {
        decoder.apply_retention();
        let config = decoder.config;
        decoder.stream.reset(config);
        Self { decoder }
    }
}

/// Runs one call on `decoder` and applies its retention policy once the
/// operation has ended, successfully or not.
fn process_on(
    decoder: &mut Decompressor,
    input: &[u8],
    output: &mut [u8],
    operation: DecodeOperation,
) -> Result<DecodeProgress, DecodeFailure> {
    let result = decoder.stream.process(input, output, operation);
    if result
        .as_ref()
        .map_or(true, |progress| progress.status == DecoderStatus::Finished)
    {
        decoder.apply_retention();
    }
    result
}

impl DecoderSession<'_> {
    /// Decodes as much of `input` into `output` as it can.
    ///
    /// Deliver `output[..progress.produced]` and advance `input` by
    /// `progress.consumed` after every call. With
    /// [`DecodeOperation::Process`] an empty `input` only drains pending
    /// output. After the first [`DecodeOperation::Finish`], keep using
    /// `Finish` with exactly the unconsumed suffix until the status is
    /// [`DecoderStatus::Finished`].
    ///
    /// In [`super::FrameMode::Single`] the session finishes right after its
    /// frame and leaves any later input unconsumed. In
    /// [`super::FrameMode::Concatenated`] a frame boundary reports
    /// `NeedsInput` until `Finish` confirms that the input ends there. A
    /// finished session reports zero progress and `Finished` again.
    ///
    /// # Errors
    ///
    /// Returns a [`DecodeFailure`] with the call's progress for invalid
    /// data, exceeded limits, failed checks, and input that ends inside a
    /// frame ([`DecodeError::UnexpectedEndOfInput`]). Changing the operation
    /// or the final input after `Finish`, or calling a failed session,
    /// reports [`DecodeError::InvalidState`]. Failures are terminal.
    ///
    /// # Examples
    ///
    /// ```
    /// use mzstd::{DecodeOperation, DecoderStatus, Decompressor};
    /// // "hello" as one raw block in a single-segment frame.
    /// let frame = [0x28, 0xb5, 0x2f, 0xfd, 0x20, 0x05, 0x29, 0x00, 0x00, b'h', b'e', b'l', b'l', b'o'];
    /// let mut decoder = Decompressor::default();
    /// let mut session = decoder.start();
    /// let mut remaining = &frame[..];
    /// let mut decoded = Vec::new();
    /// loop {
    ///     let mut buffer = [0u8; 2];
    ///     let progress = session.process(remaining, &mut buffer, DecodeOperation::Finish)?;
    ///     remaining = &remaining[progress.consumed..];
    ///     decoded.extend_from_slice(&buffer[..progress.produced]);
    ///     if progress.status == DecoderStatus::Finished {
    ///         break;
    ///     }
    /// }
    /// assert_eq!(decoded, b"hello");
    /// assert_eq!(session.total_in(), frame.len() as u64);
    /// assert_eq!(session.total_out(), 5);
    /// assert_eq!(session.frames_decoded(), 1);
    /// assert_eq!(session.window_size(), Some(5));
    /// # Ok::<(), mzstd::DecodeFailure>(())
    /// ```
    pub fn process(
        &mut self,
        input: &[u8],
        output: &mut [u8],
        operation: DecodeOperation,
    ) -> Result<DecodeProgress, DecodeFailure> {
        process_on(self.decoder, input, output, operation)
    }

    /// Delivers pending output without offering input or ending it.
    ///
    /// Shorthand for [`Self::process`] with empty input and
    /// [`DecodeOperation::Process`].
    ///
    /// # Errors
    ///
    /// As [`Self::process`]; after a `Finish` this reports
    /// [`DecodeError::InvalidState`].
    pub fn flush(&mut self, output: &mut [u8]) -> Result<DecodeProgress, DecodeFailure> {
        self.process(&[], output, DecodeOperation::Process)
    }

    /// Ends the input with nothing left to offer, delivering pending output.
    ///
    /// Shorthand for [`Self::process`] with empty input and
    /// [`DecodeOperation::Finish`]; repeat it while it reports
    /// [`DecoderStatus::NeedsOutput`].
    ///
    /// # Errors
    ///
    /// As [`Self::process`]: stopping inside a frame reports
    /// [`DecodeError::UnexpectedEndOfInput`].
    pub fn finish(&mut self, output: &mut [u8]) -> Result<DecodeProgress, DecodeFailure> {
        self.process(&[], output, DecodeOperation::Finish)
    }

    /// Whether the operation completed.
    #[must_use]
    pub fn is_finished(&self) -> bool {
        self.decoder.stream.is_finished()
    }

    /// Compressed bytes accepted, including by failing calls.
    #[must_use]
    pub fn total_in(&self) -> u64 {
        self.decoder.stream.total_in()
    }

    /// Decompressed bytes delivered, including by failing calls.
    #[must_use]
    pub fn total_out(&self) -> u64 {
        self.decoder.stream.total_out()
    }

    /// Zstandard frames completed and validated; skippable frames are not
    /// counted.
    #[must_use]
    pub fn frames_decoded(&self) -> u64 {
        self.decoder.stream.frames_decoded()
    }

    /// The window size of the most recent frame header, if any was read.
    #[must_use]
    pub fn window_size(&self) -> Option<u64> {
        self.decoder.stream.window_size()
    }
}

/// An incremental decoding operation that owns its [`Decompressor`].
///
/// The same state machine as [`DecoderSession`], without a borrow: it suits
/// wrappers that store the codec state in a struct between calls. Created by
/// [`Decompressor::into_session`]; [`Self::into_decompressor`] hands the
/// decoder back.
///
/// # Examples
///
/// ```
/// use mzstd::{DecodeOperation, DecoderStatus, Decompressor};
/// let frame = [0x28, 0xb5, 0x2f, 0xfd, 0x20, 0x02, 0x11, 0x00, 0x00, b'o', b'k'];
/// let mut session = Decompressor::default().into_session();
/// let mut output = [0u8; 8];
/// let progress = session.process(&frame, &mut output, DecodeOperation::Finish)?;
/// assert_eq!(progress.status, DecoderStatus::Finished);
/// assert_eq!(&output[..progress.produced], b"ok");
///
/// session.reinit();
/// let progress = session.process(&frame, &mut output, DecodeOperation::Finish)?;
/// assert_eq!(&output[..progress.produced], b"ok");
/// let mut decoder = session.into_decompressor();
/// assert_eq!(decoder.decompress(&frame)?, b"ok");
/// # Ok::<(), Box<dyn std::error::Error>>(())
/// ```
#[derive(Debug)]
pub struct DecoderSessionOwned {
    decoder: Decompressor,
}

impl DecoderSessionOwned {
    pub(super) fn new(mut decoder: Decompressor) -> Self {
        decoder.apply_retention();
        let config = decoder.config;
        decoder.stream.reset(config);
        Self { decoder }
    }

    /// Decodes as much of `input` into `output` as it can, exactly as
    /// [`DecoderSession::process`].
    ///
    /// # Errors
    ///
    /// As [`DecoderSession::process`].
    pub fn process(
        &mut self,
        input: &[u8],
        output: &mut [u8],
        operation: DecodeOperation,
    ) -> Result<DecodeProgress, DecodeFailure> {
        process_on(&mut self.decoder, input, output, operation)
    }

    /// Delivers pending output, as [`DecoderSession::flush`].
    ///
    /// # Errors
    ///
    /// As [`DecoderSession::process`].
    pub fn flush(&mut self, output: &mut [u8]) -> Result<DecodeProgress, DecodeFailure> {
        self.process(&[], output, DecodeOperation::Process)
    }

    /// Ends the input, as [`DecoderSession::finish`].
    ///
    /// # Errors
    ///
    /// As [`DecoderSession::process`].
    pub fn finish(&mut self, output: &mut [u8]) -> Result<DecodeProgress, DecodeFailure> {
        self.process(&[], output, DecodeOperation::Finish)
    }

    /// Whether the operation completed.
    #[must_use]
    pub fn is_finished(&self) -> bool {
        self.decoder.stream.is_finished()
    }

    /// Compressed bytes accepted, including by failing calls.
    #[must_use]
    pub fn total_in(&self) -> u64 {
        self.decoder.stream.total_in()
    }

    /// Decompressed bytes delivered, including by failing calls.
    #[must_use]
    pub fn total_out(&self) -> u64 {
        self.decoder.stream.total_out()
    }

    /// Zstandard frames completed and validated.
    #[must_use]
    pub fn frames_decoded(&self) -> u64 {
        self.decoder.stream.frames_decoded()
    }

    /// The window size of the most recent frame header, if any was read.
    #[must_use]
    pub fn window_size(&self) -> Option<u64> {
        self.decoder.stream.window_size()
    }

    /// Abandons the current operation, whatever its state, and starts a new
    /// one in place, applying the decoder's retention policy in between.
    pub fn reinit(&mut self) {
        self.decoder.apply_retention();
        let config = self.decoder.config;
        self.decoder.stream.reset(config);
    }

    /// Ends the operation and returns the decoder, ready for reuse, after
    /// applying its retention policy.
    #[must_use]
    pub fn into_decompressor(mut self) -> Decompressor {
        self.decoder.apply_retention();
        self.decoder
    }
}