Add specialized Buf::chunks_vectored for Take (#617)

Co-authored-by: Alice Ryhl <[email protected]>
Co-authored-by: Michal 'vorner' Vaner <[email protected]>

src/buf/take.rs

@@ -2,6 +2,9 @@ use crate::{Buf, Bytes};
 
 use core::cmp;
 
+#[cfg(feature = "std")]
+use std::io::IoSlice;
+
 /// A `Buf` adapter which limits the bytes read from an underlying buffer.
 ///
 /// This struct is generally created by calling `take()` on `Buf`. See
@@ -152,4 +155,50 @@ impl<T: Buf> Buf for Take<T> {
         self.limit -= len;
         r
     }
+
+    #[cfg(feature = "std")]
+    fn chunks_vectored<'a>(&'a self, dst: &mut [IoSlice<'a>]) -> usize {
+        if self.limit == 0 {
+            return 0;
+        }
+
+        const LEN: usize = 16;
+        let mut slices: [IoSlice<'a>; LEN] = [
+            IoSlice::new(&[]),
+            IoSlice::new(&[]),
+            IoSlice::new(&[]),
+            IoSlice::new(&[]),
+            IoSlice::new(&[]),
+            IoSlice::new(&[]),
+            IoSlice::new(&[]),
+            IoSlice::new(&[]),
+            IoSlice::new(&[]),
+            IoSlice::new(&[]),
+            IoSlice::new(&[]),
+            IoSlice::new(&[]),
+            IoSlice::new(&[]),
+            IoSlice::new(&[]),
+            IoSlice::new(&[]),
+            IoSlice::new(&[]),
+        ];
+
+        let cnt = self
+            .inner
+            .chunks_vectored(&mut slices[..dst.len().min(LEN)]);
+        let mut limit = self.limit;
+        for (i, (dst, slice)) in dst[..cnt].iter_mut().zip(slices.iter()).enumerate() {
+            if let Some(buf) = slice.get(..limit) {
+                // SAFETY: We could do this safely with `IoSlice::advance` if we had a larger MSRV.
+                let buf = unsafe { std::mem::transmute::<&[u8], &'a [u8]>(buf) };
+                *dst = IoSlice::new(buf);
+                return i + 1;
+            } else {
+                // SAFETY: We could do this safely with `IoSlice::advance` if we had a larger MSRV.
+                let buf = unsafe { std::mem::transmute::<&[u8], &'a [u8]>(slice) };
+                *dst = IoSlice::new(buf);
+                limit -= slice.len();
+            }
+        }
+        cnt
+    }
 }