Redo SliceIndex implementations

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

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

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