io: add ReadBuf::take (#2817)

Signed-off-by: Zahari Dichev <[email protected]>

tokio/src/io/read_buf.rs

@@ -77,6 +77,14 @@ impl<'a> ReadBuf<'a> {
         unsafe { mem::transmute::<&mut [MaybeUninit<u8>], &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.

tokio/src/io/util/take.rs

@@ -85,10 +85,7 @@ impl<R: AsyncRead> AsyncRead for Take<R> {
         }
 
         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();