rama-core 0.3.0

use std::{
    fmt,
    io::{IoSlice, Read, Write},
    pin::Pin,
    task::{Context, Poll, ready},
};

use crate::extensions::{Extensions, ExtensionsRef};
use pin_project_lite::pin_project;
use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf};

pin_project! {
    /// a stream which has some data prefixed
    /// to be read first prior to any other reading.
    ///
    /// The source of that prefix data is often the result
    /// of data which was "peeked" from the inner I/O,
    /// although that is not required.
    ///
    /// It's similar to `ChainReader`, except that writing is also
    /// supported and happening directly in function of the inner stream.
    #[derive(Debug, Clone)]
    pub struct PrefixedIo<P, S> {
        prefix_eof: bool,
        #[pin]
        prefix: P,
        #[pin]
        inner: S,
    }
}

impl<P, S> PrefixedIo<P, S> {
    /// Create a new [`PrefixedIo`] for the given prefix
    /// [`AsyncRead`] and inner [`Io`] which implements [`ExtensionsRef`].
    ///
    /// [`Io`]: super::Io
    pub fn new(prefix: P, inner: S) -> Self {
        Self {
            prefix_eof: false,
            prefix,
            inner,
        }
    }
}

impl<P, S: ExtensionsRef> ExtensionsRef for PrefixedIo<P, S> {
    fn extensions(&self) -> &Extensions {
        self.inner.extensions()
    }
}

#[warn(clippy::missing_trait_methods)]
impl<P, S> AsyncRead for PrefixedIo<P, S>
where
    P: AsyncRead,
    S: AsyncRead,
{
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<std::io::Result<()>> {
        let me = self.project();

        if !*me.prefix_eof {
            let rem = buf.remaining();
            ready!(me.prefix.poll_read(cx, buf))?;
            if buf.remaining() == rem {
                *me.prefix_eof = true;
            } else {
                return Poll::Ready(Ok(()));
            }
        }
        me.inner.poll_read(cx, buf)
    }
}

#[warn(clippy::missing_trait_methods)]
impl<P, S> AsyncBufRead for PrefixedIo<P, S>
where
    P: AsyncBufRead,
    S: AsyncBufRead,
{
    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<&[u8]>> {
        let me = self.project();

        if !*me.prefix_eof {
            match ready!(me.prefix.poll_fill_buf(cx)?) {
                [] => {
                    *me.prefix_eof = true;
                }
                buf => return Poll::Ready(Ok(buf)),
            }
        }
        me.inner.poll_fill_buf(cx)
    }

    fn consume(self: Pin<&mut Self>, amt: usize) {
        let me = self.project();
        if !*me.prefix_eof {
            me.prefix.consume(amt)
        } else {
            me.inner.consume(amt)
        }
    }
}

impl<P, S> Read for PrefixedIo<P, S>
where
    P: Read,
    S: Read,
{
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        if !self.prefix_eof {
            let n = self.prefix.read(buf)?;
            if n == 0 {
                self.prefix_eof = true;
            } else {
                return Ok(n);
            }
        }
        self.inner.read(buf)
    }
}

#[warn(clippy::missing_trait_methods)]
impl<P, S> AsyncWrite for PrefixedIo<P, S>
where
    S: AsyncWrite,
{
    fn poll_write(
        self: Pin<&mut Self>,
        cx: &mut core::task::Context<'_>,
        buf: &[u8],
    ) -> Poll<std::io::Result<usize>> {
        let me = self.project();
        me.inner.poll_write(cx, buf)
    }

    fn poll_flush(
        self: Pin<&mut Self>,
        cx: &mut core::task::Context<'_>,
    ) -> Poll<std::io::Result<()>> {
        let me = self.project();
        me.inner.poll_flush(cx)
    }

    fn poll_shutdown(
        self: Pin<&mut Self>,
        cx: &mut core::task::Context<'_>,
    ) -> Poll<std::io::Result<()>> {
        let me = self.project();
        me.inner.poll_shutdown(cx)
    }

    fn poll_write_vectored(
        self: Pin<&mut Self>,
        cx: &mut core::task::Context<'_>,
        bufs: &[IoSlice<'_>],
    ) -> Poll<Result<usize, std::io::Error>> {
        let me = self.project();
        me.inner.poll_write_vectored(cx, bufs)
    }

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

impl<P, S> Write for PrefixedIo<P, S>
where
    S: Write,
{
    #[inline]
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        self.inner.write(buf)
    }

    #[inline]
    fn flush(&mut self) -> std::io::Result<()> {
        self.inner.flush()
    }

    #[inline]
    fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
        self.inner.write_all(buf)
    }

    #[inline]
    fn write_fmt(&mut self, args: fmt::Arguments<'_>) -> std::io::Result<()> {
        self.inner.write_fmt(args)
    }

    #[inline]
    fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> std::io::Result<usize> {
        self.inner.write_vectored(bufs)
    }
}

#[cfg(test)]
mod tests {
    use crate::ServiceInput;

    use super::*;

    use std::io::Cursor;
    use tokio::io::{AsyncReadExt, AsyncWriteExt};

