qubit_io/buffered/

buffered_byte_output.rs

// =============================================================================
//    Copyright (c) 2026 Haixing Hu.
//
//    SPDX-License-Identifier: Apache-2.0
//
//    Licensed under the Apache License, Version 2.0.
// =============================================================================

use std::io::{
    Error,
    ErrorKind,
    Result,
    Seek,
    SeekFrom,
    Write,
};

use crate::Buffer;
use crate::WriteExt;
use crate::buffered::DEFAULT_BUFFER_CAPACITY;

/// Buffered byte output over a wrapped writer.
///
/// This type keeps a fixed-size byte buffer in front of an underlying writer so
/// small byte writes can be accumulated before they are written to the I/O
/// target. Large writes may bypass the buffer after pending buffered bytes
/// have been flushed.
///
/// `BufferedByteOutput` is deliberately byte-oriented. It performs no binary
/// encoding, text encoding, or record framing. Higher-level writers can either
/// use the standard [`Write`] implementation or write directly into
/// [`Self::spare_buffer_mut`] or [`Self::spare_raw_parts_mut`] and then call
/// [`Self::advance`] or [`Self::advance_unchecked`] after validating the range
/// they initialized. Callers that need to recover the wrapped writer should
/// call [`Write::flush`] first, then use [`Self::into_parts`].
#[derive(Debug)]
pub struct BufferedByteOutput<W> {
    inner: W,
    buffer: Buffer<u8>,
}

impl<W> BufferedByteOutput<W> {
    /// Creates a buffered byte output with the default capacity.
    ///
    /// # Parameters
    ///
    /// * `inner` - The writer that receives bytes when the internal buffer is
    ///   flushed.
    ///
    /// # Returns
    ///
    /// A new buffered byte output using `DEFAULT_BUFFER_CAPACITY`.
    #[inline(always)]
    #[must_use]
    pub fn new(inner: W) -> Self {
        Self::with_capacity(inner, DEFAULT_BUFFER_CAPACITY)
    }

    /// Creates a buffered byte output with at least the requested capacity.
    ///
    /// # Parameters
    ///
    /// * `inner` - The writer that receives bytes when the internal buffer is
    ///   flushed.
    /// * `capacity` - The requested internal buffer capacity in bytes.
    ///
    /// # Returns
    ///
    /// A new buffered byte output whose actual buffer capacity is
    /// `capacity.max(1)`.
    #[inline(always)]
    #[must_use]
    pub fn with_capacity(inner: W, capacity: usize) -> Self {
        Self {
            inner,
            buffer: Buffer::with_capacity(capacity),
        }
    }

    /// Returns a shared reference to the wrapped writer.
    ///
    /// # Returns
    ///
    /// An immutable reference to the underlying writer.  Pending bytes may
    /// still be present in the internal buffer and are not flushed by this
    /// method.
    #[inline(always)]
    pub const fn inner(&self) -> &W {
        &self.inner
    }

    /// Returns an exclusive reference to the wrapped writer.
    ///
    /// Pending bytes may still be present in the internal buffer and are not
    /// flushed by this method.
    ///
    /// # Returns
    ///
    /// A mutable reference to the underlying writer.
    #[inline(always)]
    pub fn inner_mut(&mut self) -> &mut W {
        &mut self.inner
    }

    /// Returns the internal buffer capacity.
    ///
    /// # Returns
    ///
    /// The total number of bytes that can be held by the internal buffer.
    #[inline(always)]
    #[must_use]
    pub fn capacity(&self) -> usize {
        self.buffer.capacity()
    }

    /// Returns the unused capacity in the internal buffer.
    ///
    /// # Returns
    ///
    /// The number of bytes that can still be appended to the internal buffer
    /// before it must be flushed.
    #[inline(always)]
    #[must_use]
    pub fn spare_capacity(&self) -> usize {
        self.buffer.spare_capacity()
    }

    /// Returns the unused portion of the internal buffer.
    ///
    /// Callers may write initialized bytes into the returned slice and then
    /// call [`Self::advance`] with the number of bytes written.
    ///
    /// # Returns
    ///
    /// A mutable slice over the spare buffer capacity.
    #[inline(always)]
    #[must_use]
    pub fn spare_buffer_mut(&mut self) -> &mut [u8] {
        let limit = self.buffer.limit();
        &mut self.buffer.data_mut()[limit..]
    }

