      1
      2
      3
      4
      5
      6
      7
      8
      9
     10
     11
     12
     13
     14
     15
     16
     17
     18
     19
     20
     21
     22
     23
     24
     25
     26
     27
     28
     29
     30
     31
     32
     33
     34
     35
     36
     37
     38
     39
     40
     41
     42
     43
     44
     45
     46
     47
     48
     49
     50
     51
     52
     53
     54
     55
     56
     57
     58
     59
     60
     61
     62
     63
     64
     65
     66
     67
     68
     69
     70
     71
     72
     73
     74
     75
     76
     77
     78
     79
     80
     81
     82
     83
     84
     85
     86
     87
     88
     89
     90
     91
     92
     93
     94
     95
     96
     97
     98
     99
    100
    101
    102
    103
    104
    105
    106
    107
    108
    109
    110
    111
    112
    113
    114
    115
    116
    117

    use crate::io::AsyncRead;

    use std::io;
    use std::ops::DerefMut;
    use std::pin::Pin;
    use std::task::{Context, Poll};

    /// Reads bytes asynchronously.
    ///
    /// This trait is analogous to [`std::io::BufRead`], but integrates with
    /// the asynchronous task system. In particular, the [`poll_fill_buf`] method,
    /// unlike [`BufRead::fill_buf`], will automatically queue the current task for wakeup
    /// and return if data is not yet available, rather than blocking the calling
    /// thread.
    ///
    /// Utilities for working with `AsyncBufRead` values are provided by
    /// [`AsyncBufReadExt`].
    ///
    /// [`std::io::BufRead`]: std::io::BufRead
    /// [`poll_fill_buf`]: AsyncBufRead::poll_fill_buf
    /// [`BufRead::fill_buf`]: std::io::BufRead::fill_buf
    /// [`AsyncBufReadExt`]: crate::io::AsyncBufReadExt
    pub trait AsyncBufRead: AsyncRead {
        /// Attempts to return the contents of the internal buffer, filling it with more data
        /// from the inner reader if it is empty.
        ///
        /// On success, returns `Poll::Ready(Ok(buf))`.
        ///
        /// If no data is available for reading, the method returns
        /// `Poll::Pending` and arranges for the current task (via
        /// `cx.waker().wake_by_ref()`) to receive a notification when the object becomes
        /// readable or is closed.
        ///
        /// This function is a lower-level call. It needs to be paired with the
        /// [`consume`] method to function properly. When calling this
        /// method, none of the contents will be "read" in the sense that later
        /// calling [`poll_read`] may return the same contents. As such, [`consume`] must
        /// be called with the number of bytes that are consumed from this buffer to
        /// ensure that the bytes are never returned twice.
        ///
        /// An empty buffer returned indicates that the stream has reached EOF.
        ///
        /// [`poll_read`]: AsyncRead::poll_read
        /// [`consume`]: AsyncBufRead::consume
        fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>>;

        /// Tells this buffer that `amt` bytes have been consumed from the buffer,
        /// so they should no longer be returned in calls to [`poll_read`].
        ///
        /// This function is a lower-level call. It needs to be paired with the
        /// [`poll_fill_buf`] method to function properly. This function does
        /// not perform any I/O, it simply informs this object that some amount of
        /// its buffer, returned from [`poll_fill_buf`], has been consumed and should
        /// no longer be returned. As such, this function may do odd things if
        /// [`poll_fill_buf`] isn't called before calling it.
        ///
        /// The `amt` must be `<=` the number of bytes in the buffer returned by
        /// [`poll_fill_buf`].
        ///
        /// [`poll_read`]: AsyncRead::poll_read
        /// [`poll_fill_buf`]: AsyncBufRead::poll_fill_buf
        fn consume(self: Pin<&mut Self>, amt: usize);
    }

    macro_rules! deref_async_buf_read {
        () => {
            fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
                Pin::new(&mut **self.get_mut()).poll_fill_buf(cx)
            }

            fn consume(mut self: Pin<&mut Self>, amt: usize) {
                Pin::new(&mut **self).consume(amt)
            }
        };
    }

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

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

    impl<P> AsyncBufRead for Pin<P>
    where
        P: DerefMut + Unpin,
        P::Target: AsyncBufRead,
    {
        fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
            self.get_mut().as_mut().poll_fill_buf(cx)
        }

        fn consume(self: Pin<&mut Self>, amt: usize) {
            self.get_mut().as_mut().consume(amt)
        }
    }

    impl AsyncBufRead for &[u8] {
        fn poll_fill_buf(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
            Poll::Ready(Ok(*self))
        }

        fn consume(mut self: Pin<&mut Self>, amt: usize) {
            *self = &self[amt..];
        }
    }

    impl<T: AsRef<[u8]> + Unpin> AsyncBufRead for io::Cursor<T> {
        fn poll_fill_buf(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
            Poll::Ready(io::BufRead::fill_buf(self.get_mut()))
        }

        fn consume(self: Pin<&mut Self>, amt: usize) {
            io::BufRead::consume(self.get_mut(), amt)
        }
    }