From 06f2569d23732f2bb48427770ae6ef1861e460cc Mon Sep 17 00:00:00 2001
From: Wonwoo Choi <chwo9843@gmail.com>
Date: Thu, 12 Sep 2019 00:02:57 +0900
Subject: [PATCH] Add BufRead::fill_buf (#176)

* Add BufRead::fill_buf

* Make FillBufFuture constructor pub(crate)

* Give more information about the transmutation source type
---
 src/io/buf_read/fill_buf.rs | 32 ++++++++++++++++++++++++++++++++
 src/io/buf_read/mod.rs      | 22 ++++++++++++++++++++++
 2 files changed, 54 insertions(+)
 create mode 100644 src/io/buf_read/fill_buf.rs

diff --git a/src/io/buf_read/fill_buf.rs b/src/io/buf_read/fill_buf.rs
new file mode 100644
index 00000000..0ce58cfb
--- /dev/null
+++ b/src/io/buf_read/fill_buf.rs
@@ -0,0 +1,32 @@
+use std::pin::Pin;
+
+use futures_io::AsyncBufRead;
+
+use crate::future::Future;
+use crate::io;
+use crate::task::{Context, Poll};
+
+#[doc(hidden)]
+#[allow(missing_debug_implementations)]
+pub struct FillBufFuture<'a, R: ?Sized> {
+    reader: &'a mut R,
+}
+
+impl<'a, R: ?Sized> FillBufFuture<'a, R> {
+    pub(crate) fn new(reader: &'a mut R) -> Self {
+        Self { reader }
+    }
+}
+
+impl<'a, R: AsyncBufRead + Unpin + ?Sized> Future for FillBufFuture<'a, R> {
+    type Output = io::Result<&'a [u8]>;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&'a [u8]>> {
+        let Self { reader } = &mut *self;
+        let result = Pin::new(reader).poll_fill_buf(cx);
+        // This is safe because:
+        // 1. The buffer is valid for the lifetime of the reader.
+        // 2. Output is unrelated to the wrapper (Self).
+        result.map_ok(|buf| unsafe { std::mem::transmute::<&'_ [u8], &'a [u8]>(buf) })
+    }
+}
diff --git a/src/io/buf_read/mod.rs b/src/io/buf_read/mod.rs
index e320375f..6018439a 100644
--- a/src/io/buf_read/mod.rs
+++ b/src/io/buf_read/mod.rs
@@ -1,7 +1,9 @@
+mod fill_buf;
 mod lines;
 mod read_line;
 mod read_until;
 
+use fill_buf::FillBufFuture;
 pub use lines::Lines;
 use read_line::ReadLineFuture;
 use read_until::ReadUntilFuture;
@@ -41,6 +43,26 @@ cfg_if! {
 /// [`futures::io::AsyncBufRead`]:
 /// https://docs.rs/futures-preview/0.3.0-alpha.17/futures/io/trait.AsyncBufRead.html
 pub trait BufRead {
+    /// 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.
+    ///
+    /// [`consume`]: #tymethod.consume
+    ///
+    /// An empty buffer returned indicates that the stream has reached EOF.
+    // TODO: write a proper doctest with `consume`
+    fn fill_buf<'a>(&'a mut self) -> ret!('a, FillBufFuture, io::Result<&'a [u8]>)
+    where
+        Self: Unpin,
+    {
+        FillBufFuture::new(self)
+    }
+
     /// Reads all bytes into `buf` until the delimiter `byte` or EOF is reached.
     ///
     /// This function will read bytes from the underlying stream until the delimiter or EOF is