From b9eade12a57b923a54376762ae908da1cff6bcd4 Mon Sep 17 00:00:00 2001
From: Stepan Koltsov <stepan.koltsov@gmail.com>
Date: Sat, 10 Apr 2021 19:56:35 +0100
Subject: [PATCH] Specialize copy_to_bytes for Chain and Take (#481)

Avoid allocation when `Take` or `Chain` is composed of `Bytes`
objects.

This works now for `Take`.

`Chain` it works if the requested bytes does not cross boundary
between `Chain` members.
---
 src/buf/chain.rs    | 20 +++++++++++++++++++-
 src/buf/take.rs     | 10 +++++++++-
 tests/test_chain.rs | 21 +++++++++++++++++++++
 tests/test_take.rs  | 20 ++++++++++++++++++++
 4 files changed, 69 insertions(+), 2 deletions(-)

diff --git a/src/buf/chain.rs b/src/buf/chain.rs
index 8a1598c..9ce5f23 100644
--- a/src/buf/chain.rs
+++ b/src/buf/chain.rs
@@ -1,5 +1,5 @@
 use crate::buf::{IntoIter, UninitSlice};
-use crate::{Buf, BufMut};
+use crate::{Buf, BufMut, Bytes};
 
 #[cfg(feature = "std")]
 use std::io::IoSlice;
@@ -170,6 +170,24 @@ where
         n += self.b.chunks_vectored(&mut dst[n..]);
         n
     }
+
+    fn copy_to_bytes(&mut self, len: usize) -> Bytes {
+        let a_rem = self.a.remaining();
+        if a_rem >= len {
+            self.a.copy_to_bytes(len)
+        } else if a_rem == 0 {
+            self.b.copy_to_bytes(len)
+        } else {
+            assert!(
+                len - a_rem <= self.b.remaining(),
+                "`len` greater than remaining"
+            );
+            let mut ret = crate::BytesMut::with_capacity(len);
+            ret.put(&mut self.a);
+            ret.put((&mut self.b).take(len - a_rem));
+            ret.freeze()
+        }
+    }
 }
 
 unsafe impl<T, U> BufMut for Chain<T, U>
diff --git a/src/buf/take.rs b/src/buf/take.rs
index 2f4c436..d3cb10a 100644
--- a/src/buf/take.rs
+++ b/src/buf/take.rs
@@ -1,4 +1,4 @@
-use crate::Buf;
+use crate::{Buf, Bytes};
 
 use core::cmp;
 
@@ -144,4 +144,12 @@ impl<T: Buf> Buf for Take<T> {
         self.inner.advance(cnt);
         self.limit -= cnt;
     }
+
+    fn copy_to_bytes(&mut self, len: usize) -> Bytes {
+        assert!(len <= self.remaining(), "`len` greater than remaining");
+
+        let r = self.inner.copy_to_bytes(len);
+        self.limit -= len;
+        r
+    }
 }
diff --git a/tests/test_chain.rs b/tests/test_chain.rs
index 500ccd4..affaf7a 100644
--- a/tests/test_chain.rs
+++ b/tests/test_chain.rs
@@ -132,3 +132,24 @@ fn vectored_read() {
         assert_eq!(iovecs[3][..], b""[..]);
     }
 }
+
+#[test]
+fn chain_get_bytes() {
+    let mut ab = Bytes::copy_from_slice(b"ab");
+    let mut cd = Bytes::copy_from_slice(b"cd");
+    let ab_ptr = ab.as_ptr();
+    let cd_ptr = cd.as_ptr();
+    let mut chain = (&mut ab).chain(&mut cd);
+    let a = chain.copy_to_bytes(1);
+    let bc = chain.copy_to_bytes(2);
+    let d = chain.copy_to_bytes(1);
+
+    assert_eq!(Bytes::copy_from_slice(b"a"), a);
+    assert_eq!(Bytes::copy_from_slice(b"bc"), bc);
+    assert_eq!(Bytes::copy_from_slice(b"d"), d);
+
+    // assert `get_bytes` did not allocate
+    assert_eq!(ab_ptr, a.as_ptr());
+    // assert `get_bytes` did not allocate
+    assert_eq!(cd_ptr.wrapping_offset(1), d.as_ptr());
+}
diff --git a/tests/test_take.rs b/tests/test_take.rs
index a23a29e..51df91d 100644
--- a/tests/test_take.rs
+++ b/tests/test_take.rs
@@ -1,6 +1,7 @@
 #![warn(rust_2018_idioms)]
 
 use bytes::buf::Buf;
+use bytes::Bytes;
 
 #[test]
 fn long_take() {
@@ -10,3 +11,22 @@ fn long_take() {
     assert_eq!(11, buf.remaining());
     assert_eq!(b"hello world", buf.chunk());
 }
+
+#[test]
+fn take_copy_to_bytes() {
+    let mut abcd = Bytes::copy_from_slice(b"abcd");
+    let abcd_ptr = abcd.as_ptr();
+    let mut take = (&mut abcd).take(2);
+    let a = take.copy_to_bytes(1);
+    assert_eq!(Bytes::copy_from_slice(b"a"), a);
+    // assert `to_bytes` did not allocate
+    assert_eq!(abcd_ptr, a.as_ptr());
+    assert_eq!(Bytes::copy_from_slice(b"bcd"), abcd);
+}
+
+#[test]
+#[should_panic]
+fn take_copy_to_bytes_panics() {
+    let abcd = Bytes::copy_from_slice(b"abcd");
+    abcd.take(2).copy_to_bytes(3);
+}