From b9c81701b81db8d580dfa5924591eb8da3ed39af Mon Sep 17 00:00:00 2001
From: Sebastian Hahn <sebastian@torproject.org>
Date: Mon, 8 Apr 2024 14:26:29 +0200
Subject: [PATCH] Address review: Add safety comments, update docs

---
 src/bytes_mut.rs | 93 ++++++++++++++++++++++++++++++------------------
 1 file changed, 58 insertions(+), 35 deletions(-)

diff --git a/src/bytes_mut.rs b/src/bytes_mut.rs
index a5e8c70e0..54a2cdcb9 100644
--- a/src/bytes_mut.rs
+++ b/src/bytes_mut.rs
@@ -827,15 +827,18 @@ impl BytesMut {
         }
     }
 
-    /// Attempts to reclaim the whole allocation of the `BytesMut`.
+    /// Attempts to cheaply 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 capacity as long as there are no
-    /// other `Bytes` or `BytesMut` handles alive.
+    /// If no other `Bytes` or `BytesMut` handles to the underlying storage of this
+    /// `BytesMut` are alive, it might be possible to cheaply reclaim the full capacity
+    /// of the allocation.
     ///
-    /// Returns `true` if the allocation was reclaimed (or reclaiming was not
-    /// necessary, because the underlying storage was empty); returns `false` if the
-    /// `BytesMut` is not empty or there are other outstanding references to the
+    /// `try_reclaim` never copies bytes internally, therefore it always returns `false`
+    /// if the `BytesMut` is not empty.
+    ///
+    /// Returns `true` if the allocation was reclaimed (or reclaiming was not necessary,
+    /// because the underlying storage was already fully available); returns `false` if the
+    /// `BytesMut` is not empty or there are other outstanding handles to the
     /// underlying storage.
     ///
     /// # Examples
@@ -845,15 +848,22 @@ impl BytesMut {
     ///
     /// let mut buf = BytesMut::with_capacity(64);
     /// assert_eq!(true, buf.try_reclaim());
+    /// assert_eq!(64, buf.capacity());
     ///
     /// buf.extend_from_slice(b"abcd");
     /// let mut split = buf.split();
+    /// assert_eq!(60, buf.capacity());
+    /// assert_eq!(4, split.capacity());
     /// assert_eq!(false, split.try_reclaim());
     /// assert_eq!(false, buf.try_reclaim());
+    ///
     /// drop(buf);
     /// assert_eq!(false, split.try_reclaim());
+    ///
     /// split.clear();
+    /// assert_eq!(4, split.capacity());
     /// assert_eq!(true, split.try_reclaim());
+    /// assert_eq!(64, split.capacity());
     /// ```
     #[must_use = "consider BytesMut::reserve if you need a specific capacity"]
     pub fn try_reclaim(&mut self) -> bool {
@@ -863,41 +873,54 @@ impl BytesMut {
 
         let kind = self.kind();
         if kind == KIND_VEC {
-            unsafe {
-                // Short-circuit if we are already at the start of the vector.
-                let off = self.get_vec_pos();
-                if off == 0 {
-                    return true;
-                }
-                // Otherwise, update info to point at the front of the vector
-                // again and reclaim capacity
-                let base_ptr = self.ptr.as_ptr().sub(off);
-                self.ptr = vptr(base_ptr);
-                self.set_vec_pos(0);
-                self.cap += off;
+            // Safety: self is of KIND_VEC, so calling get_vec_pos() is safe.
+            let off = unsafe { self.get_vec_pos() };
+
+            // Short-circuit if we are already at the start of the vector.
+            if off == 0 {
                 return true;
             }
+            // Otherwise, update info to point at the front of the vector
+            // again and reclaim capacity
+            //
+            // Safety: Offset `off` means self.ptr was moved `off` bytes from the base
+            // pointer of the allocation. Going back to the base pointer stays within
+            // that same allocation.
+            let base_ptr = unsafe { self.ptr.as_ptr().sub(off) };
+            self.ptr = vptr(base_ptr);
+
+            // Safety: Resetting the offset to 0 is safe as we reset the storage
+            // pointer to the base pointer of the allocation above
+            unsafe { self.set_vec_pos(0) }
+            self.cap += off;
+
+            return true;
         }
+
+        debug_assert_eq!(kind, KIND_ARC);
         let shared: *mut Shared = self.data;
 
-        unsafe {
-            // If there are other references to the underlying storage, we cannot reclaim
-            if !(*shared).is_unique() {
-                return false;
-            }
-            let v = &mut (*shared).vec;
-            let cap = v.capacity();
-            if cap == 0 {
-                // Short-circuit as an empty vector is trivially reclaimed already
-                return true;
-            }
+        // If there are other references to the underlying storage, we cannot reclaim
+        //
+        // Safety: self is of type KIND_ARC, so the Shared ptr is valid
+        if unsafe { !(*shared).is_unique() } {
+            return false;
+        }
 
-            // Update info to point at start of allocation and reclaim capacity
-            let ptr = v.as_mut_ptr();
-            self.ptr = vptr(ptr);
-            self.cap = cap;
-            true
+        // Safety: self is of type KIND_ARC and there are no other handles alive, so we
+        // can get a mut reference to the underlying storage
+        let v = unsafe { &mut (*shared).vec };
+        let cap = v.capacity();
+        if cap == 0 {
+            // Short-circuit as an empty vector is trivially reclaimed already
+            return true;
         }
+
+        // Update info to point at start of allocation and reclaim capacity
+        let ptr = v.as_mut_ptr();
+        self.ptr = vptr(ptr);
+        self.cap = cap;
+        true
     }
 
     // private