qubit_codec/transcode/io/

transcode_decode_input.rs

// =============================================================================
//    Copyright (c) 2026 Haixing Hu.
//
//    SPDX-License-Identifier: Apache-2.0
//
//    Licensed under the Apache License, Version 2.0.
// =============================================================================
//! Buffered input driver that decodes units into values.

use core::fmt;
use std::collections::TryReserveError;
use std::io::{
    Error,
    ErrorKind,
    Read,
    Result,
    Seek,
    SeekFrom,
};

use qubit_io::{
    Buffer,
    BufferedInput,
    Input,
    Seekable,
    UncheckedSlice,
};

use crate::{
    Codec,
    DecodeLifecycleOutput,
    DecodeLifecycleProgress,
    TranscodeFailure,
    TranscodeProgress,
    TranscodeStatus,
    Transcoder,
    codec::assert_unit_bounds,
};

use super::codec_decode_driver::CodecDecodeDriver;
use super::transcode_progress_validation::validate_decode_progress;

/// Decodes an [`Input`] unit stream into an [`Input`] value stream.
///
/// This type owns only the unit-level [`qubit_io::BufferedInput`]. Callers pass
/// a streaming [`Transcoder`] and error mapper to each decode operation, which
/// lets one buffered input drive different decoders without nesting buffers or
/// storing decoder-specific state in the buffer owner.
///
/// # Type Parameters
///
/// * `I` - Wrapped unit input.
pub struct TranscodeDecodeInput<I>
where
    I: Input,
    I::Item: Copy + Default,
{
    input: BufferedInput<I>,
}

