From 171cb57fa1baffe369e44e1bf28382dcb54b2e1c Mon Sep 17 00:00:00 2001 From: Zahari Dichev Date: Sat, 5 Sep 2020 21:09:54 +0300 Subject: [PATCH] io: add `ReadBuf::take` (#2817) Signed-off-by: Zahari Dichev --- tokio/src/io/read_buf.rs | 8 ++++++++ tokio/src/io/util/take.rs | 5 +---- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/tokio/src/io/read_buf.rs b/tokio/src/io/read_buf.rs index 03b5d05ca03..8b430ee1093 100644 --- a/tokio/src/io/read_buf.rs +++ b/tokio/src/io/read_buf.rs @@ -77,6 +77,14 @@ impl<'a> ReadBuf<'a> { unsafe { mem::transmute::<&mut [MaybeUninit], &mut [u8]>(slice) } } + /// Returns a new `ReadBuf` comprised of the unfilled section up to `n`. + #[inline] + pub fn take(&mut self, n: usize) -> ReadBuf<'_> { + let max = std::cmp::min(self.remaining(), n); + // Saftey: We don't set any of the `unfilled_mut` with `MaybeUninit::uninit`. + unsafe { ReadBuf::uninit(&mut self.unfilled_mut()[..max]) } + } + /// Returns a shared reference to the initialized portion of the buffer. /// /// This includes the filled portion. diff --git a/tokio/src/io/util/take.rs b/tokio/src/io/util/take.rs index 2abc7693172..4e424f6c8a5 100644 --- a/tokio/src/io/util/take.rs +++ b/tokio/src/io/util/take.rs @@ -85,10 +85,7 @@ impl AsyncRead for Take { } let me = self.project(); - let max = std::cmp::min(buf.remaining() as u64, *me.limit_) as usize; - // Make a ReadBuf of the unfulled section up to max - // Saftey: We don't set any of the `unfilled_mut` with `MaybeUninit::uninit`. - let mut b = unsafe { ReadBuf::uninit(&mut buf.unfilled_mut()[..max]) }; + let mut b = buf.take(*me.limit_ as usize); ready!(me.inner.poll_read(cx, &mut b))?; let n = b.filled().len();