volo 0.12.4

//! These codes are copied from `tokio/io/util/buf_reader.rs`.
//! There's some modify of this code, such as support compact.

use std::{
    cmp, fmt, io,
    pin::Pin,
    task::{Context, Poll},
};

use pin_project::pin_project;
use tokio::io::{AsyncBufRead, AsyncRead, AsyncReadExt, ReadBuf};

use crate::net::ext::AsyncExt;

// used by `BufReader` and `BufWriter`
// https://github.com/rust-lang/rust/blob/master/library/std/src/sys_common/io.rs#L1
const DEFAULT_BUF_SIZE: usize = 8 * 1024;

macro_rules! ready {
    ($e:expr $(,)?) => {
        match $e {
            std::task::Poll::Ready(t) => t,
            std::task::Poll::Pending => return std::task::Poll::Pending,
        }
    };
}

impl<R: AsyncRead + Unpin> BufReader<R> {
    pub async fn fill_buf_at_least(&mut self, len: usize) -> io::Result<&[u8]> {
        if len == 0 {
            return Ok(&[]);
        }

        if (self.len - self.pos) >= len {
            return Ok(&self.buf[self.pos..self.len]);
        }

        if len > self.cap {
            // `fill_buf_at_least` is also exposed to custom codecs, whose probe size may be larger
            // than this reader's initial capacity. Grow on demand instead of panicking, while
            // preserving bytes that have already been read but not consumed.
            self.compact();
            let mut new_cap = self.cap.max(1);
            while new_cap < len {
                new_cap = new_cap.checked_mul(2).unwrap_or(len);
            }
            let mut buf = vec![0; new_cap].into_boxed_slice();
            buf[..self.len].copy_from_slice(&self.buf[..self.len]);
            self.buf = buf;
            self.cap = new_cap;
        } else if len > (self.cap - self.pos) {
            // There is enough total capacity, but the unread bytes need to move to the front before
            // the requested amount can become contiguous.
            self.compact();
        }

        while (self.len - self.pos) < len {
            let buf = &mut self.buf[self.len..self.cap];
            let size = self.inner.read(buf).await?;
            if size == 0 {
                return Err(io::Error::new(io::ErrorKind::InvalidData, "invalid eof"));
            }
            self.len += size;
        }
        Ok(&self.buf[self.pos..self.len])
    }
}

/// The `BufReader` struct adds buffering to any reader.
///
/// It can be excessively inefficient to work directly with a [`AsyncRead`]
/// instance. A `BufReader` performs large, infrequent reads on the underlying
/// [`AsyncRead`] and maintains an in-memory buffer of the results.
///
/// `BufReader` can improve the speed of programs that make *small* and
/// *repeated* read calls to the same file or network socket. It does not
/// help when reading very large amounts at once, or reading just one or a few
/// times. It also provides no advantage when reading from a source that is
/// already in memory, like a `Vec<u8>`.
///
/// When the `BufReader` is dropped, the contents of its buffer will be
/// discarded. Creating multiple instances of a `BufReader` on the same
/// stream can cause data loss.
#[pin_project]
pub struct BufReader<R> {
    #[pin]
    pub(super) inner: R,
    pub(super) buf: Box<[u8]>,
    pub(super) pos: usize,
    pub(super) len: usize, // the current valid index
    pub(super) cap: usize,
}

impl<R: AsyncRead> BufReader<R> {
    /// Creates a new `BufReader` with a default buffer capacity. The default is currently 8 KB,
    /// but may change in the future.
    pub fn new(inner: R) -> Self {
        Self::with_capacity(DEFAULT_BUF_SIZE, inner)
    }

    /// Creates a new `BufReader` with the specified buffer capacity.
    pub fn with_capacity(capacity: usize, inner: R) -> Self {
        let buffer = vec![0; capacity];
        Self {
            inner,
            buf: buffer.into_boxed_slice(),
            pos: 0,
            len: 0,
            cap: capacity,
        }
    }

    /// Gets a reference to the underlying reader.
    ///
    /// It is inadvisable to directly read from the underlying reader.
    pub fn get_ref(&self) -> &R {
        &self.inner
    }

    /// Gets a mutable reference to the underlying reader.
    ///
    /// It is inadvisable to directly read from the underlying reader.
    pub fn get_mut(&mut self) -> &mut R {
        &mut self.inner
    }