    /// Returns raw spare-buffer parts for hot-path callers.
    ///
    /// The returned slice is the full internal backing storage. `index` is the
    /// start of the spare byte window, and `count` is the number of spare
    /// bytes. Callers that need a slice can use `&mut buffer[index..index +
    /// count]`; callers that already validated bounds can pass `buffer` and
    /// `index` directly to indexed unchecked codecs.
    ///
    /// Mutating bytes outside `index..index + count` changes pending output
    /// bytes and may corrupt the logical stream.
    ///
    /// # Returns
    ///
    /// The backing storage, the spare start index, and the spare byte count.
    #[inline(always)]
    #[must_use]
    pub fn spare_raw_parts_mut(&mut self) -> (&mut [u8], usize, usize) {
        let index = self.buffer.limit();
        let count = self.buffer.spare_capacity();
        (self.buffer.data_mut(), index, count)
    }

    /// Marks `count` bytes from [`Self::spare_buffer_mut`] as written.
    ///
    /// # Parameters
    ///
    /// * `count` - Number of bytes initialized by the caller.
    ///
    /// # Panics
    ///
    /// Panics when `count` exceeds [`Self::spare_capacity`].
    #[inline(always)]
    pub fn advance(&mut self, count: usize) {
        assert!(
            count <= self.spare_capacity(),
            "cannot advance beyond spare output buffer"
        );
        // SAFETY: The assertion proves that `count` is within spare capacity.
        unsafe {
            self.buffer.advance_unchecked(count);
        }
    }

    /// Marks spare bytes as written without checking bounds.
    ///
    /// # Parameters
    ///
    /// * `count` - Number of initialized spare bytes to make pending for
    ///   output.
    ///
    /// # Safety
    ///
    /// The caller must guarantee that `count <= self.spare_capacity()` and
    /// that the corresponding bytes returned by [`Self::spare_buffer_mut`]
    /// have been initialized.
    #[inline(always)]
    pub unsafe fn advance_unchecked(&mut self, count: usize) {
        // SAFETY: The caller guarantees that `count` is within spare capacity.
        unsafe {
            self.buffer.advance_unchecked(count);
        }
    }

    /// Writes bytes into the internal buffer without checking spare capacity.
    ///
    /// # Parameters
    ///
    /// * `input` - The source bytes.
    /// * `input_index` - The starting index in `input`.
    /// * `count` - The number of bytes to copy.
    ///
    /// # Safety
    ///
    /// The caller must ensure that `input_index..input_index + count` is valid
    /// in `input`, that `count <= self.spare_capacity()`, and that the copied
    /// source range does not overlap with the destination range in the internal
    /// buffer.
    #[inline(always)]
    unsafe fn write_to_buffer_unchecked(
        &mut self,
        input: &[u8],
        input_index: usize,
        count: usize,
    ) {
        // SAFETY: The caller upholds `Buffer::copy_from_unchecked` range and
        // non-overlap requirements.
        unsafe {
            self.buffer.copy_from_unchecked(input, input_index, count);
        }
    }
}

