Allow reclaiming the current allocation

This is based on #680, where it was noted that it is hard to use
BytesMut without additional allocations in some circumstances.

Author: shahn

src/bytes_mut.rs

@@ -11,6 +11,7 @@ use alloc::{
     vec,
     vec::Vec,
 };
+use std::num::NonZeroUsize;
 
 use crate::buf::{IntoIter, UninitSlice};
 use crate::bytes::Vtable;
@@ -819,6 +820,70 @@ impl BytesMut {
         }
     }
 
+    /// Attempts to reclaim the whole allocation of the `BytesMut`.
+    ///
+    /// If the `BytesMut` is empty but the underlying storage has been used before,
+    /// it might be possible to cheaply reclaim space by just updating a few indices.
+    /// Returns `None` if the `BytesMut` is not empty, there is nothing to reclaim
+    /// (no underlying storage has been allocated), or there are any other live
+    /// references to the underlying storage. Otherwise, returns the available
+    /// capacity after reclaiming.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bytes::BytesMut;
+    /// use core::num::NonZeroUsize;
+    ///
+    /// let mut buf = BytesMut::with_capacity(64);
+    /// assert_eq!(None, buf.try_reclaim());
+    ///
+    /// buf.extend_from_slice(b"abcd");
+    /// let mut split = buf.split();
+    /// assert_eq!(None, split.try_reclaim());
+    /// assert_eq!(None, buf.try_reclaim());
+    /// drop(buf);
+    /// assert_eq!(None, split.try_reclaim());
+    /// split.clear();
+    /// assert_eq!(Some(64), split.try_reclaim().map(NonZeroUsize::into));
+    /// ```
+    pub fn try_reclaim(&mut self) -> Option<NonZeroUsize> {
+        if !self.is_empty() {
+            return None
+        }
+
+        let kind = self.kind();
+        if kind == KIND_VEC {
+            unsafe {
+                let off = self.get_vec_pos();
+                if off == 0 {
+                    return None
+                }
+
+                let base_ptr = self.ptr.as_ptr().sub(off);
+                self.ptr = vptr(base_ptr);
+                self.set_vec_pos(0);
+                self.cap += off;
+                debug_assert!(self.capacity() > 0);
+                return Some(NonZeroUsize::new_unchecked(self.capacity()))
+            }
+        }
+        let shared: *mut Shared = self.data;
+
+        unsafe {
+            if !(*shared).is_unique() {
+                return None
+            }
+            let v = &mut (*shared).vec;
+
+            let ptr = v.as_mut_ptr();
+            self.ptr = vptr(ptr);
+            self.cap = v.capacity();
+            debug_assert!(self.capacity() > 0);
+            Some(NonZeroUsize::new_unchecked(self.capacity()))
+        }
+    }
+
     // private
 
     // For now, use a `Vec` to manage the memory for us, but we may want to

tests/test_bytes.rs

@@ -1,5 +1,6 @@
 #![warn(rust_2018_idioms)]
 
+use std::num::NonZeroUsize;
 use bytes::{Buf, BufMut, Bytes, BytesMut};
 
 use std::usize;
@@ -1172,3 +1173,43 @@ fn shared_is_unique() {
     drop(b);
     assert!(c.is_unique());
 }
+
+#[test]
+fn try_reclaim_empty() {
+    let mut buf = BytesMut::new();
+    assert_eq!(None, buf.try_reclaim());
+    buf.reserve(6);
+    assert_eq!(None, buf.try_reclaim());
+}
+
+#[test]
+fn try_reclaim_vec() {
+    let mut buf = BytesMut::with_capacity(6);
+    buf.put_slice(b"abc");
+    buf.advance(3);
+    assert_eq!(3, buf.capacity());
+    assert_eq!(Some(NonZeroUsize::new(6).unwrap()), buf.try_reclaim());
+    assert_eq!(6, buf.capacity());
+}
+
+#[test]
+fn try_reclaim_arc() {
+    let mut buf = BytesMut::with_capacity(6);
+    buf.put_slice(b"abc");
+    let x = buf.split().freeze();
+    buf.put_slice(b"def");
+    let y = buf.split().freeze();
+    let z = y.clone();
+    assert_eq!(None, buf.try_reclaim());
+    drop(x);
+    drop(z);
+    assert_eq!(None, buf.try_reclaim());
+    drop(y);
+    assert_eq!(Some(NonZeroUsize::new(6).unwrap()), buf.try_reclaim());
+    assert_eq!(6, buf.capacity());
+    assert_eq!(0, buf.len());
+    buf.put_slice(b"abc");
+    buf.put_slice(b"def");
+    assert_eq!(6, buf.capacity());
+    assert_eq!(6, buf.len());
+}