Auto merge of #126299 - scottmcm:tune-sliceindex-ubchecks, r=saethlin

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`

library/core/src/slice/index.rs

@@ -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")]
@@ -106,6 +105,47 @@ 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 debug runtime perf.
+// 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
+    let ptr = unsafe { 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
+    let ptr = unsafe { 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")]
@@ -203,13 +243,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]
@@ -227,7 +271,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)
         }
     }
 
@@ -239,7 +283,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]
@@ -265,7 +309,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
         }
@@ -275,7 +319,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
         }
@@ -292,7 +336,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]
@@ -304,14 +348,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())
         }
@@ -321,7 +365,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())
         }
@@ -338,21 +382,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
         }
     }
 
@@ -373,8 +422,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)
         }
     }
 
@@ -391,31 +442,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) }
     }
 }
 

tests/mir-opt/pre-codegen/slice_index.rs

@@ -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"]
@@ -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)
 }
 
@@ -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)
 }

tests/mir-opt/pre-codegen/slice_index.slice_get_mut_usize.PreCodegen.after.panic-abort.mir

@@ -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;
     }
 }