Add BytesMut::try_shrink_to_fit and BytesMut::try_shrink_to_capacity

src/bytes_mut.rs

@@ -858,6 +858,84 @@ impl BytesMut {
         self.reserve_inner(additional, false)
     }
 
+    /// Attempts to shrink the buffer's capacity with a lower bound.
+    ///
+    /// The capacity will remain at least as large as both the length and the
+    /// supplied capacity. This operation preserves the current contents and
+    /// returns `true` when the allocation can be reclaimed. For `BytesMut`
+    /// values backed by shared storage, this only succeeds when the current
+    /// handle is the only remaining owner of the underlying allocation;
+    /// otherwise it returns `false` and leaves the buffer unchanged.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bytes::BytesMut;
+    ///
+    /// let mut buf = BytesMut::with_capacity(64);
+    /// buf.extend_from_slice(b"hello");
+    ///
+    /// assert!(buf.try_shrink_to_capacity(16));
+    /// assert_eq!(buf.capacity(), 16);
+    /// assert_eq!(&buf[..], b"hello");
+    /// ```
+    pub fn try_shrink_to_capacity(&mut self, capacity: usize) -> bool {
+        let kind = self.kind();
+
+        if kind == KIND_VEC {
+            unsafe {
+                let off = self.get_vec_pos();
+                let v = rebuild_vec(self.ptr.as_ptr(), self.len, self.cap, off);
+                let mut v = ManuallyDrop::new(v);
+                shrink_vec(&mut v, self.ptr.as_ptr(), self.len, capacity);
+                self.ptr = vptr(v.as_mut_ptr());
+                self.cap = v.capacity();
+                self.set_vec_pos(0);
+            }
+            return true;
+        }
+
+        debug_assert_eq!(kind, KIND_ARC);
+        let shared: *mut Shared = self.data;
+
+        unsafe {
+            if (*shared).is_unique() {
+                let v = &mut (*shared).vec;
+
+                shrink_vec(v, self.ptr.as_ptr(), self.len, capacity);
+                self.ptr = vptr(v.as_mut_ptr());
+                self.cap = v.capacity();
+                return true;
+            }
+        }
+
+        false
+    }
+
+    /// Attempts to shrink the buffer's capacity to fit its current length.
+    ///
+    /// This operation preserves the current contents and returns `true` when
+    /// the allocation is successfully shrunk. For `BytesMut` values backed by
+    /// shared storage, this only succeeds when the current handle is the only
+    /// remaining owner of the underlying allocation; otherwise it returns
+    /// `false` and leaves the buffer unchanged.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bytes::BytesMut;
+    ///
+    /// let mut buf = BytesMut::with_capacity(64);
+    /// buf.extend_from_slice(b"hello");
+    ///
+    /// assert!(buf.try_shrink_to_fit());
+    /// assert_eq!(buf.capacity(), 5);
+    /// assert_eq!(&buf[..], b"hello");
+    /// ```
+    pub fn try_shrink_to_fit(&mut self) -> bool {
+        self.try_shrink_to_capacity(self.len)
+    }
+
     /// Appends given bytes to this `BytesMut`.
     ///
     /// If this `BytesMut` object does not have enough capacity, it is resized
@@ -1868,6 +1946,17 @@ unsafe fn rebuild_vec(ptr: *mut u8, mut len: usize, mut cap: usize, off: usize)
     Vec::from_raw_parts(ptr, len, cap)
 }
 
