product_os_async_executor/

read_write.rs

use alloc::boxed::Box;
use alloc::string::String;
use alloc::string::ToString;
use alloc::vec::Vec;
use bytes::{Buf, BufMut};
/// Taken with modifications from hyper
use core::fmt;
use core::ops::{Deref, DerefMut};
use core::pin::Pin;
use core::task::{ready, Context, Poll};

use core::future::Future;
use core::marker::PhantomPinned;
use core::mem;

use crate::BoxError;

/// A reference to a contiguous slice of memory for vectored I/O reads.
///
/// This is a `no_std`-compatible equivalent of `std::io::IoSliceMut`.
#[derive(Debug)]
pub struct IoSliceMut<'a>(&'a mut [u8]);

impl<'a> IoSliceMut<'a> {
    /// Creates a new `IoSliceMut` wrapping a mutable byte slice.
    pub fn new(buf: &'a mut [u8]) -> Self {
        Self(buf)
    }
}

impl Deref for IoSliceMut<'_> {
    type Target = [u8];
    fn deref(&self) -> &[u8] {
        self.0
    }
}

impl DerefMut for IoSliceMut<'_> {
    fn deref_mut(&mut self) -> &mut [u8] {
        self.0
    }
}

/// A reference to a contiguous slice of memory for vectored I/O writes.
///
/// This is a `no_std`-compatible equivalent of `std::io::IoSlice`.
#[derive(Debug, Clone, Copy)]
pub struct IoSlice<'a>(&'a [u8]);

impl<'a> IoSlice<'a> {
    /// Creates a new `IoSlice` wrapping a byte slice.
    pub const fn new(buf: &'a [u8]) -> Self {
        Self(buf)
    }
}

impl Deref for IoSlice<'_> {
    type Target = [u8];
    fn deref(&self) -> &[u8] {
        self.0
    }
}

/// Trait for async reading operations.
///
/// This trait provides an interface for asynchronously reading bytes from a source.
pub trait AsyncRead {
    /// Attempts to read bytes into a buffer.
    ///
    /// # Errors
    ///
    /// Returns an error if the read operation fails.
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut [u8],
    ) -> Poll<Result<usize, BoxError>>;

    /// Attempts to read bytes into multiple buffers (vectored I/O).
    ///
    /// # Errors
    ///
    /// Returns an error if the read operation fails.
    fn poll_read_vectored(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        bufs: &mut [IoSliceMut<'_>],
    ) -> Poll<Result<usize, BoxError>> {
        for b in bufs {
            if !b.is_empty() {
                return self.poll_read(cx, b);
            }
        }

        self.poll_read(cx, &mut [])
    }
}

/// Trait for async writing operations.
///
/// This trait provides an interface for asynchronously writing bytes to a destination.
pub trait AsyncWrite {
    /// Attempts to write bytes from a buffer.
    ///
    /// # Errors
    ///
    /// Returns an error if the write operation fails.
    fn poll_write(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<Result<usize, BoxError>>;

    /// Flushes any buffered data.
    ///
    /// # Errors
    ///
    /// Returns an error if the flush operation fails.
    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), BoxError>>;

    /// Shuts down the writer.
    ///
    /// # Errors
    ///
    /// Returns an error if the shutdown operation fails.
    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), BoxError>>;

    /// Returns whether vectored writes are supported.
    fn is_write_vectored(&self) -> bool {
        false
    }

    /// Attempts to write bytes from multiple buffers (vectored I/O).
    ///
    /// # Errors
    ///
    /// Returns an error if the write operation fails.
    fn poll_write_vectored(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        bufs: &[IoSlice],
    ) -> Poll<Result<usize, BoxError>> {
        for b in bufs {
            if !b.is_empty() {
                return self.poll_write(cx, b);
            }
        }

        self.poll_write(cx, &[])
    }
}

/// A read buffer that tracks filled and initialized portions.
///
/// This buffer manages a `Vec<u8>` and keeps track of which bytes have been
/// filled with data and which bytes have been initialized.
pub struct ReadBuf {
    raw: Vec<u8>,
    filled: usize,
    init: usize,
}

