From 3d4290410f3698ecdcc351943d4a3d405fe5c6c9 Mon Sep 17 00:00:00 2001
From: Sebastian Hahn <sebastian@torproject.org>
Date: Sun, 31 Mar 2024 10:23:08 +0200
Subject: [PATCH] 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.
---
 src/bytes_mut.rs    | 68 +++++++++++++++++++++++++++++++++++++++++++++
 tests/test_bytes.rs | 47 +++++++++++++++++++++++++++++++
 2 files changed, 115 insertions(+)

diff --git a/src/bytes_mut.rs b/src/bytes_mut.rs
index 282aaa710..df686dc80 100644
--- a/src/bytes_mut.rs
+++ b/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,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
diff --git a/tests/test_bytes.rs b/tests/test_bytes.rs
index 84c3d5a43..ba9473f71 100644
--- a/tests/test_bytes.rs
+++ b/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;
 
@@ -1172,3 +1173,49 @@ 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());
+
+    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());
+}