+// `ptr..ptr + len` must be initialized bytes within `vec`'s allocation.
+unsafe fn shrink_vec(vec: &mut Vec<u8>, ptr: *mut u8, len: usize, capacity: usize) {
+    // Offset views may point past bytes that are not part of this buffer. Move
+    // the live range to the start before shrinking.
+    if len != 0 {
+        ptr::copy(ptr, vec.as_mut_ptr(), len);
+    }
+    vec.set_len(len);
+    vec.shrink_to(capacity);
+}
+
 // ===== impl SharedVtable =====
 
 static SHARED_VTABLE: Vtable = Vtable {

tests/test_buf_mut.rs

@@ -1,7 +1,7 @@
 #![warn(rust_2018_idioms)]
 
 use bytes::buf::UninitSlice;
-use bytes::{BufMut, BytesMut};
+use bytes::{Buf, BufMut, BytesMut};
 use core::fmt::Write;
 use core::mem::MaybeUninit;
 
@@ -282,3 +282,171 @@ fn test_bytes_mut_reuse() {
     let mut buf = BytesMut::new();
     buf.put(&[1u8, 2, 3] as &[u8]);
 }
+
+#[test]
+fn test_bytes_mut_try_shrink_to_fit_empty() {
+    let mut buf = BytesMut::new();
+
+    assert!(buf.try_shrink_to_fit());
+    assert!(buf.is_empty());
+    assert_eq!(buf.capacity(), 0);
+}
+
+#[test]
+fn test_bytes_mut_try_shrink_to_fit_already_tight() {
+    let mut buf = BytesMut::from(&b"abc"[..]);
+
+    assert!(buf.try_shrink_to_fit());
+    assert_eq!(buf.capacity(), buf.len());
+    assert_eq!(&buf[..], b"abc");
+}
+
+#[test]
+fn test_bytes_mut_try_shrink_to_fit_cant_shrink_with_refs() {
+    let mut buf = BytesMut::from(&b"abc"[..]);
+    buf.reserve(3);
+
+    let mut other = buf.split_off(2);
+
+    assert!(!buf.try_shrink_to_fit());
+    assert_eq!(buf.capacity(), 2);
+    assert_eq!(&buf[..], b"ab");
+
+    assert!(!other.try_shrink_to_fit());
+    assert_eq!(other.capacity(), 6);
+    assert_eq!(other.len(), 1);
+    assert_eq!(&other[..], b"c");
+}
+
+#[test]
+fn test_bytes_mut_try_shrink_to_fit_after_split_off() {
+    let data = b"hello world";
+    let mut buf = BytesMut::with_capacity(64);
+
+    assert_eq!(buf.capacity(), 64);
+    buf.extend_from_slice(data);
+    assert_eq!(buf.capacity(), 64);
+
+    let mut other = buf.split_off(5);
+    assert_eq!(&buf[..], b"hello");
+    assert_eq!(&other[..], b" world");
+    assert_eq!(buf.capacity(), 5);
+    assert_eq!(other.capacity(), 59);
+
+    drop(buf);
+    assert!(other.try_shrink_to_fit());
+    assert_eq!(other.capacity(), 6);
+}
+
+#[test]
+fn test_bytes_mut_try_shrink_to_fit_unique_front_after_split_off() {
+    let mut buf = BytesMut::with_capacity(64);
+    buf.extend_from_slice(b"hello world");
+
+    let other = buf.split_off(5);
+    drop(other);
+
+    assert!(buf.try_shrink_to_fit());
+    assert_eq!(buf.capacity(), buf.len());
+    assert_eq!(buf.capacity(), 5);
+    assert_eq!(&buf[..], b"hello");
+}
+
+#[test]
+fn test_bytes_mut_try_shrink_to_fit_unique_empty_tail_past_len() {
+    let mut buf = BytesMut::with_capacity(64);
+    buf.extend_from_slice(b"abc");
+
+    let mut other = buf.split_off(10);
+    drop(buf);
+
+    assert!(other.try_shrink_to_fit());
+    assert!(other.is_empty());
+    assert_eq!(other.capacity(), 0);
+}
+
+#[test]
+fn test_bytes_mut_try_shrink_to_fit_after_advance() {
+    let mut buf = BytesMut::with_capacity(64);
+    buf.extend_from_slice(b"hello world");
+    buf.advance(6);
+
+    assert!(buf.try_shrink_to_fit());
+    assert_eq!(buf.capacity(), buf.len());
+    assert_eq!(buf.capacity(), 5);
+    assert_eq!(&buf[..], b"world");
+}
+
+#[test]
+fn test_bytes_mut_try_shrink_to_fit_empty_after_advance() {
+    let mut buf = BytesMut::with_capacity(64);
+    buf.extend_from_slice(b"abc");
+    buf.advance(3);
+
+    assert!(buf.try_shrink_to_fit());
+    assert!(buf.is_empty());
+    assert_eq!(buf.capacity(), 0);
+}
+
+#[test]
+fn test_bytes_mut_try_shrink_to_fit_cant_shrink_with_bytes_ref() {
+    let mut buf = BytesMut::from(&b"abcdef"[..]);
+    buf.reserve(8);
+
+    let bytes = buf.split_off(3).freeze();
+
+    assert!(!buf.try_shrink_to_fit());
+    assert_eq!(buf.capacity(), 3);
+    assert_eq!(&buf[..], b"abc");
+    assert_eq!(&bytes[..], b"def");
+}
+
+#[test]
+fn test_bytes_mut_try_shrink_to_capacity() {
+    let mut buf = BytesMut::with_capacity(64);
+    buf.extend_from_slice(b"hello");
+
+    assert!(buf.try_shrink_to_capacity(16));
+    assert_eq!(buf.capacity(), 16);
+    assert_eq!(&buf[..], b"hello");
+}
+
+#[test]
+fn test_bytes_mut_try_shrink_to_capacity_does_not_shrink_below_len() {
+    let mut buf = BytesMut::with_capacity(64);
+    buf.extend_from_slice(b"hello world");
+
+    assert!(buf.try_shrink_to_capacity(5));
+    assert_eq!(buf.capacity(), buf.len());
+    assert_eq!(buf.capacity(), 11);
+    assert_eq!(&buf[..], b"hello world");
+}
+
+#[test]
+fn test_bytes_mut_try_shrink_to_capacity_after_advance() {
+    let mut buf = BytesMut::with_capacity(64);
+    buf.extend_from_slice(b"hello world");
+    buf.advance(6);
+
+    assert!(buf.try_shrink_to_capacity(16));
+    assert_eq!(buf.capacity(), 16);
+    assert_eq!(&buf[..], b"world");
+}
+
+#[test]
+fn test_bytes_mut_try_shrink_to_capacity_cant_shrink_with_refs() {
+    let mut buf = BytesMut::from(&b"abcdef"[..]);
+    buf.reserve(8);
+
+    let mut other = buf.split_off(3);
+    let buf_cap = buf.capacity();
+    let other_cap = other.capacity();
+
+    assert!(!buf.try_shrink_to_capacity(4));
+    assert_eq!(buf.capacity(), buf_cap);
+    assert_eq!(&buf[..], b"abc");
+
+    assert!(!other.try_shrink_to_capacity(4));
+    assert_eq!(other.capacity(), other_cap);
+    assert_eq!(&other[..], b"def");
+}