impl<W> BufferedByteOutput<W>
where
    W: Write,
{
    /// Consumes this buffered output without flushing pending bytes.
    ///
    /// This method performs no I/O. Pending bytes that have been accepted into
    /// the internal buffer but not written to the wrapped writer are returned
    /// as the second tuple item.
    ///
    /// # Returns
    ///
    /// The wrapped writer and pending bytes in logical write order.
    #[inline(always)]
    #[must_use]
    pub fn into_parts(self) -> (W, Vec<u8>) {
        let pending = self.pending_slice().to_vec();
        (self.inner, pending)
    }

    /// Ensures that at least `count` bytes are available in the spare buffer.
    ///
    /// # Parameters
    ///
    /// * `count` - Number of spare bytes required.
    ///
    /// # Errors
    ///
    /// Returns any non-interrupted I/O error produced while flushing buffered
    /// bytes. Returns [`ErrorKind::InvalidInput`] if `count` exceeds the buffer
    /// capacity. Returns [`ErrorKind::InvalidData`] if the wrapped writer
    /// reports more bytes than the pending buffer range contained.
    pub fn ensure_spare_capacity(&mut self, count: usize) -> Result<()> {
        if count > self.buffer.capacity() {
            return Err(Error::new(
                ErrorKind::InvalidInput,
                "requested spare capacity exceeds buffered output capacity",
            ));
        }
        if self.spare_capacity() < count {
            self.flush_buffer()?;
        }
        Ok(())
    }

    /// Writes all bytes through the internal buffer.
    ///
    /// Small inputs are appended to the internal buffer.  Inputs that do not
    /// fit may flush the buffer first, and inputs at least as large as the
    /// buffer may be written directly to the wrapped writer.
    ///
    /// # Parameters
    ///
    /// * `input` - The bytes to write.
    ///
    /// # Returns
    ///
    /// `Ok(())` after all bytes from `input` have been accepted.
    ///
    /// # Errors
    ///
    /// Returns any I/O error produced while flushing pending bytes or writing a
    /// large input directly to the wrapped writer. Flush failures include
    /// [`ErrorKind::WriteZero`] if the writer reports that zero bytes were
    /// written before the buffer is drained, and [`ErrorKind::InvalidData`] if
    /// it reports more bytes than the requested range contained.
    #[inline]
    fn write_all_buffered(&mut self, input: &[u8]) -> Result<()> {
        if input.len() < self.spare_capacity() {
            // SAFETY: The branch proves that the input fits in spare capacity.
            unsafe {
                self.write_to_buffer_unchecked(input, 0, input.len());
            }
            Ok(())
        } else {
            self.write_all_cold(input)
        }
    }

    /// Handles slow-path raw writes that must flush or bypass the buffer.
    ///
    /// # Parameters
    ///
    /// * `input` - The bytes to write after the fast path determined that they
    ///   do not fit comfortably in the current spare buffer capacity.
    ///
    /// # Returns
    ///
    /// `Ok(())` after all bytes from `input` have been accepted either by the
    /// buffer or by the wrapped writer.
    ///
    /// # Errors
    ///
    /// Returns any I/O error produced while flushing pending bytes or writing a
    /// large input directly to the wrapped writer. Flush failures include
    /// [`ErrorKind::WriteZero`] if the writer reports that zero bytes were
    /// written before the buffer is drained, and [`ErrorKind::InvalidData`] if
    /// it reports more bytes than the requested range contained.
    #[cold]
    #[inline(never)]
    fn write_all_cold(&mut self, input: &[u8]) -> Result<()> {
        if input.len() > self.spare_capacity() {
            self.flush_buffer()?;
        }
        if input.len() >= self.buffer.capacity() {
            // SAFETY: The range covers the full source slice.
            unsafe { self.write_all_inner_unchecked(input, 0, input.len()) }
        } else {
            // SAFETY: After the optional flush, any input smaller than the
            // buffer capacity fits in the empty or sufficiently spare buffer.
            unsafe {
                self.write_to_buffer_unchecked(input, 0, input.len());
            }
            Ok(())
        }
    }

    /// Handles slow-path raw writes for [`Write::write`] semantics.
    ///
    /// The method preserves `Write::write` behavior: it may accept fewer bytes
    /// than the input length when the write is delegated directly to the
    /// wrapped writer.
    ///
    /// # Parameters
    ///
    /// * `input` - The bytes to write after the fast path determined that they
    ///   do not fit comfortably in the current spare buffer capacity.
    ///
    /// # Returns
    ///
    /// The number of bytes accepted.  Buffered writes return `input.len()`;
    /// direct writes return the byte count reported by the wrapped writer.
    ///
    /// # Errors
    ///
    /// Returns any I/O error produced while flushing pending bytes or writing a
    /// large input directly to the wrapped writer. Flush failures include
    /// [`ErrorKind::WriteZero`] if the writer reports that zero bytes were
    /// written before the buffer is drained, and [`ErrorKind::InvalidData`] if
    /// it reports more bytes than the requested range contained.
    #[cold]
    #[inline(never)]
    fn write_cold(&mut self, input: &[u8]) -> Result<usize> {
        if input.len() > self.spare_capacity() {
            self.flush_buffer()?;
        }
        if input.len() >= self.buffer.capacity() {
            // SAFETY: The range covers the full source slice.
            unsafe { self.write_inner_unchecked(input, 0, input.len()) }
        } else {
            // SAFETY: After the optional flush, any input smaller than the
            // buffer capacity fits in the empty or sufficiently spare buffer.
            unsafe {
                self.write_to_buffer_unchecked(input, 0, input.len());
            }
            Ok(input.len())
        }
    }

    /// Flushes buffered bytes to the wrapped writer.
    ///
    /// The method retries interrupted writes.  If an error occurs after some
    /// bytes have been written, the already-written bytes are removed from the
    /// front of the buffer and the unwritten suffix is kept for a later retry.
    ///
    /// # Returns
    ///
    /// `Ok(())` once all currently buffered bytes have been written to the
    /// wrapped writer.
    ///
    /// # Errors
    ///
    /// Returns any non-interrupted I/O error produced by the wrapped writer.
    /// Returns [`ErrorKind::WriteZero`] if the writer reports a zero-length
    /// write before all buffered bytes are drained. Returns
    /// [`ErrorKind::InvalidData`] if the writer reports more bytes than the
    /// pending buffer range contained.
    pub fn flush_buffer(&mut self) -> Result<()> {
        while !self.buffer.is_empty() {
            let position = self.buffer.position();
            let available = self.buffer.available();
            // SAFETY: `position..position + available` is the current readable
            // range maintained by `Buffer`.
            match unsafe {
                self.inner.write_unchecked(
                    self.buffer.data(),
                    position,
                    available,
                )
            } {
                Ok(0) => {
                    self.buffer.compact();
                    return Err(Error::new(
                        ErrorKind::WriteZero,
                        "failed to write buffered data",
                    ));
                }
                Ok(written) => {
                    if let Err(error) = validate_write_count(written, available)
                    {
                        self.buffer.compact();
                        return Err(error);
                    }
                    // SAFETY: The validated count is in `0..=available`.
                    unsafe {
                        self.buffer.consume_unchecked(written);
                    }
                }
                Err(error) if error.kind() == ErrorKind::Interrupted => {}
                Err(error) => {
                    self.buffer.compact();
                    return Err(error);
                }
            }
        }
        self.buffer.clear();
        Ok(())
    }

    /// Flushes buffered bytes and then flushes the wrapped writer.
    ///
    /// # Returns
    ///
    /// `Ok(())` once pending buffered bytes have been written and the wrapped
    /// writer's own flush operation succeeds.
    ///
    /// # Errors
    ///
    /// Returns any non-interrupted I/O error produced while flushing buffered
    /// bytes, [`ErrorKind::WriteZero`] if the wrapped writer cannot make
    /// progress while draining the buffer, [`ErrorKind::InvalidData`] if the
    /// writer reports an impossible byte count, or any error returned by
    /// [`Write::flush`] on the wrapped writer.
    #[inline(always)]
    fn flush_all(&mut self) -> Result<()> {
        self.flush_buffer().and_then(|()| self.inner.flush())
    }

    /// Writes bytes from the input slice and reports the accepted byte count.
    ///
    /// This is the buffered implementation for [`Write::write`]-style callers.
    /// Small inputs are appended to the buffer and reported as fully accepted;
    /// large inputs may be delegated to the wrapped writer after pending bytes
    /// are flushed.
    ///
    /// # Parameters
    ///
    /// * `input` - The bytes to write.
    ///
    /// # Returns
    ///
    /// The number of bytes accepted.  Buffered writes return `input.len()`;
    /// direct writes return the byte count reported by the wrapped writer.
    ///
    /// # Errors
    ///
    /// Returns any I/O error produced while flushing pending bytes or writing a
    /// large input directly to the wrapped writer. Flush failures include
    /// [`ErrorKind::WriteZero`] if the writer reports that zero bytes were
    /// written before the buffer is drained, and [`ErrorKind::InvalidData`] if
    /// it reports more bytes than the requested range contained.
    #[inline]
    fn write_from(&mut self, input: &[u8]) -> Result<usize> {
        if input.len() < self.spare_capacity() {
            // SAFETY: The branch proves that the input fits in spare capacity.
            unsafe {
                self.write_to_buffer_unchecked(input, 0, input.len());
            }
            Ok(input.len())
        } else {
            self.write_cold(input)
        }
    }

    /// Flushes pending bytes before seeking the wrapped writer.
    ///
    /// # Parameters
    ///
    /// * `position` - The target seek position passed to the wrapped writer.
    ///
    /// # Returns
    ///
    /// The new stream position reported by the wrapped writer.
    ///
    /// # Errors
    ///
    /// Returns any non-interrupted I/O error produced while flushing buffered
    /// bytes, [`ErrorKind::WriteZero`] if the wrapped writer cannot make
    /// progress while draining the buffer, [`ErrorKind::InvalidData`] if the
    /// writer reports an impossible byte count, or any error returned by
    /// [`Seek::seek`] on the wrapped writer.
    #[inline(always)]
    fn flush_then_seek(&mut self, position: SeekFrom) -> Result<u64>
    where
        W: Seek,
    {
        self.flush_buffer().and_then(|()| self.inner.seek(position))
    }

    /// Returns pending bytes currently stored in the internal buffer.
    ///
    /// # Returns
    ///
    /// A slice over bytes accepted by this output but not yet written to the
    /// wrapped writer.
    #[inline(always)]
    fn pending_slice(&self) -> &[u8] {
        &self.buffer.data()[self.buffer.position()..self.buffer.limit()]
    }

    /// Writes bytes to the wrapped writer and validates the reported count.
    ///
    /// # Parameters
    ///
    /// * `input` - Source storage.
    /// * `input_index` - Start index inside `input`.
    /// * `count` - Maximum number of bytes to write.
    ///
    /// # Returns
    ///
    /// The number of bytes accepted by the wrapped writer.
    ///
    /// # Errors
    ///
    /// Returns the wrapped writer's I/O error, or [`ErrorKind::InvalidData`]
    /// if it reports a byte count larger than `count`.
    ///
    /// # Safety
    ///
    /// The caller must guarantee that `input_index..input_index + count` is a
    /// valid range inside `input` and that the addition does not overflow.
    #[inline(always)]
    unsafe fn write_inner_unchecked(
        &mut self,
        input: &[u8],
        input_index: usize,
        count: usize,
    ) -> Result<usize> {
        // SAFETY: The caller guarantees the source range is valid.
        let written =
            unsafe { self.inner.write_unchecked(input, input_index, count) }?;
        validate_write_count(written, count)?;
        Ok(written)
    }

    /// Writes all bytes in an indexed source range to the wrapped writer.
    ///
    /// # Parameters
    ///
    /// * `input` - Source storage.
    /// * `input_index` - Start index inside `input`.
    /// * `count` - Number of bytes to write.
    ///
    /// # Errors
    ///
    /// Returns the wrapped writer's I/O error, [`ErrorKind::WriteZero`] if the
    /// writer cannot make progress, or [`ErrorKind::InvalidData`] if it
    /// reports an impossible byte count.
    ///
    /// # Safety
    ///
    /// The caller must guarantee that `input_index..input_index + count` is a
    /// valid range inside `input` and that the addition does not overflow.
    unsafe fn write_all_inner_unchecked(
        &mut self,
        input: &[u8],
        input_index: usize,
        count: usize,
    ) -> Result<()> {
        let mut written = 0;
        while written < count {
            let remaining = count - written;
            // SAFETY: `written < count`, so this suffix remains inside the
            // caller-validated source range.
            match unsafe {
                self.write_inner_unchecked(
                    input,
                    input_index + written,
                    remaining,
                )
            } {
                Ok(0) => {
                    return Err(Error::new(
                        ErrorKind::WriteZero,
                        "failed to write whole buffer",
                    ));
                }
                Ok(count) => written += count,
                Err(error) if error.kind() == ErrorKind::Interrupted => {}
                Err(error) => return Err(error),
            }
        }
        Ok(())
    }
}

