product-os-async-executor 0.0.18

use std::prelude::v1::*;

/// Taken with modifications from hyper
use std::fmt;
use std::ops::DerefMut;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::boxed::Box;
use std::vec::Vec;
use bytes::{Buf, BufMut};
use futures_core::ready;

use pin_project_lite::pin_project;
use std::future::Future;
use std::marker::PhantomPinned;
use std::mem;

use crate::BoxError;

pub use ioslice::{IoSlice, IoSliceMut};


/// 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.
    ///
    /// All bytes are considered initialized but unfilled.
    #[inline]
    pub fn new(raw: &mut [u8]) -> Self {
        let len = raw.len();
        Self {
            // SAFETY: We never de-init the bytes ourselves.
            raw: 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> {
        // SAFETY: We only slice the filled part of the buffer, which is always valid
        &mut self.raw
    }

    /// Returns the filled portion of the buffer as a slice.
    #[inline]
    pub fn filled(&self) -> &[u8] {
        // SAFETY: We only slice the filled part of the buffer, which is always valid
        self.raw.as_slice()
    }

    /// Returns a cursor over the unfilled portion of the buffer.
    #[inline]
    pub fn unfilled(&mut self) -> ReadBufCursor<'_> {
        ReadBufCursor {
            // SAFETY: self.buf is never re-assigned, so it's safe to narrow
            // the lifetime.
            buf: unsafe {
                mem::transmute::<&mut Self, &mut Self>(
                    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
    ///
    /// This function returns a mutable reference to the underlying buffer.
    /// 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();
        // Cannot overflow, asserted above
        let end = self.buf.filled + amt;

        // Safety: the length is asserted above
        unsafe {
            self.buf.raw[self.buf.filled..end]
                .as_mut_ptr()
                .cast::<u8>()
                .copy_from_nonoverlapping(buf.as_ptr(), amt);
        }

        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()`
/// TODO: use `Pin::as_deref_mut()` instead once 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()
}

/*
pub trait IoSlice: DerefMut + Send + Sync {
    fn is_empty(&self) -> bool;
    fn to_slice(&self) -> &[u8];
}
*/


/// Polls reading from an `AsyncRead` into a buffer.
///
/// This function attempts to read data from the async reader into the provided buffer.
///
/// # Errors
///
/// Returns an error if the read operation fails.
///
/// # Panics
///
/// May panic if the buffer pointer changes unexpectedly (internal consistency check).
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 n = {
        let mut dst = buf.chunk().to_vec();

        ready!(io.poll_read(cx, &mut dst)?);

        let buf = &mut ReadBuf::uninit(dst);
        let ptr = buf.filled().as_ptr();

        // Ensure the pointer does not change from under us
        assert_eq!(ptr, buf.filled().as_ptr());
        buf.filled().len()
    };

    // Safety: This is guaranteed to be the number of initialized (and read)
    // bytes due to the invariants provided by `ReadBuf::filled`.
    unsafe {
        buf.advance_mut(n);
    }

    Poll::Ready(Ok(n))
}


/// Polls writing from a buffer to an `AsyncWrite`.
///
/// This function attempts to write data from the provided buffer to the async writer.
///
/// # 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>> {
    const MAX_BUFS: usize = 64;

    if !buf.has_remaining() {
        return Poll::Ready(Ok(0));
    }

    let n = if io.is_write_vectored() {
        let mut slices = Vec::new();
        for _ in 0..MAX_BUFS {
            slices.push(IoSlice::new(&[]));
        }

        let cnt = {
            if slices.is_empty() {
                0
            }
            else if buf.has_remaining() {
                slices[0] = IoSlice::new(buf.chunk());
                1
            } else {
                0
            }
        };

        ready!(io.poll_write_vectored(cx, &slices[..cnt]))?
    } 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.
    ///
    /// This method will continuously call `poll_write` 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)
    }
}




pin_project! {
    /// 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],
        // Make this future `!Unpin` for compatibility with async trait methods.
        #[pin]
        _pin: PhantomPinned,
    }
}

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 me = self.project();
        while !me.buf.is_empty() {
            let n = ready!(Pin::new(&mut *me.writer).poll_write(cx, me.buf))?;
            {
                let (_, rest) = mem::take(&mut *me.buf).split_at(n);
                *me.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 {}