diff --git a/src/bytes_mut.rs b/src/bytes_mut.rs index 689f4b478..87efa5816 100644 --- a/src/bytes_mut.rs +++ b/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, 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 { diff --git a/tests/test_buf_mut.rs b/tests/test_buf_mut.rs index f1bd3b0a0..71771a378 100644 --- a/tests/test_buf_mut.rs +++ b/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"); +}