impl<I> TranscodeDecodeInput<I>
where
    I: Input,
    I::Item: Copy + Default,
{
    /// Creates a decoder input with the default unit buffer capacity.
    ///
    /// # Parameters
    ///
    /// * `inner` - Unit input read by this adapter.
    ///
    /// # Returns
    ///
    /// A new buffered decoder input.
    #[must_use]
    pub fn new(inner: I) -> Self {
        Self {
            input: BufferedInput::new(inner),
        }
    }

    /// Creates a decoder input with a unit buffer of at least `capacity`.
    ///
    /// # Parameters
    ///
    /// * `inner` - Unit input read by this adapter.
    /// * `capacity` - Requested internal unit buffer capacity.
    ///
    /// # Returns
    ///
    /// A new buffered decoder input.
    #[must_use]
    pub fn with_capacity(inner: I, capacity: usize) -> Self {
        Self {
            input: BufferedInput::with_capacity(inner, capacity),
        }
    }

    /// Creates a decoder input with a unit buffer of at least `capacity`.
    ///
    /// # Parameters
    ///
    /// * `inner` - Unit input read by this adapter.
    /// * `capacity` - Requested internal unit buffer capacity.
    ///
    /// # Errors
    ///
    /// Returns an allocation error when the requested buffer cannot be
    /// allocated.
    #[inline]
    pub fn try_with_capacity(
        inner: I,
        capacity: usize,
    ) -> std::result::Result<Self, TryReserveError> {
        Ok(Self {
            input: BufferedInput::try_with_capacity(inner, capacity)?,
        })
    }

    /// Returns a shared reference to the wrapped unit input.
    ///
    /// The wrapped input can be physically positioned after units that remain
    /// unread in this adapter's internal buffer.
    ///
    /// # Returns
    ///
    /// A shared reference to the wrapped unit input.
    #[must_use]
    pub const fn inner(&self) -> &I {
        self.input.inner()
    }

    /// Returns the number of unread units currently buffered.
    ///
    /// # Returns
    ///
    /// The number of unread units in the internal buffer.
    #[must_use]
    pub fn unread_len(&self) -> usize {
        self.input.unread_len()
    }

    /// Returns the currently buffered unread units.
    ///
    /// # Returns
    ///
    /// Returns a shared slice over the unread portion of the internal unit
    /// buffer. The slice is valid until this adapter is mutated.
    #[must_use]
    pub fn unread(&self) -> &[I::Item] {
        self.input.unread()
    }

    /// Returns the internal unit buffer capacity.
    ///
    /// # Returns
    ///
    /// The maximum number of units retained in the internal buffer.
    #[must_use]
    pub fn capacity(&self) -> usize {
        self.input.capacity()
    }

    /// Refills the internal buffer until at least `count` unread units are
    /// available.
    ///
    /// # Parameters
    ///
    /// * `count` - Minimum number of unread units required.
    ///
    /// # Errors
    ///
    /// Returns allocation errors mapped to [`ErrorKind::OutOfMemory`], or I/O
    /// errors from the wrapped input while refilling.
    pub fn fill_until(&mut self, count: usize) -> std::io::Result<bool> {
        if count > self.input.capacity() {
            self.input
                .try_reserve_capacity(count)
                .map_err(allocation_to_io_error)?;
        }
        self.input.fill_until(count)
    }

    /// Performs one decode step and preserves the transcoder stop status.
    ///
    /// `None` means that the wrapped input reached EOF with no buffered units.
    /// Consumed input is committed before this method returns. Callers can use
    /// the returned [`TranscodeProgress`] to distinguish `NeedOutput` from an
    /// incomplete input tail without guessing from the unread buffer alone.
    pub fn transcode_step<D, M, Value>(
        &mut self,
        decoder: &mut D,
        map_error: &mut M,
        output: &mut [Value],
        output_index: usize,
        count: usize,
    ) -> Result<Option<TranscodeProgress>>
    where
        D: Transcoder<Input = I::Item, Output = Value>,
        M: FnMut(D::Error) -> Error,
    {
        let output_end = UncheckedSlice::checked_range_end(
            output.len(),
            output_index,
            count,
            "decoded output range exceeds destination buffer",
        )?;
        if count == 0 {
            return Ok(Some(TranscodeProgress::complete(0, 0)));
        }
        if self.unread_len() == 0 && !self.fill_more()? {
            return Ok(None);
        }
        let available_input = self.unread_len();
        let progress = decoder
            .transcode(
                self.unread(),
                0,
                &mut output[..output_end],
                output_index,
            )
            .map_err(&mut *map_error)
            .and_then(|progress| {
                validate_decode_progress(
                    progress,
                    0,
                    available_input,
                    output_index,
                    count,
                )
            })?;
        self.consume(progress.read());
        Ok(Some(progress))
    }

    /// Consumes unread units from the current buffer window.
    ///
    /// # Parameters
    ///
    /// * `count` - Number of unread units to discard.
    ///
    /// # Panics
    ///
    /// Panics when `count` exceeds [`Self::unread_len`].
    pub fn consume(&mut self, count: usize) {
        assert!(
            count <= self.unread_len(),
            "cannot consume beyond buffered input",
        );
        // SAFETY: The caller-provided count is within the unread window.
        unsafe {
            self.input.consume(count);
        }
    }

    /// Copies unread units into an indexed output range without consuming them.
    ///
    /// # Parameters
    ///
    /// * `output` - Destination storage that receives a copy of unread units.
    /// * `output_index` - Start index inside `output`.
    /// * `count` - Number of unread units to copy.
    ///
    /// # Safety
    ///
    /// The caller must guarantee that `output_index..output_index + count` is
    /// a valid range inside `output`, that the addition does not overflow, that
    /// `count <= self.unread_len()`, and that the destination range does not
    /// overlap with the unread units stored inside this buffer.
    pub unsafe fn copy_unread_to(
        &self,
        output: &mut [I::Item],
        output_index: usize,
        count: usize,
    ) {
        // SAFETY: The caller guarantees the destination range and non-overlap
        // requirements for the unread copy.
        let unread = self.unread();
        unsafe {
            UncheckedSlice::copy_nonoverlapping(
                unread,
                0,
                output,
                output_index,
                count,
            );
        }
    }

    /// Consumes this adapter and returns its parts.
    ///
    /// # Returns
    ///
    /// The wrapped input and the buffer holding unread units.
    #[must_use = "the returned input and unread buffer must be handled"]
    pub fn into_parts(self) -> (I, Buffer<I::Item>) {
        self.input.into_parts()
    }

    /// Reads buffered units into an indexed output range.
    ///
    /// # Parameters
    ///
    /// * `output` - Destination unit storage.
    /// * `output_index` - Start index inside `output`.
    /// * `count` - Maximum number of units to read.
    ///
    /// # Returns
    ///
    /// The number of units copied into `output`.
    ///
    /// # Errors
    ///
    /// Returns input or buffer validation errors from the wrapped
    /// [`qubit_io::BufferedInput`].
    ///
    /// # Safety
    ///
    /// The caller must guarantee that `output_index..output_index + count` is
    /// a valid range inside `output` and that the addition does not overflow.
    pub unsafe fn read_unchecked(
        &mut self,
        output: &mut [I::Item],
        output_index: usize,
        count: usize,
    ) -> Result<usize> {
        debug_assert!(
            UncheckedSlice::range_fits(output.len(), output_index, count),
            "unchecked read output range exceeds destination buffer",
        );
        if count == 0 {
            return Ok(0);
        }
        // SAFETY: The caller guarantees the destination range is valid.
        unsafe { self.input.read_unchecked(output, output_index, count) }
    }

    /// Decodes one codec value from the buffered unit input.
    ///
    /// This strict convenience method supports codecs whose decode reset and
    /// finish phases do not emit values. Use
    /// [`Self::read_decoded_lifecycle_with`] when those phases have output.
    ///
    /// # Parameters
    ///
    /// * `codec` - Codec used for this single-value decode.
    /// * `map_error` - Mapper for codec-domain invalid-input errors.
    ///
    /// # Returns
    ///
    /// Returns one decoded codec value.
    ///
    /// # Errors
    ///
    /// Returns [`ErrorKind::Unsupported`] before invoking the codec or reading
    /// input when reset or finish may emit values. Otherwise, returns I/O
    /// errors from the wrapped input, `UnexpectedEof` when EOF occurs before a
    /// complete value is available, `InvalidData` when the codec reports an
    /// impossible incomplete state, or the error returned by `map_error` for
    /// codec reset, invalid input, or codec finish failures.
    ///
    /// # Panics
    ///
    /// Panics when the codec reports more reset or finish values than its
    /// declared bounds.
    pub fn read_decoded_with<C, M>(
        &mut self,
        codec: &mut C,
        map_error: M,
    ) -> Result<C::Value>
    where
        C: Codec<Unit = I::Item>,
        M: FnMut(C::DecodeError) -> Error,
    {
        TranscodeFailure::ensure_no_decode_lifecycle_output::<C>()
            .map_err(|error| Error::new(ErrorKind::Unsupported, error))?;
        let progress = self.read_decoded_lifecycle_with_scratch_impl(
            codec,
            &mut [],
            &mut [],
            map_error,
        )?;
        let (value, reset_written, finish_written) = progress.into_parts();
        debug_assert_eq!(0, reset_written);
        debug_assert_eq!(0, finish_written);
        Ok(value)
    }

    /// Decodes one complete codec lifecycle and preserves every output phase.
    ///
    /// The method allocates independent storage for values emitted by
    /// `decode_reset` and `decode_finish`, so neither phase can overwrite the
    /// other.
    ///
    /// # Parameters
    ///
    /// - `codec`: Codec used to decode one complete lifecycle.
    /// - `map_error`: Maps codec decode errors to I/O errors.
    ///
    /// # Returns
    ///
    /// Returns reset output, the main decoded value, and finish output.
    ///
    /// # Errors
    ///
    /// Returns I/O errors from the wrapped input, `UnexpectedEof` when EOF
    /// occurs before a complete value is available, `InvalidData` when the
    /// codec reports an impossible incomplete state, or the error returned
    /// by `map_error` for codec reset, invalid input, or codec finish failures.
    ///
    /// # Panics
    ///
    /// Panics when the codec reports more reset or finish values than its
    /// declared bounds.
    pub fn read_decoded_lifecycle_with<C, M>(
        &mut self,
        codec: &mut C,
        map_error: M,
    ) -> Result<DecodeLifecycleOutput<C::Value>>
    where
        C: Codec<Unit = I::Item>,
        C::Value: Default,
        M: FnMut(C::DecodeError) -> Error,
    {
        assert_unit_bounds::<C>();
        let mut reset_output = Vec::new();
        reset_output.resize_with(C::MAX_DECODE_RESET_VALUES, C::Value::default);
        let mut finish_output = Vec::new();
        finish_output
            .resize_with(C::MAX_DECODE_FINISH_VALUES, C::Value::default);
        let progress = self.read_decoded_lifecycle_with_scratch_impl(
            codec,
            &mut reset_output,
            &mut finish_output,
            map_error,
        )?;
        let (value, reset_written, finish_written) = progress.into_parts();
        reset_output.truncate(reset_written);
        finish_output.truncate(finish_written);
        Ok(DecodeLifecycleOutput::new(
            reset_output,
            value,
            finish_output,
        ))
    }

    /// Decodes one complete codec lifecycle into separate caller storage.
    ///
    /// # Parameters
    ///
    /// - `codec`: Codec used to decode one complete lifecycle.
    /// - `reset_output`: Reusable storage for values emitted by
    ///   [`Codec::decode_reset`].
    /// - `finish_output`: Reusable storage for values emitted by
    ///   [`Codec::decode_finish`].
    /// - `map_error`: Maps codec decode errors to I/O errors.
    ///
    /// # Returns
    ///
    /// Returns the main decoded value and the initialized lengths of both
    /// lifecycle output buffers.
    ///
    /// # Errors
    ///
    /// Returns [`ErrorKind::InvalidInput`] when either output buffer is shorter
    /// than the corresponding codec bound. Returns I/O errors from the wrapped
    /// input, `UnexpectedEof` when EOF occurs before a complete value is
    /// available, `InvalidData` when the codec reports an impossible
    /// incomplete state, or the error returned by `map_error` for codec reset,
    /// invalid input, or codec finish failures.
    ///
    /// # Panics
    ///
    /// Panics when the codec reports more reset or finish values than its
    /// declared bounds.
    pub fn read_decoded_lifecycle_with_scratch<C, M>(
        &mut self,
        codec: &mut C,
        reset_output: &mut [C::Value],
        finish_output: &mut [C::Value],
        map_error: M,
    ) -> Result<DecodeLifecycleProgress<C::Value>>
    where
        C: Codec<Unit = I::Item>,
        M: FnMut(C::DecodeError) -> Error,
    {
        self.read_decoded_lifecycle_with_scratch_impl(
            codec,
            reset_output,
            finish_output,
            map_error,
        )
    }

    fn read_decoded_lifecycle_with_scratch_impl<C, M>(
        &mut self,
        codec: &mut C,
        reset_output: &mut [C::Value],
        finish_output: &mut [C::Value],
        mut map_error: M,
    ) -> Result<DecodeLifecycleProgress<C::Value>>
    where
        C: Codec<Unit = I::Item>,
        M: FnMut(C::DecodeError) -> Error,
    {
        assert_unit_bounds::<C>();
        if reset_output.len() < C::MAX_DECODE_RESET_VALUES {
            return Err(Error::new(
                ErrorKind::InvalidInput,
                "decode reset output is shorter than the codec reset bound",
            ));
        }
        if finish_output.len() < C::MAX_DECODE_FINISH_VALUES {
            return Err(Error::new(
                ErrorKind::InvalidInput,
                "decode finish output is shorter than the codec finish bound",
            ));
        }
        let reset_written = unsafe {
            // SAFETY: The reset output length check above reserves the codec's
            // declared decode-reset output bound.
            codec.decode_reset(reset_output, 0)
        }
        .map_err(&mut map_error)?;
        assert!(
            reset_written <= C::MAX_DECODE_RESET_VALUES,
            "Codec::decode_reset wrote beyond its reset bound",
        );

        let value = CodecDecodeDriver::new(&mut self.input)
            .read_one(codec, &mut map_error)?;

        let finish_written = unsafe {
            // SAFETY: The finish output length check above reserves the
            // codec's declared decode-finish output bound.
            codec.decode_finish(finish_output, 0)
        }
        .map_err(&mut map_error)?;
        assert!(
            finish_written <= C::MAX_DECODE_FINISH_VALUES,
            "Codec::decode_finish wrote beyond its finish bound",
        );
        Ok(DecodeLifecycleProgress::new(
            value,
            reset_written,
            finish_written,
        ))
    }

    /// Runs decoder reset into an indexed output range.
    ///
    /// This method does not read or consume buffered input units.
    ///
    /// # Errors
    ///
    /// Returns invalid output ranges, insufficient declared output capacity,
    /// capacity errors, or reset errors mapped by `map_error`.
    ///
    /// # Panics
    ///
    /// Panics when the decoder reports writing more values than its declared
    /// reset bound.
    pub fn reset<D, M, Value>(
        &self,
        decoder: &mut D,
        map_error: &mut M,
        output: &mut [Value],
        output_index: usize,
        count: usize,
    ) -> Result<usize>
    where
        D: Transcoder<Input = I::Item, Output = Value>,
        M: FnMut(D::Error) -> Error,
    {
        let required = decoder
            .max_reset_output_len()
            .map_err(capacity_to_io_error)?;
        let output_end = UncheckedSlice::checked_range_end(
            output.len(),
            output_index,
            count,
            "reset output range exceeds destination buffer",
        )?;
        if count < required {
            return Err(Error::new(
                ErrorKind::InvalidData,
                "insufficient output for decoder reset bound",
            ));
        }
        let output = &mut output[..output_end];
        let written = decoder
            .reset(output, output_index)
            .map_err(&mut *map_error)?;
        assert!(written <= required, "reset wrote beyond its bound");
        Ok(written)
    }

    /// Decodes values into an indexed output range using a streaming
    /// [`Transcoder`].
    ///
    /// # Parameters
    ///
    /// * `decoder` - Streaming decoder used for this operation.
    /// * `map_error` - Function mapping transcode errors into I/O errors.
    /// * `output` - Destination value storage.
    /// * `output_index` - Start index inside `output`.
    /// * `count` - Maximum number of values to write.
    ///
    /// # Returns
    ///
    /// The number of values written. Incomplete EOF tails are left buffered
    /// and reported as `Ok(written)`, so callers can apply their own EOF
    /// policy.
    ///
    /// # Errors
    ///
    /// Returns input errors, invalid output ranges, capacity errors from the
    /// internal buffer, or transcode errors mapped by `map_error`.
    pub fn transcode<D, M, Value>(
        &mut self,
        decoder: &mut D,
        map_error: &mut M,
        output: &mut [Value],
        output_index: usize,
        count: usize,
    ) -> Result<usize>
    where
        D: Transcoder<Input = I::Item, Output = Value>,
        M: FnMut(D::Error) -> Error,
    {
        let output_end = UncheckedSlice::checked_range_end(
            output.len(),
            output_index,
            count,
            "decoded output range exceeds destination buffer",
        )?;
        if count == 0 {
            return Ok(0);
        }
        let output = &mut output[..output_end];
        let mut written_total = 0;
        loop {
            if self.unread_len() == 0 && !self.fill_more()? {
                return Ok(written_total);
            }
            let units = self.unread();
            let available_input = units.len();
            let remaining_output = count - written_total;
            let progress = decoder
                .transcode(units, 0, output, output_index + written_total)
                .map_err(&mut *map_error)?;
            let progress = validate_decode_progress(
                progress,
                0,
                available_input,
                output_index + written_total,
                remaining_output,
            )?;
            let consumed = progress.read();
            let written = progress.written();
            self.consume(consumed);
            written_total += written;
            match progress.status() {
                TranscodeStatus::Complete => {
                    if written_total == count {
                        return Ok(written_total);
                    }
                }
                TranscodeStatus::NeedOutput { .. } => {
                    return Ok(written_total);
                }
                TranscodeStatus::NeedInput { required, .. } => {
                    if self.fill_until(required.get())? {
                        continue;
                    }
                    return Ok(written_total);
                }
            }
        }
    }

    /// Finishes a streaming decoder into an indexed output range.
    ///
    /// # Parameters
    ///
    /// * `decoder` - Streaming decoder whose final output is being collected.
    /// * `map_error` - Function mapping transcode errors into I/O errors.
    /// * `output` - Destination value storage.
    /// * `output_index` - Start index inside `output`.
    /// * `count` - Maximum number of finish values to write.
    ///
    /// # Returns
    ///
    /// The number of values written by the decoder finish operation.
    ///
    /// # Errors
    ///
    /// Returns invalid output ranges, capacity errors, or transcode
    /// finalization errors mapped to I/O errors.
    ///
    /// # Panics
    ///
    /// Panics when the decoder reports writing more values than its
    /// [`Transcoder::max_finish_output_len`] bound.
    pub fn finish<D, M, Value>(
        &self,
        decoder: &mut D,
        map_error: &mut M,
        output: &mut [Value],
        output_index: usize,
        count: usize,
    ) -> Result<usize>
    where
        D: Transcoder<Input = I::Item, Output = Value>,
        M: FnMut(D::Error) -> Error,
    {
        let required = decoder
            .max_finish_output_len()
            .map_err(capacity_to_io_error)?;
        // Validate the caller-supplied count range first (InvalidInput).
        let output_end = UncheckedSlice::checked_range_end(
            output.len(),
            output_index,
            count,
            "finish output range exceeds destination buffer",
        )?;
        // `count` is the caller's declared writable finish range.  The
        // destination slice may be larger, but passing extra capacity to the
        // decoder would allow it to write beyond the range the caller granted.
        if count < required {
            return Err(Error::new(
                ErrorKind::InvalidData,
                "insufficient output for decoder finish bound",
            ));
        }
        let output = &mut output[..output_end];
        let written = decoder
            .finish(output, output_index)
            .map_err(&mut *map_error)?;
        assert!(written <= required, "finish wrote beyond its bound");
        Ok(written)
    }
}

