Skip to content

Commit

Permalink
Move to unchecked_sub as suggested by @the8472
Browse files Browse the repository at this point in the history
Signed-off-by: Alex Saveau <saveau.alexandre@gmail.com>
  • Loading branch information
SUPERCILEX committed Dec 31, 2023
1 parent 05a8576 commit d46ba1d
Show file tree
Hide file tree
Showing 6 changed files with 96 additions and 47 deletions.
8 changes: 4 additions & 4 deletions library/alloc/src/collections/vec_deque/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -756,8 +756,8 @@ impl<T, A: Allocator> VecDeque<T, A> {
let old_cap = self.capacity();

if new_cap > old_cap {
self.buf.reserve_exact(self.len, additional);
unsafe {
self.buf.reserve_exact(self.len, additional);
self.handle_capacity_increase(old_cap);
}
}
Expand Down Expand Up @@ -787,8 +787,8 @@ impl<T, A: Allocator> VecDeque<T, A> {
if new_cap > old_cap {
// we don't need to reserve_exact(), as the size doesn't have
// to be a power of 2.
self.buf.reserve(self.len, additional);
unsafe {
self.buf.reserve(self.len, additional);
self.handle_capacity_increase(old_cap);
}
}
Expand Down Expand Up @@ -838,8 +838,8 @@ impl<T, A: Allocator> VecDeque<T, A> {
let old_cap = self.capacity();

if new_cap > old_cap {
self.buf.try_reserve_exact(self.len, additional)?;
unsafe {
self.buf.try_reserve_exact(self.len, additional)?;
self.handle_capacity_increase(old_cap);
}
}
Expand Down Expand Up @@ -886,8 +886,8 @@ impl<T, A: Allocator> VecDeque<T, A> {
let old_cap = self.capacity();

if new_cap > old_cap {
self.buf.try_reserve(self.len, additional)?;
unsafe {
self.buf.try_reserve(self.len, additional)?;
self.handle_capacity_increase(old_cap);
}
}
Expand Down
50 changes: 32 additions & 18 deletions library/alloc/src/raw_vec.rs
Original file line number Diff line number Diff line change
Expand Up @@ -276,10 +276,6 @@ impl<T, A: Allocator> RawVec<T, A> {
/// *O*(1) behavior. Will limit this behavior if it would needlessly cause
/// itself to panic.
///
/// If `len` exceeds `self.capacity()`, this may fail to actually allocate
/// the requested space. This is not really unsafe, but the unsafe
/// code *you* write that relies on the behavior of this function may break.
///
/// This is ideal for implementing a bulk-push operation like `extend`.
///
/// # Panics
Expand All @@ -289,9 +285,13 @@ impl<T, A: Allocator> RawVec<T, A> {
/// # Aborts
///
/// Aborts on OOM.
///
/// # Safety
///
/// `len` must be less than or equal to the capacity of this [`RawVec`].
#[cfg(not(no_global_oom_handling))]
#[inline]
pub fn reserve(&mut self, len: usize, additional: usize) {
pub unsafe fn reserve(&mut self, len: usize, additional: usize) {
// Callers expect this function to be very cheap when there is already sufficient capacity.
// Therefore, we move all the resizing and error-handling logic from grow_amortized and
// handle_reserve behind a call, while making sure that this function is likely to be
Expand All @@ -305,7 +305,7 @@ impl<T, A: Allocator> RawVec<T, A> {
handle_reserve(slf.grow_amortized(len, additional));
}

if self.needs_to_grow(len, additional) {
if unsafe { self.needs_to_grow(len, additional) } {
do_reserve_and_handle(self, len, additional);
}
unsafe {
Expand All @@ -323,8 +323,16 @@ impl<T, A: Allocator> RawVec<T, A> {
}

/// The same as `reserve`, but returns on errors instead of panicking or aborting.
pub fn try_reserve(&mut self, len: usize, additional: usize) -> Result<(), TryReserveError> {
if self.needs_to_grow(len, additional) {
///
/// # Safety
///
/// `len` must be less than or equal to the capacity of this [`RawVec`].
pub unsafe fn try_reserve(
&mut self,
len: usize,
additional: usize,
) -> Result<(), TryReserveError> {
if unsafe { self.needs_to_grow(len, additional) } {
self.grow_amortized(len, additional)?;
}
unsafe {
Expand All @@ -340,29 +348,35 @@ impl<T, A: Allocator> RawVec<T, A> {
/// exactly the amount of memory necessary, but in principle the allocator
/// is free to give back more than we asked for.
///
/// If `len` exceeds `self.capacity()`, this may fail to actually allocate
/// the requested space. This is not really unsafe, but the unsafe code
/// *you* write that relies on the behavior of this function may break.
///
/// # Panics
///
/// Panics if the new capacity exceeds `isize::MAX` _bytes_.
///
/// # Aborts
///
/// Aborts on OOM.
///
/// # Safety
///
/// `len` must be less than or equal to the capacity of this [`RawVec`].
#[cfg(not(no_global_oom_handling))]
pub fn reserve_exact(&mut self, len: usize, additional: usize) {
handle_reserve(self.try_reserve_exact(len, additional));
pub unsafe fn reserve_exact(&mut self, len: usize, additional: usize) {
unsafe {
handle_reserve(self.try_reserve_exact(len, additional));
}
}

/// The same as `reserve_exact`, but returns on errors instead of panicking or aborting.
pub fn try_reserve_exact(
///
/// # Safety
///
/// `len` must be less than or equal to the capacity of this [`RawVec`].
pub unsafe fn try_reserve_exact(
&mut self,
len: usize,
additional: usize,
) -> Result<(), TryReserveError> {
if self.needs_to_grow(len, additional) {
if unsafe { self.needs_to_grow(len, additional) } {
self.grow_exact(len, additional)?;
}
unsafe {
Expand Down Expand Up @@ -391,8 +405,8 @@ impl<T, A: Allocator> RawVec<T, A> {
impl<T, A: Allocator> RawVec<T, A> {
/// Returns if the buffer needs to grow to fulfill the needed extra capacity.
/// Mainly used to make inlining reserve-calls possible without inlining `grow`.
pub(crate) fn needs_to_grow(&self, len: usize, additional: usize) -> bool {
additional > self.capacity().wrapping_sub(len)
pub(crate) unsafe fn needs_to_grow(&self, len: usize, additional: usize) -> bool {
unsafe { additional > self.capacity().unchecked_sub(len) }
}

/// # Safety:
Expand Down
53 changes: 37 additions & 16 deletions library/alloc/src/raw_vec/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,10 @@ fn allocator_param() {
let a = BoundedAlloc { fuel: Cell::new(500) };
let mut v: RawVec<u8, _> = RawVec::with_capacity_in(50, a);
assert_eq!(v.alloc.fuel.get(), 450);
v.reserve(50, 150); // (causes a realloc, thus using 50 + 150 = 200 units of fuel)
unsafe {
// (causes a realloc, thus using 50 + 150 = 200 units of fuel)
v.reserve(50, 150);
}
assert_eq!(v.alloc.fuel.get(), 250);
}

Expand All @@ -52,25 +55,35 @@ fn reserve_does_not_overallocate() {
{
let mut v: RawVec<u32> = RawVec::new();
// First, `reserve` allocates like `reserve_exact`.
v.reserve(0, 9);
unsafe {
v.reserve(0, 9);
}
assert_eq!(9, v.capacity());
}

{
let mut v: RawVec<u32> = RawVec::new();
v.reserve(0, 7);
unsafe {
v.reserve(0, 7);
}
assert_eq!(7, v.capacity());
// 97 is more than double of 7, so `reserve` should work
// like `reserve_exact`.
v.reserve(7, 90);
unsafe {
v.reserve(7, 90);
}
assert_eq!(97, v.capacity());
}

{
let mut v: RawVec<u32> = RawVec::new();
v.reserve(0, 12);
unsafe {
v.reserve(0, 12);
}
assert_eq!(12, v.capacity());
v.reserve(12, 3);
unsafe {
v.reserve(12, 3);
}
// 3 is less than half of 12, so `reserve` must grow
// exponentially. At the time of writing this test grow
// factor is 2, so new capacity is 24, however, grow factor
Expand Down Expand Up @@ -116,24 +129,28 @@ fn zst() {

// Check all these operations work as expected with zero-sized elements.

assert!(!v.needs_to_grow(100, usize::MAX - 100));
assert!(v.needs_to_grow(101, usize::MAX - 100));
assert!(unsafe { !v.needs_to_grow(100, usize::MAX - 100) });
assert!(unsafe { v.needs_to_grow(101, usize::MAX - 100) });
zst_sanity(&v);

v.reserve(100, usize::MAX - 100);
unsafe {
v.reserve(100, usize::MAX - 100);
}
//v.reserve(101, usize::MAX - 100); // panics, in `zst_reserve_panic` below
zst_sanity(&v);

v.reserve_exact(100, usize::MAX - 100);
unsafe {
v.reserve_exact(100, usize::MAX - 100);
}
//v.reserve_exact(101, usize::MAX - 100); // panics, in `zst_reserve_exact_panic` below
zst_sanity(&v);

assert_eq!(v.try_reserve(100, usize::MAX - 100), Ok(()));
assert_eq!(v.try_reserve(101, usize::MAX - 100), cap_err);
assert_eq!(unsafe { v.try_reserve(100, usize::MAX - 100) }, Ok(()));
assert_eq!(unsafe { v.try_reserve(101, usize::MAX - 100) }, cap_err);
zst_sanity(&v);

assert_eq!(v.try_reserve_exact(100, usize::MAX - 100), Ok(()));
assert_eq!(v.try_reserve_exact(101, usize::MAX - 100), cap_err);
assert_eq!(unsafe { v.try_reserve_exact(100, usize::MAX - 100) }, Ok(()));
assert_eq!(unsafe { v.try_reserve_exact(101, usize::MAX - 100) }, cap_err);
zst_sanity(&v);

assert_eq!(v.grow_amortized(100, usize::MAX - 100), cap_err);
Expand All @@ -151,7 +168,9 @@ fn zst_reserve_panic() {
let mut v: RawVec<ZST> = RawVec::new();
zst_sanity(&v);

v.reserve(101, usize::MAX - 100);
unsafe {
v.reserve(101, usize::MAX - 100);
}
}

#[test]
Expand All @@ -160,7 +179,9 @@ fn zst_reserve_exact_panic() {
let mut v: RawVec<ZST> = RawVec::new();
zst_sanity(&v);

v.reserve_exact(101, usize::MAX - 100);
unsafe {
v.reserve_exact(101, usize::MAX - 100);
}
}

#[test]
Expand Down
8 changes: 4 additions & 4 deletions library/alloc/src/string.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1152,7 +1152,7 @@ impl String {
self.vec.reserve(additional);
unsafe {
// Inform the optimizer that the reservation has succeeded or wasn't needed
intrinsics::assume(additional <= self.capacity().wrapping_sub(self.len()));
intrinsics::assume(additional <= self.capacity().unchecked_sub(self.len()));
}
}

Expand Down Expand Up @@ -1206,7 +1206,7 @@ impl String {
self.vec.reserve_exact(additional);
unsafe {
// Inform the optimizer that the reservation has succeeded or wasn't needed
intrinsics::assume(additional <= self.capacity().wrapping_sub(self.len()));
intrinsics::assume(additional <= self.capacity().unchecked_sub(self.len()));
}
}

Expand Down Expand Up @@ -1246,7 +1246,7 @@ impl String {
self.vec.try_reserve(additional)?;
unsafe {
// Inform the optimizer that the reservation has succeeded or wasn't needed
intrinsics::assume(additional <= self.capacity().wrapping_sub(self.len()));
intrinsics::assume(additional <= self.capacity().unchecked_sub(self.len()));
}
Ok(())
}
Expand Down Expand Up @@ -1293,7 +1293,7 @@ impl String {
self.vec.try_reserve_exact(additional)?;
unsafe {
// Inform the optimizer that the reservation has succeeded or wasn't needed
intrinsics::assume(additional <= self.capacity().wrapping_sub(self.len()));
intrinsics::assume(additional <= self.capacity().unchecked_sub(self.len()));
}
Ok(())
}
Expand Down
20 changes: 16 additions & 4 deletions library/alloc/src/vec/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -909,7 +909,10 @@ impl<T, A: Allocator> Vec<T, A> {
#[stable(feature = "rust1", since = "1.0.0")]
#[inline]
pub fn reserve(&mut self, additional: usize) {
self.buf.reserve(self.len, additional);
// SAFETY: len <= capacity
unsafe {
self.buf.reserve(self.len, additional);
}
unsafe {
// Inform the optimizer that the reservation has succeeded or wasn't needed
intrinsics::assume(!self.buf.needs_to_grow(self.len, additional));
Expand Down Expand Up @@ -944,7 +947,10 @@ impl<T, A: Allocator> Vec<T, A> {
#[stable(feature = "rust1", since = "1.0.0")]
#[inline]
pub fn reserve_exact(&mut self, additional: usize) {
self.buf.reserve_exact(self.len, additional);
// SAFETY: len <= capacity
unsafe {
self.buf.reserve_exact(self.len, additional);
}
unsafe {
// Inform the optimizer that the reservation has succeeded or wasn't needed
intrinsics::assume(!self.buf.needs_to_grow(self.len, additional));
Expand Down Expand Up @@ -986,7 +992,10 @@ impl<T, A: Allocator> Vec<T, A> {
#[stable(feature = "try_reserve", since = "1.57.0")]
#[inline]
pub fn try_reserve(&mut self, additional: usize) -> Result<(), TryReserveError> {
self.buf.try_reserve(self.len, additional)?;
// SAFETY: len <= capacity
unsafe {
self.buf.try_reserve(self.len, additional)?;
}
unsafe {
// Inform the optimizer that the reservation has succeeded or wasn't needed
intrinsics::assume(!self.buf.needs_to_grow(self.len, additional));
Expand Down Expand Up @@ -1035,7 +1044,10 @@ impl<T, A: Allocator> Vec<T, A> {
#[stable(feature = "try_reserve", since = "1.57.0")]
#[inline]
pub fn try_reserve_exact(&mut self, additional: usize) -> Result<(), TryReserveError> {
self.buf.try_reserve_exact(self.len, additional)?;
// SAFETY: len <= capacity
unsafe {
self.buf.try_reserve_exact(self.len, additional)?;
}
unsafe {
// Inform the optimizer that the reservation has succeeded or wasn't needed
intrinsics::assume(!self.buf.needs_to_grow(self.len, additional));
Expand Down
4 changes: 3 additions & 1 deletion library/alloc/src/vec/splice.rs
Original file line number Diff line number Diff line change
Expand Up @@ -126,7 +126,9 @@ impl<T, A: Allocator> Drain<'_, T, A> {
unsafe fn move_tail(&mut self, additional: usize) {
let vec = unsafe { self.vec.as_mut() };
let len = self.tail_start + self.tail_len;
vec.buf.reserve(len, additional);
unsafe {
vec.buf.reserve(len, additional);
}

let new_tail_start = self.tail_start + additional;
unsafe {
Expand Down

0 comments on commit d46ba1d

Please sign in to comment.