implicitly grow BytesMut; add BufMutExt::chain_mut (#316)

carllerche · web-flow · commit 8135c1f60692 · 2019-11-20T12:11:40.000-08:00
This brings `BytesMut` in line with `Vec<u8>` behavior. This also fixes an existing bug in BytesMut::bytes_mut that exposes invalid slices. The bug was recently introduced and was only on master and never released to `crates.io`. In order to fix a test, `BufMutExt::chain_mut` is provided. Withou this, it is not possible to chain two `&mut [u8]`. Closes #170
diff --git a/src/buf/buf_mut.rs b/src/buf/buf_mut.rs
@@ -270,7 +270,7 @@ pub trait BufMut {
     fn put_slice(&mut self, src: &[u8]) {
         let mut off = 0;
 
-        assert!(self.remaining_mut() >= src.len(), "buffer overflow");
+        assert!(self.remaining_mut() >= src.len(), "buffer overflow; remaining = {}; src = {}", self.remaining_mut(), src.len());
 
         while off < src.len() {
             let cnt;
diff --git 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 adapter 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
@@ -1,4 +1,5 @@
-use core::{cmp, fmt, hash, isize, mem, slice, usize};
+use core::{cmp, fmt, hash, isize, slice, usize};
+use core::mem::{self, ManuallyDrop};
 use core::ops::{Deref, DerefMut};
 use core::ptr::{self, NonNull};
 use core::iter::{FromIterator, Iterator};
@@ -559,17 +560,15 @@ impl BytesMut {
                     self.cap += off;
                 } else {
                     // No space - allocate more
-                    let mut v = rebuild_vec(self.ptr.as_ptr(), self.len, self.cap, off);
+                    let mut v = ManuallyDrop::new(rebuild_vec(self.ptr.as_ptr(), self.len, self.cap, off));
                     v.reserve(additional);
 
                     // Update the info
                     self.ptr = vptr(v.as_mut_ptr().offset(off as isize));
                     self.len = v.len() - off;
                     self.cap = v.capacity() - off;
-
-                    // Drop the vec reference
-                    mem::forget(v);
                 }
+
                 return;
             }
         }
@@ -582,7 +581,8 @@ impl BytesMut {
         // allocating a new vector with the requested capacity.
         //
         // Compute the new capacity
-        let mut new_cap = len + additional;
+        let mut new_cap = len.checked_add(additional).expect("overflow");
+
         let original_capacity;
         let original_capacity_repr;
 
@@ -618,16 +618,18 @@ impl BytesMut {
                 // There are some situations, using `reserve_exact` that the
                 // buffer capacity could be below `original_capacity`, so do a
                 // check.
+                let double = v.capacity().checked_shl(1).unwrap_or(new_cap);
+
                 new_cap = cmp::max(
-                    cmp::max(v.capacity() << 1, new_cap),
+                    cmp::max(double, new_cap),
                     original_capacity);
             } else {
                 new_cap = cmp::max(new_cap, original_capacity);
             }
         }
 
         // Create a new vector to store the data
-        let mut v = Vec::with_capacity(new_cap);
+        let mut v = ManuallyDrop::new(Vec::with_capacity(new_cap));
 
         // Copy the bytes
         v.extend_from_slice(self.as_ref());
@@ -642,10 +644,6 @@ impl BytesMut {
         self.ptr = vptr(v.as_mut_ptr());
         self.len = v.len();
         self.cap = v.capacity();
-
-        // Forget the vector handle
-        mem::forget(v);
-
     }
     /// Appends given bytes to this object.
     ///
@@ -924,20 +922,27 @@ impl Buf for BytesMut {
 impl BufMut for BytesMut {
     #[inline]
     fn remaining_mut(&self) -> usize {
-        self.capacity() - self.len()
+        usize::MAX - self.len()
     }
 
     #[inline]
     unsafe fn advance_mut(&mut self, cnt: usize) {
         let new_len = self.len() + cnt;
-        assert!(new_len <= self.cap);
+        assert!(new_len <= self.cap, "new_len = {}; capacity = {}", new_len, self.cap);
         self.len = new_len;
     }
 
     #[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)
+            let ptr = self.ptr.as_ptr().offset(self.len as isize);
+            let len = self.cap - self.len;
+
+            slice::from_raw_parts_mut(ptr as *mut mem::MaybeUninit<u8>, len)
         }
     }
 }
diff --git 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
@@ -2,6 +2,8 @@
 
 use bytes::{Bytes, BytesMut, Buf, BufMut};
 
+use std::usize;
+
 const LONG: &'static [u8] = b"mary had a little lamb, little lamb, little lamb";
 const SHORT: &'static [u8] = b"hello world";
 
@@ -93,8 +95,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]
@@ -820,3 +822,34 @@ fn empty_slice_ref_catches_not_an_empty_subset() {
 
     bytes.slice_ref(slice);
 }
+
+#[test]
+fn bytes_buf_mut_advance() {
+    let mut bytes = BytesMut::with_capacity(1024);
+
+    unsafe {
+        let ptr = bytes.bytes_mut().as_ptr();
+        assert_eq!(1024, bytes.bytes_mut().len());
+
+        bytes.advance_mut(10);
+
+        let next = bytes.bytes_mut().as_ptr();
+        assert_eq!(1024 - 10, bytes.bytes_mut().len());
+        assert_eq!(ptr.offset(10), next);
+
+        // advance to the end
+        bytes.advance_mut(1024 - 10);
+
+        // The buffer size is doubled
+        assert_eq!(1024, bytes.bytes_mut().len());
+    }
+}
+
+#[test]
+#[should_panic]
+fn bytes_reserve_overflow() {
+    let mut bytes = BytesMut::with_capacity(1024);
+    bytes.put_slice(b"hello world");
+
+    bytes.reserve(usize::MAX);
+}
diff --git 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]);