impl<I> TranscodeDecodeInput<I>
where
    I: Input<Item = u8> + Seekable<Unit = u8>,
{
    /// Seeks the wrapped byte input and discards buffered bytes after success.
    ///
    /// # Parameters
    ///
    /// * `position` - Target seek position.
    ///
    /// # Returns
    ///
    /// The new stream position reported by the wrapped input.
    ///
    /// # Errors
    ///
    /// Returns seek errors from the wrapped input.
    pub fn seek(&mut self, position: SeekFrom) -> Result<u64> {
        self.input.seek_to(position)
    }
}

impl<I> TranscodeDecodeInput<I>
where
    I: Input,
    I::Item: Copy + Default,
{
    /// Refills the underlying buffer.
    fn fill_more(&mut self) -> Result<bool> {
        self.input.fill_more()
    }
}

impl<I> Read for TranscodeDecodeInput<I>
where
    I: Input<Item = u8>,
{
    /// Reads raw bytes through the internal buffer.
    fn read(&mut self, output: &mut [u8]) -> Result<usize> {
        // SAFETY: The full output slice is a valid destination range.
        unsafe { self.read_unchecked(output, 0, output.len()) }
    }
}

impl<I> Seek for TranscodeDecodeInput<I>
where
    I: Input<Item = u8> + Seekable<Unit = u8>,
{
    /// Seeks the wrapped byte input and discards buffered bytes after success.
    fn seek(&mut self, position: SeekFrom) -> Result<u64> {
        self.seek(position)
    }
}

impl<I> fmt::Debug for TranscodeDecodeInput<I>
where
    I: Input,
    I::Item: Copy + Default,
    BufferedInput<I>: fmt::Debug,
{
    /// Formats this buffered decode input for debugging.
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("TranscodeDecodeInput")
            .field("input", &self.input)
            .finish()
    }
}

/// Converts a capacity planning failure into an I/O error.
fn capacity_to_io_error(error: crate::CapacityError) -> Error {
    Error::new(ErrorKind::InvalidData, error)
}

/// Converts an allocation failure into an I/O boundary error.
fn allocation_to_io_error(error: std::collections::TryReserveError) -> Error {
    Error::new(ErrorKind::OutOfMemory, error)
}