    /// Gets a pinned mutable reference to the underlying reader.
    ///
    /// It is inadvisable to directly read from the underlying reader.
    pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut R> {
        self.project().inner
    }

    /// Consumes this `BufReader`, returning the underlying reader.
    ///
    /// Note that any leftover data in the internal buffer is lost.
    pub fn into_inner(self) -> R {
        self.inner
    }

    /// Returns a reference to the internally buffered data.
    ///
    /// Unlike `fill_buf`, this will not attempt to fill the buffer if it is empty.
    pub fn buffer(&self) -> &[u8] {
        &self.buf[self.pos..self.len]
    }

    /// 整理 buffer，移动到最前。
    pub fn compact(&mut self) {
        if self.len == self.pos {
            self.pos = 0;
            self.len = 0;
            return;
        }

        let len = self.len - self.pos;
        let dst = self.buf.as_mut_ptr();
        let src = unsafe { dst.add(self.pos) } as *const u8;

        unsafe {
            std::ptr::copy(src, dst, len);
        }

        self.pos = 0;
        self.len = len;
    }

    pub fn clear(&mut self) {
        self.pos = 0;
        self.len = 0;
    }

    /// Invalidates all data in the internal buffer.
    #[inline]
    fn discard_buffer(self: Pin<&mut Self>) {
        let me = self.project();
        *me.pos = 0;
        *me.len = 0;
    }
}

impl<R: AsyncRead> AsyncRead for BufReader<R> {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        // If we don't have any buffered data and we're doing a massive read
        // (larger than our internal buffer), bypass our internal buffer
        // entirely.
        if self.pos == self.len && buf.remaining() >= self.buf.len() {
            let res = ready!(self.as_mut().get_pin_mut().poll_read(cx, buf));
            self.discard_buffer();
            return Poll::Ready(res);
        }
        let rem = ready!(self.as_mut().poll_fill_buf(cx))?;
        let amt = cmp::min(rem.len(), buf.remaining());
        buf.put_slice(&rem[..amt]);
        self.consume(amt);
        Poll::Ready(Ok(()))
    }
}

impl<R: AsyncRead> AsyncBufRead for BufReader<R> {
    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
        let me = self.project();

        // Always return buffered data before polling the underlying reader. In particular, an
        // EOF or error from the underlying reader must not hide data that has already been read.
        if *me.pos < *me.len {
            return Poll::Ready(Ok(&me.buf[*me.pos..*me.len]));
        }

        // If we've reached the end of our internal buffer then we need to fetch
        // some more data from the underlying reader.
        // Branch using `>=` instead of the more correct `==`
        // to tell the compiler that the pos..cap slice is always valid.
        if *me.pos >= *me.cap {
            debug_assert!(*me.pos == *me.cap);
            let mut buf = ReadBuf::new(me.buf);
            ready!(me.inner.poll_read(cx, &mut buf))?;
            *me.len = buf.filled().len();
            *me.pos = 0;
        } else if *me.len < *me.cap {
            let mut buf = ReadBuf::new(&mut me.buf[*me.len..*me.cap]);
            ready!(me.inner.poll_read(cx, &mut buf))?;
            *me.len += buf.filled().len();
        }
        Poll::Ready(Ok(&me.buf[*me.pos..*me.len]))
    }

    fn consume(self: Pin<&mut Self>, amt: usize) {
        let me = self.project();
        *me.pos = cmp::min(*me.pos + amt, *me.len);
    }
}

impl<R: AsyncExt + Send + Sync> AsyncExt for BufReader<R> {
    async fn ready(&self, interest: tokio::io::Interest) -> io::Result<tokio::io::Ready> {
        self.inner.ready(interest).await
    }

    #[cfg(feature = "shmipc")]
    fn is_shmipc(&self) -> bool {
        self.inner.is_shmipc()
    }

    #[cfg(feature = "shmipc")]
    fn shmipc_helper(&self) -> crate::net::shmipc::ShmipcHelper {
        self.inner.shmipc_helper()
    }
}

impl<R: fmt::Debug> fmt::Debug for BufReader<R> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("BufReader")
            .field("reader", &self.inner)
            .field(
                "buffer",
                &format_args!("{}/{}", self.len - self.pos, self.buf.len()),
            )
            .finish()
    }
}

