add specialized Buf::chunks_vectored for Take

src/buf/take.rs

@@ -152,4 +152,30 @@ impl<T: Buf> Buf for Take<T> {
         self.limit -= len;
         r
     }
+
+    #[cfg(feature = "std")]
+    fn chunks_vectored<'a>(&'a self, dst: &mut [std::io::IoSlice<'a>]) -> usize {
+        let cnt = self.inner.chunks_vectored(dst);
+        let mut len = 0;
+        for (n, io) in dst[0..cnt].iter_mut().enumerate() {
+            let max = self.limit - len;
+            if max == 0 {
+                return n;
+            }
+            if io.len() > max {
+                // In this case, `IoSlice` is longer than our max, so we need to truncate it to the max.
+                //
+                // We need to work around the fact here that even though `IoSlice<'a>` has the correct
+                // lifetime, its `Deref` impl strips it. So we need to reassamble the slice to add the
+                // correct lifetime that allows us to call `IoSlice::<'a>::new` with it.
+                //
+                // TODO: remove `unsafe` as soon as `IoSlice::as_bytes` is available (rust-lang/rust#111277)
+                let buf = unsafe { std::slice::from_raw_parts::<'a, u8>(io.as_ptr(), max) };
+                *io = std::io::IoSlice::new(buf);
+                return n + 1;
+            }
+            len += io.len();
+        }
+        cnt
+    }
 }

tests/test_take.rs

@@ -30,3 +30,55 @@ fn take_copy_to_bytes_panics() {
     let abcd = Bytes::copy_from_slice(b"abcd");
     abcd.take(2).copy_to_bytes(3);
 }
+
+#[cfg(feature = "std")]
+#[test]
+fn take_chunks_vectored() {
+    fn chain() -> impl Buf {
+        Bytes::from([1, 2, 3].to_vec()).chain(Bytes::from([4, 5, 6].to_vec()))
+    }
+
+    {
+        let mut dst = [std::io::IoSlice::new(&[]); 2];
+        let take = chain().take(0);
+        assert_eq!(take.chunks_vectored(&mut dst), 0);
+    }
+
+    {
+        let mut dst = [std::io::IoSlice::new(&[]); 2];
+        let take = chain().take(1);
+        assert_eq!(take.chunks_vectored(&mut dst), 1);
+        assert_eq!(&*dst[0], &[1]);
+    }
+
+    {
+        let mut dst = [std::io::IoSlice::new(&[]); 2];
+        let take = chain().take(3);
+        assert_eq!(take.chunks_vectored(&mut dst), 1);
+        assert_eq!(&*dst[0], &[1, 2, 3]);
+    }
+
+    {
+        let mut dst = [std::io::IoSlice::new(&[]); 2];
+        let take = chain().take(4);
+        assert_eq!(take.chunks_vectored(&mut dst), 2);
+        assert_eq!(&*dst[0], &[1, 2, 3]);
+        assert_eq!(&*dst[1], &[4]);
+    }
+
+    {
+        let mut dst = [std::io::IoSlice::new(&[]); 2];
+        let take = chain().take(6);
+        assert_eq!(take.chunks_vectored(&mut dst), 2);
+        assert_eq!(&*dst[0], &[1, 2, 3]);
+        assert_eq!(&*dst[1], &[4, 5, 6]);
+    }
+
+    {
+        let mut dst = [std::io::IoSlice::new(&[]); 2];
+        let take = chain().take(7);
+        assert_eq!(take.chunks_vectored(&mut dst), 2);
+        assert_eq!(&*dst[0], &[1, 2, 3]);
+        assert_eq!(&*dst[1], &[4, 5, 6]);
+    }
+}