Skip to content

Commit

Permalink
Redo SliceIndex implementations
Browse files Browse the repository at this point in the history
  • Loading branch information
@scottmcm
scottmcm committed Jun 12, 2024
1 parent 4314bd7 commit b0d8bf8
Show file tree
Hide file tree
Showing 8 changed files with 375 additions and 51 deletions.
122 changes: 90 additions & 32 deletions library/core/src/slice/index.rs
View file
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
26 changes: 25 additions & 1 deletion tests/mir-opt/pre-codegen/slice_index.rs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
// skip-filecheck
//@ compile-flags: -O -C debuginfo=0 -Zmir-opt-level=2 -Z ub-checks=yes
// EMIT_MIR_FOR_EACH_PANIC_STRATEGY

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)
}
43 changes: 42 additions & 1 deletion tests/mir-opt/pre-codegen/slice_index.slice_get_mut_usize.PreCodegen.after.panic-abort.mir
View file
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 b0d8bf8

Please sign in to comment.