impl ReadBuf {
    /// Creates a new `ReadBuf` from a mutable slice.
    ///
    /// The data is copied into an owned `Vec<u8>`. All bytes are considered
    /// initialized but unfilled.
    #[inline]
    pub fn new(raw: &mut [u8]) -> Self {
        let len = raw.len();
        Self {
            raw: raw.to_vec(),
            filled: 0,
            init: len,
        }
    }

    /// Creates a new `ReadBuf` with uninitialized bytes.
    ///
    /// All bytes are considered uninitialized and unfilled.
    #[inline]
    pub const fn uninit(raw: Vec<u8>) -> Self {
        Self {
            raw,
            filled: 0,
            init: 0,
        }
    }

    /// Returns a mutable reference to the underlying buffer.
    #[inline]
    pub fn raw(&mut self) -> &mut Vec<u8> {
        &mut self.raw
    }

    /// Returns the filled portion of the buffer as a slice.
    #[inline]
    pub fn filled(&self) -> &[u8] {
        &self.raw[..self.filled]
    }

    /// Returns a cursor over the unfilled portion of the buffer.
    #[inline]
    pub fn unfilled(&mut self) -> ReadBufCursor<'_> {
        ReadBufCursor { buf: &mut *self }
    }

    /// Sets the number of initialized bytes.
    ///
    /// # Safety
    ///
    /// The caller must ensure that at least `n` bytes have been initialized.
    #[inline]
    #[allow(dead_code)]
    pub(crate) unsafe fn set_init(&mut self, n: usize) {
        self.init = self.init.max(n);
    }

    /// Sets the number of filled bytes.
    ///
    /// # Safety
    ///
    /// The caller must ensure that at least `n` bytes have been filled with valid data.
    #[inline]
    #[allow(dead_code)]
    pub(crate) unsafe fn set_filled(&mut self, n: usize) {
        self.filled = self.filled.max(n);
    }

    #[inline]
    #[allow(dead_code)]
    pub(crate) const fn len(&self) -> usize {
        self.filled
    }

    #[inline]
    #[allow(dead_code)]
    pub(crate) const fn init_len(&self) -> usize {
        self.init
    }

    /// Returns the remaining capacity in the buffer.
    #[inline]
    fn remaining(&self) -> usize {
        self.capacity() - self.filled
    }

    /// Returns the total capacity of the buffer.
    #[inline]
    fn capacity(&self) -> usize {
        self.raw.len()
    }
}

impl core::fmt::Debug for ReadBuf {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ReadBuf")
            .field("filled", &self.filled)
            .field("init", &self.init)
            .field("capacity", &self.capacity())
            .field("raw", &self.raw)
            .finish()
    }
}

/// Cursor over the unfilled portion of a read buffer.
///
/// This provides methods to advance and modify the unfilled portion of the buffer.
#[derive(Debug)]
pub struct ReadBufCursor<'a> {
    buf: &'a mut ReadBuf,
}

impl ReadBufCursor<'_> {
    /// Returns a mutable reference to the underlying buffer.
    ///
    /// # Safety
    ///
    /// The caller must ensure that the buffer is properly initialized.
    #[inline]
    pub unsafe fn as_mut(&mut self) -> &mut Vec<u8> {
        &mut self.buf.raw
    }

    /// Returns the filled portion of the buffer.
    #[inline]
    pub fn filled(&self) -> &[u8] {
        self.buf.filled()
    }

    /// Advances the cursor by `n` bytes.
    ///
    /// # Safety
    ///
    /// The caller must ensure that `n` bytes have been initialized in the buffer.
    ///
    /// # Panics
    ///
    /// Panics if advancing by `n` would cause an overflow.
    #[inline]
    pub unsafe fn advance(&mut self, n: usize) {
        self.buf.filled = self.buf.filled.checked_add(n).expect("overflow");
        self.buf.init = self.buf.filled.max(self.buf.init);
    }

    /// Returns the remaining capacity in the buffer.
    #[inline]
    pub fn remaining(&self) -> usize {
        self.buf.remaining()
    }

    /// Copies bytes from the provided buffer into this cursor.
    ///
    /// # Panics
    ///
    /// Panics if the buffer doesn't have enough remaining capacity.
    #[inline]
    pub fn put_slice(&mut self, buf: &[u8]) {
        assert!(
            self.buf.remaining() >= buf.len(),
            "buf.len() must fit in remaining()"
        );

        let amt = buf.len();
        let end = self.buf.filled + amt;

        self.buf.raw[self.buf.filled..end].copy_from_slice(buf);

        if self.buf.init < end {
            self.buf.init = end;
        }
        self.buf.filled = end;
    }
}

