From 3d4290410f3698ecdcc351943d4a3d405fe5c6c9 Mon Sep 17 00:00:00 2001 From: Sebastian Hahn 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 { + 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()); +}