impl<W> Write for BufferedByteOutput<W>
where
    W: Write,
{
    /// Writes bytes through the internal buffer.
    #[inline(always)]
    fn write(&mut self, buffer: &[u8]) -> Result<usize> {
        self.write_from(buffer)
    }

    /// Writes all bytes through the internal buffer.
    #[inline(always)]
    fn write_all(&mut self, buffer: &[u8]) -> Result<()> {
        self.write_all_buffered(buffer)
    }

    /// Flushes the internal buffer and then the wrapped writer.
    #[inline(always)]
    fn flush(&mut self) -> Result<()> {
        self.flush_all()
    }
}

impl<W> Seek for BufferedByteOutput<W>
where
    W: Write + Seek,
{
    /// Flushes pending bytes before seeking the wrapped writer.
    #[inline(always)]
    fn seek(&mut self, position: SeekFrom) -> Result<u64> {
        self.flush_then_seek(position)
    }
}

/// Validates a byte count returned by a wrapped writer.
///
/// # Parameters
///
/// * `written` - Byte count reported by the wrapped writer.
/// * `requested` - Maximum byte count requested from the wrapped writer.
///
/// # Errors
///
/// Returns [`ErrorKind::InvalidData`] when the wrapped writer reports more
/// bytes than the source range contained.
#[inline(always)]
fn validate_write_count(written: usize, requested: usize) -> Result<()> {
    if written > requested {
        return Err(Error::new(
            ErrorKind::InvalidData,
            format!(
                "writer reported {written} bytes for a {requested}-byte buffer"
            ),
        ));
    }
    Ok(())
}