io: add fill_buf and consume

tokio/src/io/util/async_buf_read_ext.rs

@@ -1,3 +1,4 @@
+use crate::io::util::fill_buf::{fill_buf, FillBuf};
 use crate::io::util::lines::{lines, Lines};
 use crate::io::util::read_line::{read_line, ReadLine};
 use crate::io::util::read_until::{read_until, ReadUntil};
@@ -206,6 +207,59 @@ cfg_io_util! {
             split(self, byte)
         }
 
+        /// Returns the contents of the internal buffer, filling it with more
+        /// data from the inner reader if it is empty.
+        ///
+        /// 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
+        /// `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.
+        ///
+        /// Equivalent to:
+        ///
+        /// ```ignore
+        /// async fn fill_buf(&mut self) -> io::Result<&[u8]>;
+        /// ```
+        ///
+        /// # Errors
+        ///
+        /// This function will return an I/O error if the underlying reader was
+        /// read, but returned an error.
+        ///
+        /// [`consume`]: crate::io::AsyncBufReadExt::consume
+        fn fill_buf<'a>(&'a mut self) -> FillBuf<'a, Self>
+        where
+            Self: Unpin,
+        {
+            fill_buf(self)
+        }
+
+        /// Tells this buffer that `amt` bytes have been consumed from the
+        /// buffer, so they should no longer be returned in calls to [`read`].
+        ///
+        /// This function is a lower-level call. It needs to be paired with the
+        /// [`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 [`fill_buf`], has been consumed and should
+        /// no longer be returned. As such, this function may do odd things if
+        /// [`fill_buf`] isn't called before calling it.
+        ///
+        /// The `amt` must be less than the number of bytes in the buffer
+        /// returned by [`fill_buf`].
+        ///
+        /// [`read`]: crate::io::AsyncReadExt::read
+        /// [`fill_buf`]: crate::io::AsyncBufReadExt::fill_buf
+        fn consume(&mut self, amt: usize)
+        where
+            Self: Unpin,
+        {
+            std::pin::Pin::new(self).consume(amt)
+        }
+
         /// Returns a stream over the lines of this reader.
         /// This method is the async equivalent to [`BufRead::lines`](std::io::BufRead::lines).
         ///

tokio/src/io/util/fill_buf.rs

@@ -0,0 +1,49 @@
+use crate::io::AsyncBufRead;
+
+use pin_project_lite::pin_project;
+use std::future::Future;
+use std::io;
+use std::marker::PhantomPinned;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+pin_project! {
+    /// Future for the [`fill_buf`](crate::io::AsyncBufReadExt::fill_buf) method.
+    #[derive(Debug)]
+    #[must_use = "futures do nothing unless you `.await` or poll them"]
+    pub struct FillBuf<'a, R: ?Sized> {
+        reader: Option<&'a mut R>,
+        #[pin]
+        _pin: PhantomPinned,
+    }
+}
+
+pub(crate) fn fill_buf<'a, R>(reader: &'a mut R) -> FillBuf<'a, R>
+where
+    R: AsyncBufRead + ?Sized + Unpin,
+{
+    FillBuf {
+        reader: Some(reader),
+        _pin: PhantomPinned,
+    }
+}
+
+impl<'a, R: AsyncBufRead + ?Sized + Unpin> Future for FillBuf<'a, R> {
+    type Output = io::Result<&'a [u8]>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let me = self.project();
+
+        let reader = me.reader.take().expect("Polled after completion.");
+        match Pin::new(&mut *reader).poll_fill_buf(cx) {
+            Poll::Ready(_) => match Pin::new(reader).poll_fill_buf(cx) {
+                Poll::Ready(slice) => Poll::Ready(slice),
+                Poll::Pending => panic!("poll_fill_buf returned Pending while having data"),
+            },
+            Poll::Pending => {
+                *me.reader = Some(reader);
+                Poll::Pending
+            }
+        }
+    }
+}

tokio/src/io/util/mod.rs

@@ -49,6 +49,7 @@ cfg_io_util! {
     mod read_exact;
     mod read_int;
     mod read_line;
+    mod fill_buf;
 
     mod read_to_end;
     mod vec_with_initialized;