#[cfg(test)]
mod tests {
    use tokio::io::AsyncBufReadExt;

    use super::*;

    struct DataThenError {
        data: Option<Vec<u8>>,
        reads: usize,
    }

    impl AsyncRead for DataThenError {
        fn poll_read(
            self: Pin<&mut Self>,
            _cx: &mut Context<'_>,
            buf: &mut ReadBuf<'_>,
        ) -> Poll<io::Result<()>> {
            let this = self.get_mut();
            this.reads += 1;

            if let Some(data) = this.data.take() {
                assert!(buf.remaining() >= data.len());
                buf.put_slice(&data);
                Poll::Ready(Ok(()))
            } else {
                Poll::Ready(Err(io::Error::new(
                    io::ErrorKind::UnexpectedEof,
                    "stream closed",
                )))
            }
        }
    }

    #[cfg(test)]
    fn is_unpin<T: Unpin>() {}

    #[test]
    fn assert_unpin() {
        is_unpin::<BufReader<()>>();
    }

    #[test]
    fn test_compact() {
        let mut v = vec![0; 16];
        v[0] = 1;
        v[1] = 2;
        v[2] = 3;
        v[3] = 4;
        v[4] = 5;
        let mut buf = BufReader {
            inner: tokio::io::empty(),
            buf: v.into_boxed_slice(),
            pos: 0,
            len: 5,
            cap: 16,
        };
        buf.compact();
        assert_eq!(buf.pos, 0);
        assert_eq!(buf.len, 5);
        assert_eq!(buf.buf[buf.pos..buf.len], [1, 2, 3, 4, 5]);
        let mut v = vec![0; 16];
        v[0] = 1;
        v[1] = 2;
        v[2] = 3;
        v[3] = 4;
        v[4] = 5;
        v[5] = 6;
        v[6] = 7;
        v[7] = 8;
        v[8] = 9;
        v[9] = 10;
        let mut buf = BufReader {
            inner: tokio::io::empty(),
            buf: v.into_boxed_slice(),
            pos: 5,
            len: 10,
            cap: 16,
        };
        buf.compact();
        assert_eq!(buf.pos, 0);
        assert_eq!(buf.len, 5);
        assert_eq!(buf.buf[buf.pos..buf.len], [6, 7, 8, 9, 10]);
    }

    #[tokio::test]
    async fn fill_buf_at_least_grows_and_preserves_unread_data() {
        let data = b"abcdefgh";
        let mut reader = BufReader::with_capacity(4, &data[..]);

        assert_eq!(reader.fill_buf_at_least(4).await.unwrap(), b"abcd");
        reader.consume(2);

        assert_eq!(reader.fill_buf_at_least(6).await.unwrap(), b"cdefgh");
        assert_eq!(reader.cap, 8);
        assert_eq!(reader.pos, 0);
        assert_eq!(reader.len, 6);
    }

    #[tokio::test]
    async fn test_buffered_ttheader_data_is_delivered_before_close_error() {
        // Simulate a complete 1123-byte fallback TTHeader event immediately followed by close.
        const FRAME_LEN: usize = 1123;
        const PAYLOAD_LEN: usize = FRAME_LEN - 4;

        let mut frame = vec![0x5a; FRAME_LEN];
        frame[..4].copy_from_slice(&(PAYLOAD_LEN as u32).to_be_bytes());
        frame[4..6].copy_from_slice(&[0x10, 0x00]);
        let expected_payload = frame[4..].to_vec();

        let inner = DataThenError {
            data: Some(frame),
            reads: 0,
        };
        let mut reader = BufReader::new(inner);

        let header = reader.fill_buf_at_least(6).await.unwrap();
        assert_eq!(
            u32::from_be_bytes(header[..4].try_into().unwrap()) as usize,
            PAYLOAD_LEN
        );
        assert_eq!(&header[4..6], &[0x10, 0x00]);

        reader.consume(4);
        let mut payload = vec![0; PAYLOAD_LEN];
        reader.read_exact(&mut payload).await.unwrap();

        assert_eq!(payload, expected_payload);
        assert_eq!(reader.get_ref().reads, 1);

        let err = reader.read_u8().await.unwrap_err();
        assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof);
        assert_eq!(reader.get_ref().reads, 2);
    }
}