Skip to content

Commit

Permalink
Auto merge of rust-lang#126299 - scottmcm:tune-sliceindex-ubchecks, r…
Browse files Browse the repository at this point in the history
…=<try>

Remove superfluous UbChecks from `SliceIndex` methods

The current implementation calls the unsafe ones from the safe ones, but that means they end up emitting UbChecks that are impossible to hit, since we just checked those things.

This PR adds some new module-local helpers for the code shared between them, so the safe methods can be small enough to inline by avoiding those extra checks, while the unsafe methods still help catch length mistakes.

r? `@saethlin`
  • Loading branch information
bors committed Jun 12, 2024
2 parents 9a7bf4a + b0d8bf8 commit 280a96f
Show file tree
Hide file tree
Showing 8 changed files with 376 additions and 52 deletions.
122 changes: 90 additions & 32 deletions library/core/src/slice/index.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@
use crate::intrinsics::const_eval_select;
use crate::ops;
use crate::ptr;
use crate::ub_checks::assert_unsafe_precondition;

#[stable(feature = "rust1", since = "1.0.0")]
Expand Down Expand Up @@ -106,6 +105,51 @@ const fn slice_end_index_overflow_fail() -> ! {
panic!("attempted to index slice up to maximum usize");
}

// The UbChecks are great for catching bugs in the unsafe methods, but including
// them in safe indexing is unnecessary and hurts inlining and compile-time.
// Both the safe and unsafe public methods share these helpers,
// which use intrinsics directly to get *no* extra checks.

#[inline(always)]
const unsafe fn get_noubcheck<T>(ptr: *const [T], index: usize) -> *const T {
let ptr = ptr as *const T;
// SAFETY: The caller already checked these preconditions
unsafe { crate::intrinsics::offset(ptr, index) }
}

#[inline(always)]
const unsafe fn get_mut_noubcheck<T>(ptr: *mut [T], index: usize) -> *mut T {
let ptr = ptr as *mut T;
// SAFETY: The caller already checked these preconditions
unsafe { crate::intrinsics::offset(ptr, index) }
}

#[inline(always)]
const unsafe fn get_offset_len_noubcheck<T>(
ptr: *const [T],
offset: usize,
len: usize,
) -> *const [T] {
// SAFETY: The caller already checked these preconditions
unsafe {
let ptr = get_noubcheck(ptr, offset);
crate::intrinsics::aggregate_raw_ptr(ptr, len)
}
}

#[inline(always)]
const unsafe fn get_offset_len_mut_noubcheck<T>(
ptr: *mut [T],
offset: usize,
len: usize,
) -> *mut [T] {
// SAFETY: The caller already checked these preconditions
unsafe {
let ptr = get_mut_noubcheck(ptr, offset);
crate::intrinsics::aggregate_raw_ptr(ptr, len)
}
}

mod private_slice_index {
use super::ops;
#[stable(feature = "slice_get_slice", since = "1.28.0")]
Expand Down Expand Up @@ -203,13 +247,17 @@ unsafe impl<T> SliceIndex<[T]> for usize {
#[inline]
fn get(self, slice: &[T]) -> Option<&T> {
// SAFETY: `self` is checked to be in bounds.
if self < slice.len() { unsafe { Some(&*self.get_unchecked(slice)) } } else { None }
if self < slice.len() { unsafe { Some(&*get_noubcheck(slice, self)) } } else { None }
}

#[inline]
fn get_mut(self, slice: &mut [T]) -> Option<&mut T> {
// SAFETY: `self` is checked to be in bounds.
if self < slice.len() { unsafe { Some(&mut *self.get_unchecked_mut(slice)) } } else { None }
if self < slice.len() {
// SAFETY: `self` is checked to be in bounds.
unsafe { Some(&mut *get_mut_noubcheck(slice, self)) }
} else {
None
}
}

#[inline]
Expand All @@ -227,7 +275,7 @@ unsafe impl<T> SliceIndex<[T]> for usize {
// Use intrinsics::assume instead of hint::assert_unchecked so that we don't check the
// precondition of this function twice.
crate::intrinsics::assume(self < slice.len());
slice.as_ptr().add(self)
get_noubcheck(slice, self)
}
}

Expand All @@ -239,7 +287,7 @@ unsafe impl<T> SliceIndex<[T]> for usize {
(this: usize = self, len: usize = slice.len()) => this < len
);
// SAFETY: see comments for `get_unchecked` above.
unsafe { slice.as_mut_ptr().add(self) }
unsafe { get_mut_noubcheck(slice, self) }
}