macro_rules! deref_async_read {
    () => {
        fn poll_read(
            mut self: Pin<&mut Self>,
            cx: &mut Context<'_>,
            buf: &mut [u8],
        ) -> Poll<Result<usize, BoxError>> {
            Pin::new(&mut **self).poll_read(cx, buf)
        }
    };
}

impl<T: ?Sized + AsyncRead + Unpin> AsyncRead for Box<T> {
    deref_async_read!();
}

impl<T: ?Sized + AsyncRead + Unpin> AsyncRead for &mut T {
    deref_async_read!();
}

impl<P> AsyncRead for Pin<P>
where
    P: DerefMut,
    P::Target: AsyncRead,
{
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut [u8],
    ) -> Poll<Result<usize, BoxError>> {
        pin_as_deref_mut(self).poll_read(cx, buf)
    }
}

macro_rules! deref_async_write {
    () => {
        fn poll_write(
            mut self: Pin<&mut Self>,
            cx: &mut Context<'_>,
            buf: &[u8],
        ) -> Poll<Result<usize, BoxError>> {
            Pin::new(&mut **self).poll_write(cx, buf)
        }

        fn poll_write_vectored(
            mut self: Pin<&mut Self>,
            cx: &mut Context<'_>,
            bufs: &[IoSlice],
        ) -> Poll<Result<usize, BoxError>> {
            Pin::new(&mut **self).poll_write_vectored(cx, bufs)
        }

        fn is_write_vectored(&self) -> bool {
            (**self).is_write_vectored()
        }

        fn poll_flush(
            mut self: Pin<&mut Self>,
            cx: &mut Context<'_>,
        ) -> Poll<Result<(), BoxError>> {
            Pin::new(&mut **self).poll_flush(cx)
        }

        fn poll_shutdown(
            mut self: Pin<&mut Self>,
            cx: &mut Context<'_>,
        ) -> Poll<Result<(), BoxError>> {
            Pin::new(&mut **self).poll_shutdown(cx)
        }
    };
}

impl<T: ?Sized + AsyncWrite + Unpin> AsyncWrite for Box<T> {
    deref_async_write!();
}

impl<T: ?Sized + AsyncWrite + Unpin> AsyncWrite for &mut T {
    deref_async_write!();
}

impl<P> AsyncWrite for Pin<P>
where
    P: DerefMut,
    P::Target: AsyncWrite,
{
    fn poll_write(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<Result<usize, BoxError>> {
        pin_as_deref_mut(self).poll_write(cx, buf)
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), BoxError>> {
        pin_as_deref_mut(self).poll_flush(cx)
    }

    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), BoxError>> {
        pin_as_deref_mut(self).poll_shutdown(cx)
    }

    fn is_write_vectored(&self) -> bool {
        (**self).is_write_vectored()
    }

    fn poll_write_vectored(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        bufs: &[IoSlice],
    ) -> Poll<Result<usize, BoxError>> {
        pin_as_deref_mut(self).poll_write_vectored(cx, bufs)
    }
}

/// Polyfill for `Pin::as_deref_mut()`
/// NOTE: `Pin::as_deref_mut()` is not yet stabilized as of Rust 1.81.
/// This polyfill should be replaced once the method is stabilized.
fn pin_as_deref_mut<P: DerefMut>(pin: Pin<&mut Pin<P>>) -> Pin<&mut P::Target> {
    // SAFETY: we go directly from Pin<&mut Pin<P>> to Pin<&mut P::Target>, without moving or
    // giving out the &mut Pin<P> in the process. See Pin::as_deref_mut() for more detail.
    unsafe { pin.get_unchecked_mut() }.as_mut()
}

