From e93b8670dbd1545350464883d5405f2042126c59 Mon Sep 17 00:00:00 2001
From: Carl Lerche <me@carllerche.com>
Date: Wed, 20 Nov 2019 00:35:54 -0800
Subject: [PATCH] implicitly grow BytesMut; add BufMutExt::chain_mut

This brings `BytesMut` in line with `Vec<u8>` behavior.

In order to fix a test, `BufMutExt::chain_mut` is provided. Withou this,
it is not possible to chain two `&mut [u8]`.

Closes #170
---
 src/buf/ext/mod.rs    | 26 ++++++++++++++++++++++++++
 src/bytes_mut.rs      |  6 +++++-
 tests/test_buf_mut.rs |  2 +-
 tests/test_bytes.rs   |  4 ++--
 tests/test_chain.rs   | 13 +++++--------
 5 files changed, 39 insertions(+), 12 deletions(-)

diff --git a/src/buf/ext/mod.rs b/src/buf/ext/mod.rs
index 10b9a34c0..452535a70 100644
--- a/src/buf/ext/mod.rs
+++ b/src/buf/ext/mod.rs
@@ -124,6 +124,32 @@ pub trait BufMutExt: BufMut {
     fn writer(self) -> Writer<Self> where Self: Sized {
         writer::new(self)
     }
+
+    /// Creates an adaptor which will chain this buffer with another.
+    ///
+    /// The returned `BufMut` instance will first write to all bytes from
+    /// `self`. Afterwards, it will write to `next`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bytes::{BufMut, buf::BufMutExt};
+    ///
+    /// let mut a = [0u8; 5];
+    /// let mut b = [0u8; 6];
+    ///
+    /// let mut chain = (&mut a[..]).chain_mut(&mut b[..]);
+    ///
+    /// chain.put_slice(b"hello world");
+    ///
+    /// assert_eq!(&a[..], b"hello");
+    /// assert_eq!(&b[..], b" world");
+    /// ```
+    fn chain_mut<U: BufMut>(self, next: U) -> Chain<Self, U>
+        where Self: Sized
+    {
+        Chain::new(self, next)
+    }
 }
 
 impl<B: BufMut + ?Sized> BufMutExt for B {}
diff --git a/src/bytes_mut.rs b/src/bytes_mut.rs
index 6cf791a01..d030dcb9d 100644
--- a/src/bytes_mut.rs
+++ b/src/bytes_mut.rs
@@ -924,7 +924,7 @@ impl Buf for BytesMut {
 impl BufMut for BytesMut {
     #[inline]
     fn remaining_mut(&self) -> usize {
-        self.capacity() - self.len()
+        usize::MAX - self.len()
     }
 
     #[inline]
@@ -936,6 +936,10 @@ impl BufMut for BytesMut {
 
     #[inline]
     fn bytes_mut(&mut self) -> &mut [mem::MaybeUninit<u8>] {
+        if self.capacity() == self.len() {
+            self.reserve(64);
+        }
+
         unsafe {
             slice::from_raw_parts_mut(self.ptr.as_ptr().offset(self.len as isize) as *mut mem::MaybeUninit<u8>, self.cap)
         }
diff --git a/tests/test_buf_mut.rs b/tests/test_buf_mut.rs
index 0d372d1ec..dd414e5a2 100644
--- a/tests/test_buf_mut.rs
+++ b/tests/test_buf_mut.rs
@@ -74,7 +74,7 @@ fn test_bufs_vec_mut() {
     // with no capacity
     let mut buf = BytesMut::new();
     assert_eq!(buf.capacity(), 0);
-    assert_eq!(0, buf.bytes_vectored_mut(&mut dst[..]));
+    assert_eq!(1, buf.bytes_vectored_mut(&mut dst[..]));
 
     // with capacity
     let mut buf = BytesMut::with_capacity(64);
diff --git a/tests/test_bytes.rs b/tests/test_bytes.rs
index 9fa8019ee..918210695 100644
--- a/tests/test_bytes.rs
+++ b/tests/test_bytes.rs
@@ -93,8 +93,8 @@ fn fmt_write() {
 
 
     let mut c = BytesMut::with_capacity(64);
-    write!(c, "{}", s).unwrap_err();
-    assert!(c.is_empty());
+    write!(c, "{}", s).unwrap();
+    assert_eq!(c, s[..].as_bytes());
 }
 
 #[test]
diff --git a/tests/test_chain.rs b/tests/test_chain.rs
index 877618843..332571d8b 100644
--- a/tests/test_chain.rs
+++ b/tests/test_chain.rs
@@ -1,7 +1,7 @@
 #![deny(warnings, rust_2018_idioms)]
 
-use bytes::{Buf, BufMut, Bytes, BytesMut};
-use bytes::buf::BufExt;
+use bytes::{Buf, BufMut, Bytes};
+use bytes::buf::{BufExt, BufMutExt};
 use std::io::IoSlice;
 
 #[test]
@@ -15,20 +15,17 @@ fn collect_two_bufs() {
 
 #[test]
 fn writing_chained() {
-    let mut a = BytesMut::with_capacity(64);
-    let mut b = BytesMut::with_capacity(64);
+    let mut a = [0u8; 64];
+    let mut b = [0u8; 64];
 
     {
-        let mut buf = (&mut a).chain(&mut b);
+        let mut buf = (&mut a[..]).chain_mut(&mut b[..]);
 
         for i in 0u8..128 {
             buf.put_u8(i);
         }
     }
 
-    assert_eq!(64, a.len());
-    assert_eq!(64, b.len());
-
     for i in 0..64 {
         let expect = i as u8;
         assert_eq!(expect, a[i]);