    async fn test_multi_read_async<const N: usize>(
        mut stream: impl AsyncRead + Unpin,
        cases: &[&str],
    ) {
        let mut buf = [0u8; N];

        for (i, case) in cases.iter().enumerate() {
            let n = stream.read(&mut buf).await.unwrap();
            assert_eq!(
                n,
                case.len(),
                "[{N}][async] step #{} for cases: {:?}",
                i + 1,
                cases
            );
            assert_eq!(
                &buf[..n],
                case.as_bytes(),
                "[{N}][async] step #{} for cases: {:?}",
                i + 1,
                cases
            );
        }
    }

    fn test_multi_read_sync<const N: usize>(mut stream: impl Read, cases: &[&str]) {
        let mut buf = [0u8; N];

        for (i, case) in cases.iter().enumerate() {
            let n = stream.read(&mut buf).unwrap();
            assert_eq!(
                n,
                case.len(),
                "[{N}][sync] step #{} for cases: {:?}",
                i + 1,
                cases
            );
            assert_eq!(
                &buf[..n],
                case.as_bytes(),
                "[{N}][sync] step #{} for cases: {:?}",
                i + 1,
                cases
            );
        }
    }

    #[tokio::test]
    async fn test_prefix_stream_read() {
        #[derive(Debug)]
        struct TestCase<const N: usize> {
            prefix_data: &'static str,
            inner_data: &'static str,
            expected_reads: &'static [&'static str],
        }

        impl<const N: usize> TestCase<N> {
            async fn test_sync_and_async(&self) {
                let new_stream = || {
                    let prefix_data = Cursor::new(self.prefix_data);
                    let inner_data = Cursor::new(self.inner_data);
                    PrefixedIo::new(prefix_data, ServiceInput::new(inner_data))
                };

                test_multi_read_async::<N>(&mut new_stream(), self.expected_reads).await;
                test_multi_read_sync::<N>(&mut new_stream(), self.expected_reads);
            }
        }

        TestCase::<10> {
            prefix_data: "hello",
            inner_data: " world",
            expected_reads: &["hello", " world", ""],
        }
        .test_sync_and_async()
        .await;

        TestCase::<5> {
            prefix_data: "hello world",
            inner_data: "next data",
            expected_reads: &["hello", " worl", "d", "next ", "data", ""],
        }
        .test_sync_and_async()
        .await;

        TestCase::<2> {
            prefix_data: "peek",
            inner_data: "inner",
            expected_reads: &["pe", "ek", "in", "ne", "r", ""],
        }
        .test_sync_and_async()
        .await;

        TestCase::<8> {
            prefix_data: "",
            inner_data: "inner data",
            expected_reads: &["inner da", "ta", ""],
        }
        .test_sync_and_async()
        .await;

        TestCase::<10> {
            prefix_data: "",
            inner_data: "inner data",
            expected_reads: &["inner data", ""],
        }
        .test_sync_and_async()
        .await;

        TestCase::<12> {
            prefix_data: "",
            inner_data: "inner data",
            expected_reads: &["inner data", ""],
        }
        .test_sync_and_async()
        .await;
    }

    fn new_prefix_write_stream() -> PrefixedIo<Cursor<Vec<u8>>, ServiceInput<Cursor<Vec<u8>>>> {
        let prefix_data = Cursor::new(Vec::new());
        let inner_data = Cursor::new(Vec::new());
        PrefixedIo::new(prefix_data, ServiceInput::new(inner_data))
    }

    async fn test_multi_write_async(mut stream: impl AsyncWrite + Unpin, cases: &[&str]) {
        for case in cases {
            stream.write_all(case.as_bytes()).await.unwrap();
        }
    }

    fn test_multi_write_sync(mut stream: impl Write, cases: &[&str]) {
        for case in cases {
            stream.write_all(case.as_bytes()).unwrap();
        }
    }

    #[tokio::test]
    async fn test_prefix_stream_write() {
        #[derive(Debug)]
        struct TestCase<'a> {
            writes: &'a [&'static str],
        }

        impl TestCase<'_> {
            async fn test_sync_and_async(&self) {
                let mut stream = new_prefix_write_stream();
                test_multi_write_async(&mut stream, self.writes).await;

                assert!(!stream.prefix_eof, "[async] writes: {:?}", self.writes);
                assert_eq!(
                    stream.prefix.position(),
                    0,
                    "[async] writes: {:?}",
                    self.writes
                );
                assert!(
                    stream.prefix.into_inner().is_empty(),
                    "[async] writes: {:?}",
                    self.writes
                );

                assert_eq!(
                    self.writes.join(""),
                    String::from_utf8(stream.inner.input.into_inner()).unwrap(),
                    "[async] writes: {:?}",
                    self.writes,
                );

                let mut stream = new_prefix_write_stream();
                test_multi_write_sync(&mut stream, self.writes);

                assert!(!stream.prefix_eof, "[sync] writes: {:?}", self.writes);
                assert_eq!(
                    stream.prefix.position(),
                    0,
                    "[sync] writes: {:?}",
                    self.writes
                );
                assert!(
                    stream.prefix.into_inner().is_empty(),
                    "[sync] writes: {:?}",
                    self.writes,
                );

                assert_eq!(
                    self.writes.join(""),
                    String::from_utf8(stream.inner.input.into_inner()).unwrap(),
                    "[sync] writes: {:?}",
                    self.writes
                );
            }
        }

        for writes in [
            vec![],
            vec![""],
            vec!["test", " ", "data"],
            vec!["test data"],
        ] {
            TestCase {
                writes: writes.as_slice(),
            }
            .test_sync_and_async()
            .await;
        }
    }
}