Specialize BytesMut::put::<Bytes> (#793)

src/bytes_mut.rs

@@ -1201,11 +1201,21 @@ unsafe impl BufMut for BytesMut {
     where
         Self: Sized,
     {
-        while src.has_remaining() {
+        // When capacity is zero, try reusing allocation of `src`.
-            let s = src.chunk();
+        if self.capacity() == 0 {
-            let l = s.len();
+            let src_copy = src.copy_to_bytes(src.remaining());
-            self.extend_from_slice(s);
+            drop(src);
-            src.advance(l);
+            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);
+            }
         }
     }
 

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);
+}