From ef7f25736cfe48db5b42f2aa3225cace4126d2ac Mon Sep 17 00:00:00 2001
From: Stepan Koltsov <stepan.koltsov@gmail.com>
Date: Thu, 14 Aug 2025 09:55:18 +0100
Subject: [PATCH] Specialize BytesMut::put::<Bytes> (#793)

---
 src/bytes_mut.rs    | 20 +++++++++++++++-----
 tests/test_bytes.rs | 19 +++++++++++++++++++
 2 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/src/bytes_mut.rs b/src/bytes_mut.rs
index 6a4b9e7..ad3aa77 100644
--- a/src/bytes_mut.rs
+++ b/src/bytes_mut.rs
@@ -1201,11 +1201,21 @@ unsafe impl BufMut for BytesMut {
     where
         Self: Sized,
     {
-        while src.has_remaining() {
-            let s = src.chunk();
-            let l = s.len();
-            self.extend_from_slice(s);
-            src.advance(l);
+        // When capacity is zero, try reusing allocation of `src`.
+        if self.capacity() == 0 {
+            let src_copy = src.copy_to_bytes(src.remaining());
+            drop(src);
+            match src_copy.try_into_mut() {
+                Ok(bytes_mut) => *self = bytes_mut,
+                Err(bytes) => self.extend_from_slice(&bytes),
+            }
+        } else {
+            while src.has_remaining() {
+                let s = src.chunk();
+                let l = s.len();
+                self.extend_from_slice(s);
+                src.advance(l);
+            }
         }
     }
 
diff --git a/tests/test_bytes.rs b/tests/test_bytes.rs
index 85c0079..ec9a60e 100644
--- a/tests/test_bytes.rs
+++ b/tests/test_bytes.rs
@@ -1688,3 +1688,22 @@ fn owned_safe_drop_on_as_ref_panic() {
     assert!(result.is_err());
     assert_eq!(drop_counter.get(), 1);
 }
+
+/// Test `BytesMut::put` reuses allocation of `Bytes`.
+#[test]
+fn bytes_mut_put_bytes_specialization() {
+    let mut vec = Vec::with_capacity(1234);
+    vec.push(10);
+    let capacity = vec.capacity();
+    assert!(capacity >= 1234);
+
+    // Make `Bytes` backed by `Vec`.
+    let bytes = Bytes::from(vec);
+    let mut bytes_mut = BytesMut::new();
+    bytes_mut.put(bytes);
+
+    // Check contents is correct.
+    assert_eq!(&[10], bytes_mut.as_ref());
+    // If allocation is reused, capacity should be equal to original vec capacity.
+    assert_eq!(bytes_mut.capacity(), capacity);
+}