#[inline]
Expand All @@ -265,7 +313,7 @@ unsafe impl<T> SliceIndex<[T]> for ops::IndexRange {
fn get(self, slice: &[T]) -> Option<&[T]> {
if self.end() <= slice.len() {
// SAFETY: `self` is checked to be valid and in bounds above.
unsafe { Some(&*self.get_unchecked(slice)) }
unsafe { Some(&*get_offset_len_noubcheck(slice, self.start(), self.len())) }
} else {
None
}
Expand All @@ -275,7 +323,7 @@ unsafe impl<T> SliceIndex<[T]> for ops::IndexRange {
fn get_mut(self, slice: &mut [T]) -> Option<&mut [T]> {
if self.end() <= slice.len() {
// SAFETY: `self` is checked to be valid and in bounds above.
unsafe { Some(&mut *self.get_unchecked_mut(slice)) }
unsafe { Some(&mut *get_offset_len_mut_noubcheck(slice, self.start(), self.len())) }
} else {
None
}
Expand All @@ -292,7 +340,7 @@ unsafe impl<T> SliceIndex<[T]> for ops::IndexRange {
// cannot be longer than `isize::MAX`. They also guarantee that
// `self` is in bounds of `slice` so `self` cannot overflow an `isize`,
// so the call to `add` is safe.
unsafe { ptr::slice_from_raw_parts(slice.as_ptr().add(self.start()), self.len()) }
unsafe { get_offset_len_noubcheck(slice, self.start(), self.len()) }
}

#[inline]
Expand All @@ -304,14 +352,14 @@ unsafe impl<T> SliceIndex<[T]> for ops::IndexRange {
);

// SAFETY: see comments for `get_unchecked` above.
unsafe { ptr::slice_from_raw_parts_mut(slice.as_mut_ptr().add(self.start()), self.len()) }
unsafe { get_offset_len_mut_noubcheck(slice, self.start(), self.len()) }
}

#[inline]
fn index(self, slice: &[T]) -> &[T] {
if self.end() <= slice.len() {
// SAFETY: `self` is checked to be valid and in bounds above.
unsafe { &*self.get_unchecked(slice) }
unsafe { &*get_offset_len_noubcheck(slice, self.start(), self.len()) }
} else {
slice_end_index_len_fail(self.end(), slice.len())
}
Expand All @@ -321,7 +369,7 @@ unsafe impl<T> SliceIndex<[T]> for ops::IndexRange {
fn index_mut(self, slice: &mut [T]) -> &mut [T] {
if self.end() <= slice.len() {
// SAFETY: `self` is checked to be valid and in bounds above.
unsafe { &mut *self.get_unchecked_mut(slice) }
unsafe { &mut *get_offset_len_mut_noubcheck(slice, self.start(), self.len()) }
} else {
slice_end_index_len_fail(self.end(), slice.len())
}
Expand All @@ -338,21 +386,26 @@ unsafe impl<T> SliceIndex<[T]> for ops::Range<usize> {

#[inline]
fn get(self, slice: &[T]) -> Option<&[T]> {
if self.start > self.end || self.end > slice.len() {
None
} else {
// Using checked_sub is a safe way to get `SubUnchecked` in MIR
if let Some(new_len) = usize::checked_sub(self.end, self.start)
&& self.end < slice.len()
{
// SAFETY: `self` is checked to be valid and in bounds above.
unsafe { Some(&*self.get_unchecked(slice)) }
unsafe { Some(&*get_offset_len_noubcheck(slice, self.start, new_len)) }
} else {
None
}
}

#[inline]
fn get_mut(self, slice: &mut [T]) -> Option<&mut [T]> {
if self.start > self.end || self.end > slice.len() {
None
} else {
if let Some(new_len) = usize::checked_sub(self.end, self.start)
&& self.end < slice.len()
{
// SAFETY: `self` is checked to be valid and in bounds above.
unsafe { Some(&mut *self.get_unchecked_mut(slice)) }
unsafe { Some(&mut *get_offset_len_mut_noubcheck(slice, self.start, new_len)) }
} else {
None
}
}

Expand All @@ -373,8 +426,10 @@ unsafe impl<T> SliceIndex<[T]> for ops::Range<usize> {
// `self` is in bounds of `slice` so `self` cannot overflow an `isize`,
// so the call to `add` is safe and the length calculation cannot overflow.
unsafe {
let new_len = self.end.unchecked_sub(self.start);
ptr::slice_from_raw_parts(slice.as_ptr().add(self.start), new_len)
// Using the intrinsic avoids a superfluous UB check,
// since the one on this method already checked `end >= start`.
let new_len = crate::intrinsics::unchecked_sub(self.end, self.start);
get_offset_len_noubcheck(slice, self.start, new_len)
}
}

Expand All @@ -391,31 +446,34 @@ unsafe impl<T> SliceIndex<[T]> for ops::Range<usize> {
);
// SAFETY: see comments for `get_unchecked` above.
unsafe {
let new_len = self.end.unchecked_sub(self.start);
ptr::slice_from_raw_parts_mut(slice.as_mut_ptr().add(self.start), new_len)
let new_len = crate::intrinsics::unchecked_sub(self.end, self.start);
get_offset_len_mut_noubcheck(slice, self.start, new_len)
}
}

#[inline(always)]
fn index(self, slice: &[T]) -> &[T] {
if self.start > self.end {
slice_index_order_fail(self.start, self.end);
} else if self.end > slice.len() {
// Using checked_sub is a safe way to get `SubUnchecked` in MIR
let Some(new_len) = usize::checked_sub(self.end, self.start) else {
slice_index_order_fail(self.start, self.end)
};
if self.end > slice.len() {
slice_end_index_len_fail(self.end, slice.len());
}
// SAFETY: `self` is checked to be valid and in bounds above.
unsafe { &*self.get_unchecked(slice) }
unsafe { &*get_offset_len_noubcheck(slice, self.start, new_len) }
}

#[inline]
fn index_mut(self, slice: &mut [T]) -> &mut [T] {
if self.start > self.end {
slice_index_order_fail(self.start, self.end);
} else if self.end > slice.len() {
let Some(new_len) = usize::checked_sub(self.end, self.start) else {
slice_index_order_fail(self.start, self.end)
};
if self.end > slice.len() {
slice_end_index_len_fail(self.end, slice.len());
}
// SAFETY: `self` is checked to be valid and in bounds above.
unsafe { &mut *self.get_unchecked_mut(slice) }
unsafe { &mut *get_offset_len_mut_noubcheck(slice, self.start, new_len) }
}
}

Expand Down
28 changes: 26 additions & 2 deletions tests/mir-opt/pre-codegen/slice_index.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
// skip-filecheck
//@ compile-flags: -O -C debuginfo=0 -Zmir-opt-level=2
//@ compile-flags: -O -C debuginfo=0 -Zmir-opt-level=2 -Z ub-checks=yes
// EMIT_MIR_FOR_EACH_PANIC_STRATEGY

#![crate_type = "lib"]
Expand All @@ -9,21 +8,39 @@ use std::ops::Range;

// EMIT_MIR slice_index.slice_index_usize.PreCodegen.after.mir
pub fn slice_index_usize(slice: &[u32], index: usize) -> u32 {
// CHECK-LABEL: slice_index_usize
// CHECK: [[LEN:_[0-9]+]] = Len((*_1))
// CHECK: Lt(_2, [[LEN]])
// CHECK-NOT: precondition_check
// CHECK: _0 = (*_1)[_2];
slice[index]
}

// EMIT_MIR slice_index.slice_get_mut_usize.PreCodegen.after.mir
pub fn slice_get_mut_usize(slice: &mut [u32], index: usize) -> Option<&mut u32> {
// CHECK-LABEL: slice_get_mut_usize
// CHECK: [[LEN:_[0-9]+]] = Len((*_1))
// CHECK: Lt(_2, move [[LEN]])
// CHECK-NOT: precondition_check
slice.get_mut(index)
}

// EMIT_MIR slice_index.slice_index_range.PreCodegen.after.mir
pub fn slice_index_range(slice: &[u32], index: Range<usize>) -> &[u32] {
// CHECK-LABEL: slice_index_range
&slice[index]
}

// EMIT_MIR slice_index.slice_get_unchecked_mut_range.PreCodegen.after.mir
pub unsafe fn slice_get_unchecked_mut_range(slice: &mut [u32], index: Range<usize>) -> &mut [u32] {
// CHECK-LABEL: slice_get_unchecked_mut_range
// CHECK: [[START:_[0-9]+]] = move (_2.0: usize);
// CHECK: [[END:_[0-9]+]] = move (_2.1: usize);
// CHECK: precondition_check
// CHECK: [[LEN:_[0-9]+]] = SubUnchecked([[END]], [[START]]);
// CHECK: [[PTR:_[0-9]+]] = Offset({{_[0-9]+}}, [[START]]);
// CHECK: [[SLICE:_[0-9]+]] = *mut [u32] from ([[PTR]], [[LEN]])
// CHECK: _0 = &mut (*[[SLICE]]);
slice.get_unchecked_mut(index)
}

Expand All @@ -32,5 +49,12 @@ pub unsafe fn slice_ptr_get_unchecked_range(
slice: *const [u32],
index: Range<usize>,
) -> *const [u32] {
// CHECK-LABEL: slice_ptr_get_unchecked_range
// CHECK: [[START:_[0-9]+]] = move (_2.0: usize);
// CHECK: [[END:_[0-9]+]] = move (_2.1: usize);
// CHECK: precondition_check
// CHECK: [[LEN:_[0-9]+]] = SubUnchecked([[END]], [[START]]);
// CHECK: [[PTR:_[0-9]+]] = Offset({{_[0-9]+}}, [[START]]);
// CHECK: _0 = *const [u32] from ([[PTR]], [[LEN]])
slice.get_unchecked(index)
}
Original file line number Diff line number Diff line change
Expand Up @@ -5,13 +5,54 @@ fn slice_get_mut_usize(_1: &mut [u32], _2: usize) -> Option<&mut u32> {
debug index => _2;
let mut _0: std::option::Option<&mut u32>;
scope 1 (inlined core::slice::<impl [u32]>::get_mut::<usize>) {
scope 2 (inlined <usize as SliceIndex<[u32]>>::get_mut) {
let mut _3: usize;
let mut _4: bool;
let mut _5: *mut [u32];
let mut _7: *mut u32;
let mut _8: &mut u32;
scope 3 (inlined core::slice::index::get_mut_noubcheck::<u32>) {
let _6: *mut u32;
scope 4 {
}
}
}
}

bb0: {
_0 = <usize as SliceIndex<[u32]>>::get_mut(move _2, move _1) -> [return: bb1, unwind unreachable];
StorageLive(_7);
StorageLive(_4);
StorageLive(_3);
_3 = Len((*_1));
_4 = Lt(_2, move _3);
switchInt(move _4) -> [0: bb1, otherwise: bb2];
}

bb1: {
StorageDead(_3);
_0 = const Option::<&mut u32>::None;
goto -> bb3;
}

bb2: {
StorageDead(_3);
StorageLive(_8);
StorageLive(_5);
_5 = &raw mut (*_1);
StorageLive(_6);
_6 = _5 as *mut u32 (PtrToPtr);
_7 = Offset(_6, _2);
StorageDead(_6);
StorageDead(_5);
_8 = &mut (*_7);
_0 = Option::<&mut u32>::Some(move _8);
StorageDead(_8);
goto -> bb3;
}

bb3: {
StorageDead(_4);
StorageDead(_7);
return;
}
}
Loading

0 comments on commit 280a96f

Please sign in to comment.