Allow reclaiming the current allocation

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

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;
@@ -827,6 +828,73 @@ 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 cap = v.capacity();
+            if cap == 0 {
+                return None;
+            }
+
+            let ptr = v.as_mut_ptr();
+            self.ptr = vptr(ptr);
+            self.cap = cap;
+            Some(NonZeroUsize::new_unchecked(cap))
+        }
+    }
+
     // private
 
     // For now, use a `Vec` to manage the memory for us, but we may want to

tests/test_bytes.rs

@@ -1,6 +1,7 @@
 #![warn(rust_2018_idioms)]
 
 use bytes::{Buf, BufMut, Bytes, BytesMut};
+use std::num::NonZeroUsize;
 
 use std::usize;
 
@@ -1283,3 +1284,49 @@ fn test_bytes_make_mut_promotable_even_arc_offset() {
     assert_eq!(b2m, vec[20..]);
     assert_eq!(b1m, vec[..20]);
 }
+
+#[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());
+
+    let mut buf = BytesMut::new();
+    let mut split = buf.split();
+    drop(buf);
+    assert_eq!(0, split.capacity());
+    assert_eq!(None, split.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());
+}