/// Polls reading from an `AsyncRead` into a buffer.
///
/// # Errors
///
/// Returns an error if the read operation fails.
pub fn poll_read_buf<T: AsyncRead, B: Buf + BufMut>(
    io: Pin<&mut T>,
    cx: &mut Context<'_>,
    buf: &mut B,
) -> Poll<Result<usize, BoxError>> {
    if !buf.has_remaining_mut() {
        return Poll::Ready(Ok(0));
    }

    let dst = buf.chunk_mut();
    // SAFETY: BufMut guarantees the returned UninitSlice is valid for writes
    let dst_slice = unsafe { core::slice::from_raw_parts_mut(dst.as_mut_ptr(), dst.len()) };

    let n = ready!(io.poll_read(cx, dst_slice)?);

    // SAFETY: poll_read guarantees that n bytes have been written to the buffer
    unsafe {
        buf.advance_mut(n);
    }

    Poll::Ready(Ok(n))
}

/// Polls writing from a buffer to an `AsyncWrite`.
///
/// # Errors
///
/// Returns an error if the write operation fails.
pub fn poll_write_buf<T: AsyncWrite, B: Buf + BufMut>(
    io: Pin<&mut T>,
    cx: &mut Context<'_>,
    buf: &mut B,
) -> Poll<Result<usize, BoxError>> {
    if !buf.has_remaining() {
        return Poll::Ready(Ok(0));
    }

    let n = if io.is_write_vectored() {
        let slice = IoSlice::new(buf.chunk());
        let slices = [slice];
        ready!(io.poll_write_vectored(cx, &slices))?
    } else {
        ready!(io.poll_write(cx, buf.chunk()))?
    };

    buf.advance(n);

    Poll::Ready(Ok(n))
}

/// Extension trait for `AsyncWrite` providing utility methods.
pub trait AsyncWriteExt: AsyncWrite {
    /// Writes all bytes from the source buffer to this writer.
    ///
    /// Calls `poll_write` repeatedly until the entire buffer has been written.
    ///
    /// # Errors
    ///
    /// Returns an error if writing fails or if zero bytes are written (indicating the writer is closed).
    fn write_all<'a>(&'a mut self, src: &'a [u8]) -> WriteAll<'a, Self>
    where
        Self: Unpin,
    {
        write_all(self, src)
    }
}

impl<T: AsyncWrite + ?Sized> AsyncWriteExt for T {}

/// A future that writes all bytes from a buffer to a writer.
///
/// This is created by the [`AsyncWriteExt::write_all`] method.
#[derive(Debug)]
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct WriteAll<'a, W: ?Sized> {
    writer: &'a mut W,
    buf: &'a [u8],
    _pin: PhantomPinned,
}

impl<'a, W: ?Sized> WriteAll<'a, W> {
    #[allow(clippy::mut_mut)]
    fn project(self: Pin<&mut Self>) -> (&mut &'a mut W, &mut &'a [u8]) {
        // SAFETY: `writer` and `buf` are not structurally pinned (they are references).
        // Only `_pin: PhantomPinned` is structurally pinned but we never move it.
        let this = unsafe { self.get_unchecked_mut() };
        (&mut this.writer, &mut this.buf)
    }
}

pub(crate) fn write_all<'a, W>(writer: &'a mut W, buf: &'a [u8]) -> WriteAll<'a, W>
where
    W: AsyncWrite + Unpin + ?Sized,
{
    WriteAll {
        writer,
        buf,
        _pin: PhantomPinned,
    }
}

impl<W> Future for WriteAll<'_, W>
where
    W: AsyncWrite + Unpin + ?Sized,
{
    type Output = Result<(), BoxError>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), BoxError>> {
        let (writer, buf) = self.project();
        while !buf.is_empty() {
            let n = ready!(Pin::new(&mut **writer).poll_write(cx, buf))?;
            {
                let (_, rest) = mem::take(buf).split_at(n);
                *buf = rest;
            }
            if n == 0 {
                return Poll::Ready(Err(Box::new(IoError(String::from("Write Zero")))));
            }
        }

        Poll::Ready(Ok(()))
    }
}

/// An I/O error with a message.
#[derive(Debug)]
pub struct IoError(pub String);

impl IoError {
    #[allow(dead_code)]
    pub(crate) fn new(s: &str) -> Self {
        Self(s.to_string())
    }
}

impl fmt::Display for IoError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "{}", self.0)
    }
}

impl core::error::Error for IoError {}