From 0bb015588a28a2f7cfd505d611975bfb2123d044 Mon Sep 17 00:00:00 2001
From: Taiki Endo <te316e89@gmail.com>
Date: Sun, 4 Aug 2019 12:15:50 +0900
Subject: [PATCH] codec: add AsyncBufRead/BufRead implementations (#1385)

* AsyncBufRead for FramedWrite2<T>
* BufRead for FramedWrite2<T>
* AsyncBufRead for Fuse<T, U>
* BufRead for Fuse<T, U>
---
 tokio-codec/src/framed.rs       | 27 +++++++++++++++++++++++++--
 tokio-codec/src/framed_write.rs | 29 ++++++++++++++++++++++++++---
 2 files changed, 51 insertions(+), 5 deletions(-)

diff --git a/tokio-codec/src/framed.rs b/tokio-codec/src/framed.rs
index 00e477578..707353dbc 100644
--- a/tokio-codec/src/framed.rs
+++ b/tokio-codec/src/framed.rs
@@ -5,13 +5,13 @@ use crate::encoder::Encoder;
 use crate::framed_read::{framed_read2, framed_read2_with_buffer, FramedRead2};
 use crate::framed_write::{framed_write2, framed_write2_with_buffer, FramedWrite2};
 
-use tokio_io::{AsyncRead, AsyncWrite};
+use tokio_io::{AsyncBufRead, AsyncRead, AsyncWrite};
 
 use bytes::BytesMut;
 use futures_core::Stream;
 use futures_sink::Sink;
 use std::fmt;
-use std::io::{self, Read, Write};
+use std::io::{self, BufRead, Read, Write};
 use std::pin::Pin;
 use std::task::{Context, Poll};
 
@@ -209,6 +209,16 @@ impl<T: Read, U> Read for Fuse<T, U> {
     }
 }
 
+impl<T: BufRead, U> BufRead for Fuse<T, U> {
+    fn fill_buf(&mut self) -> io::Result<&[u8]> {
+        self.0.fill_buf()
+    }
+
+    fn consume(&mut self, amt: usize) {
+        self.0.consume(amt)
+    }
+}
+
 impl<T: AsyncRead + Unpin, U: Unpin> AsyncRead for Fuse<T, U> {
     unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [u8]) -> bool {
         self.0.prepare_uninitialized_buffer(buf)
@@ -223,6 +233,19 @@ impl<T: AsyncRead + Unpin, U: Unpin> AsyncRead for Fuse<T, U> {
     }
 }
 
+impl<T: AsyncBufRead + Unpin, U: Unpin> AsyncBufRead for Fuse<T, U> {
+    fn poll_fill_buf<'a>(
+        self: Pin<&'a mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<io::Result<&'a [u8]>> {
+        pin!(self.get_mut().0).poll_fill_buf(cx)
+    }
+
+    fn consume(self: Pin<&mut Self>, amt: usize) {
+        pin!(self.get_mut().0).consume(amt)
+    }
+}
+
 impl<T: Write, U> Write for Fuse<T, U> {
     fn write(&mut self, src: &[u8]) -> io::Result<usize> {
         self.0.write(src)
diff --git a/tokio-codec/src/framed_write.rs b/tokio-codec/src/framed_write.rs
index c150847be..34d246065 100644
--- a/tokio-codec/src/framed_write.rs
+++ b/tokio-codec/src/framed_write.rs
@@ -4,14 +4,14 @@ use super::framed::Fuse;
 use crate::decoder::Decoder;
 use crate::encoder::Encoder;
 
-use tokio_io::{AsyncRead, AsyncWrite};
+use tokio_io::{AsyncBufRead, AsyncRead, AsyncWrite};
 
 use bytes::BytesMut;
 use futures_core::{ready, Stream};
 use futures_sink::Sink;
 use log::trace;
 use std::fmt;
-use std::io::{self, Read};
+use std::io::{self, BufRead, Read};
 use std::pin::Pin;
 use std::task::{Context, Poll};
 
@@ -255,6 +255,16 @@ impl<T: Read> Read for FramedWrite2<T> {
     }
 }
 
+impl<T: BufRead> BufRead for FramedWrite2<T> {
+    fn fill_buf(&mut self) -> io::Result<&[u8]> {
+        self.inner.fill_buf()
+    }
+
+    fn consume(&mut self, amt: usize) {
+        self.inner.consume(amt)
+    }
+}
+
 impl<T: AsyncRead + Unpin> AsyncRead for FramedWrite2<T> {
     unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [u8]) -> bool {
         self.inner.prepare_uninitialized_buffer(buf)
@@ -265,6 +275,19 @@ impl<T: AsyncRead + Unpin> AsyncRead for FramedWrite2<T> {
         cx: &mut Context<'_>,
         buf: &mut [u8],
     ) -> Poll<Result<usize, io::Error>> {
-        pin!(Pin::get_mut(self).inner).poll_read(cx, buf)
+        pin!(self.get_mut().inner).poll_read(cx, buf)
+    }
+}
+
+impl<T: AsyncBufRead + Unpin> AsyncBufRead for FramedWrite2<T> {
+    fn poll_fill_buf<'a>(
+        self: Pin<&'a mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<io::Result<&'a [u8]>> {
+        pin!(self.get_mut().inner).poll_fill_buf(cx)
+    }
+
+    fn consume(self: Pin<&mut Self>, amt: usize) {
+        pin!(self.get_mut().inner).consume(amt)
     }
 }