diff options
Diffstat (limited to 'src/io/util/read_exact.rs')
-rw-r--r-- | src/io/util/read_exact.rs | 46 |
1 files changed, 19 insertions, 27 deletions
diff --git a/src/io/util/read_exact.rs b/src/io/util/read_exact.rs index 86b8412..1e8150e 100644 --- a/src/io/util/read_exact.rs +++ b/src/io/util/read_exact.rs @@ -1,7 +1,9 @@ -use crate::io::AsyncRead; +use crate::io::{AsyncRead, ReadBuf}; +use pin_project_lite::pin_project; use std::future::Future; use std::io; +use std::marker::PhantomPinned; use std::marker::Unpin; use std::pin::Pin; use std::task::{Context, Poll}; @@ -17,12 +19,12 @@ where { ReadExact { reader, - buf, - pos: 0, + buf: ReadBuf::new(buf), + _pin: PhantomPinned, } } -cfg_io_util! { +pin_project! { /// Creates a future which will read exactly enough bytes to fill `buf`, /// returning an error if EOF is hit sooner. /// @@ -31,8 +33,10 @@ cfg_io_util! { #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct ReadExact<'a, A: ?Sized> { reader: &'a mut A, - buf: &'a mut [u8], - pos: usize, + buf: ReadBuf<'a>, + // Make this future `!Unpin` for compatibility with async trait methods. + #[pin] + _pin: PhantomPinned, } } @@ -46,32 +50,20 @@ where { type Output = io::Result<usize>; - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>> { + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>> { + let mut me = self.project(); + loop { // if our buffer is empty, then we need to read some data to continue. - if self.pos < self.buf.len() { - let me = &mut *self; - let n = ready!(Pin::new(&mut *me.reader).poll_read(cx, &mut me.buf[me.pos..]))?; - me.pos += n; - if n == 0 { + let rem = me.buf.remaining(); + if rem != 0 { + ready!(Pin::new(&mut *me.reader).poll_read(cx, &mut me.buf))?; + if me.buf.remaining() == rem { return Err(eof()).into(); } - } - - if self.pos >= self.buf.len() { - return Poll::Ready(Ok(self.pos)); + } else { + return Poll::Ready(Ok(me.buf.capacity())); } } } } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn assert_unpin() { - use std::marker::PhantomPinned; - crate::is_unpin::<ReadExact<'_, PhantomPinned>>(); - } -} |