From 26c14f000d42f86b8642c75604a4a93dbe603e1b Mon Sep 17 00:00:00 2001 From: Sebastian Hahn Date: Mon, 10 Jun 2024 00:15:45 +0200 Subject: [PATCH] Allow copying bytes in `BytesMut::try_reclaim` This has the nice side effect of making the implementation simpler with regards to unsafe code, as it mostly follows what `reserve` does. --- src/bytes_mut.rs | 94 +++++++++++++-------------------------------- tests/test_bytes.rs | 12 +++++- 2 files changed, 36 insertions(+), 70 deletions(-) diff --git a/src/bytes_mut.rs b/src/bytes_mut.rs index b00ea03b6..cb7d16bc4 100644 --- a/src/bytes_mut.rs +++ b/src/bytes_mut.rs @@ -589,12 +589,13 @@ impl BytesMut { return; } - self.reserve_inner(additional); + // will always succeed + let _ = self.reserve_inner(additional, true); } - // In separate function to allow the short-circuits in `reserve` to - // be inline-able. Significantly helps performance. - fn reserve_inner(&mut self, additional: usize) { + // In separate function to allow the short-circuits in `reserve` and `try_reclaim` to + // be inline-able. Significantly helps performance. Returns false if it did not succeed. + fn reserve_inner(&mut self, additional: usize, allocate: bool) -> bool { let len = self.len(); let kind = self.kind(); @@ -639,6 +640,9 @@ impl BytesMut { // can gain capacity back. self.cap += off; } else { + if !allocate { + return false; + } // Not enough space, or reusing might be too much overhead: // allocate more space! let mut v = @@ -651,7 +655,7 @@ impl BytesMut { debug_assert_eq!(self.len, v.len() - off); } - return; + return true; } } @@ -693,6 +697,9 @@ impl BytesMut { self.ptr = vptr(ptr); self.cap = v.capacity(); } else { + if !allocate { + return false; + } // calculate offset let off = (self.ptr.as_ptr() as usize) - (v.as_ptr() as usize); @@ -731,9 +738,12 @@ impl BytesMut { self.cap = v.capacity() - off; } - return; + return true; } } + if !allocate { + return false; + } let original_capacity_repr = unsafe { (*shared).original_capacity_repr }; let original_capacity = original_capacity_from_repr(original_capacity_repr); @@ -756,21 +766,23 @@ impl BytesMut { self.ptr = vptr(v.as_mut_ptr()); self.cap = v.capacity(); debug_assert_eq!(self.len, v.len()); + return true; } - /// Attempts to cheaply reclaim capacity for at least `additional` more bytes to be inserted - /// into the given `BytesMut` and returns `true` if it succeeded. + /// Attempts to cheaply reclaim already allocated capacity for at least `additional` more + /// bytes to be inserted into the given `BytesMut` and returns `true` if it succeeded. + /// + /// `try_reclaim` behaves exactly like `reserve`, except that it never allocates new storage + /// and returns a `bool` indicating whether it was successful in doing so: /// /// `try_reclaim` returns false under these conditions: /// - The spare capacity left is less than `additional` bytes AND /// - The existing allocation cannot be reclaimed cheaply or it was less than /// `additional` bytes in size /// - /// Reclaiming the allocation cheaply is possible if the `BytesMut` is currently empty and - /// there are no outstanding references through other `BytesMut`s or `Bytes` which point to the - /// same underlying storage. - /// - /// A call to `try_reclaim` never copies inside the allocation nor allocates new storage. + /// Reclaiming the allocation cheaply is possible if the `BytesMut` has no outstanding + /// references through other `BytesMut`s or `Bytes` which point to the same underlying + /// storage. /// /// # Examples /// @@ -814,61 +826,7 @@ impl BytesMut { return true; } - if !self.is_empty() { - // try_reclaim never copies any bytes but there are some bytes stored already - return false; - } - - let kind = self.kind(); - if kind == KIND_VEC { - // Safety: self is of KIND_VEC, so calling get_vec_pos() is safe. - let off = unsafe { self.get_vec_pos() }; - // The whole allocation is too small to fit the request - if additional > rem + off { - return false; - } - - // 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; - - // If there are other references to the underlying storage, we cannot reclaim it - // - // Safety: self is of type KIND_ARC, so the Shared ptr is valid - if unsafe { !(*shared).is_unique() } { - return false; - } - - // Safety: self is of type KIND_ARC and there are no other handles alive, so we - // can get a mut reference to the underlying storageq - let v = unsafe { &mut (*shared).vec }; - let cap = v.capacity(); - if additional > cap { - // The underlying storage does not have enough capacity - 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 + self.reserve_inner(additional, false) } /// Appends given bytes to this `BytesMut`. diff --git a/tests/test_bytes.rs b/tests/test_bytes.rs index 2fb547ec5..a20156345 100644 --- a/tests/test_bytes.rs +++ b/tests/test_bytes.rs @@ -1310,7 +1310,13 @@ fn try_reserve_vec() { let mut buf = BytesMut::with_capacity(6); buf.put_slice(b"abc"); assert_eq!(false, buf.try_reclaim(6)); - buf.advance(3); + buf.advance(2); + assert_eq!(4, buf.capacity()); + // We can reclaim 5 bytes, because the byte in the buffer can be moved to the front. 6 bytes + // cannot be reclaimed because there is already one byte stored + assert_eq!(false, buf.try_reclaim(6)); + assert_eq!(true, buf.try_reclaim(5)); + buf.advance(1); assert_eq!(true, buf.try_reclaim(6)); assert_eq!(6, buf.capacity()); } @@ -1336,6 +1342,8 @@ fn try_reserve_arc() { assert_eq!(6, buf.capacity()); assert_eq!(6, buf.len()); assert_eq!(false, buf.try_reclaim(6)); - buf.advance(6); + buf.advance(4); + assert_eq!(true, buf.try_reclaim(4)); + buf.advance(2); assert_eq!(true, buf.try_reclaim(6)); }