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

src/bytes_mut.rs

@@ -1201,6 +1201,15 @@ unsafe impl BufMut for BytesMut {
     where
         Self: Sized,
     {
+        // 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();
@@ -1208,6 +1217,7 @@ unsafe impl BufMut for BytesMut {
                 src.advance(l);
             }
         }
+    }
 
     fn put_slice(&mut self, src: &[u8]) {
         self.extend_from_slice(src);

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