From 6b911d129f05ba0e9e970903806e924c21d95a10 Mon Sep 17 00:00:00 2001
From: Sam Rijs <srijs@airpost.net>
Date: Mon, 12 Jun 2023 01:22:39 +0200
Subject: [PATCH 1/2] add specialized `Buf::chunks_vectored` for `Take`

---
 src/buf/take.rs    | 25 +++++++++++++++++++++++
 tests/test_take.rs | 51 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 76 insertions(+)

diff --git a/src/buf/take.rs b/src/buf/take.rs
index d3cb10ab6..77799bb5e 100644
--- a/src/buf/take.rs
+++ b/src/buf/take.rs
@@ -152,4 +152,29 @@ impl<T: Buf> Buf for Take<T> {
         self.limit -= len;
         r
     }
+
+    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
+    }
 }
diff --git a/tests/test_take.rs b/tests/test_take.rs
index 51df91d14..767b19f1a 100644
--- a/tests/test_take.rs
+++ b/tests/test_take.rs
@@ -30,3 +30,54 @@ fn take_copy_to_bytes_panics() {
     let abcd = Bytes::copy_from_slice(b"abcd");
     abcd.take(2).copy_to_bytes(3);
 }
+
+#[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]);
+    }
+}

From f8ae0bbd0d7940cad6090280c03b13c4a15a1f6a Mon Sep 17 00:00:00 2001
From: Sam Rijs <srijs@airpost.net>
Date: Tue, 13 Jun 2023 00:53:05 +0200
Subject: [PATCH 2/2] make compatible with no-std

---
 src/buf/take.rs    | 1 +
 tests/test_take.rs | 1 +
 2 files changed, 2 insertions(+)

diff --git a/src/buf/take.rs b/src/buf/take.rs
index 77799bb5e..8171ab400 100644
--- a/src/buf/take.rs
+++ b/src/buf/take.rs
@@ -153,6 +153,7 @@ impl<T: Buf> Buf for Take<T> {
         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;
diff --git a/tests/test_take.rs b/tests/test_take.rs
index 767b19f1a..0c0159be1 100644
--- a/tests/test_take.rs
+++ b/tests/test_take.rs
@@ -31,6 +31,7 @@ fn take_copy_to_bytes_panics() {
     abcd.take(2).copy_to_bytes(3);
 }
 
+#[cfg(feature = "std")]
 #[test]
 fn take_chunks_vectored() {
     fn chain() -> impl Buf {