Auto merge of #139118 - scottmcm:slice-get-unchecked-intrinsic, r=workingjubilee

`slice.get(i)` should use a slice projection in MIR, like `slice[i]` does

`slice[i]` is built-in magic, so ends up being quite different from `slice.get(i)` in MIR, even though they're both doing nearly identical operations -- checking the length of the slice then getting a ref/ptr to the element if it's in-bounds.

This PR adds a `slice_get_unchecked` intrinsic for `impl SliceIndex for usize` to use to fix that, so it no longer needs to do a bunch of lines of pointer math and instead just gets the obvious single statement.  (This is *not* used for the range versions, since `slice[i..]` and `slice[..k]` can't use the mir Slice projection as they're using fenceposts, not indices.)

I originally tried to do this with some kind of GVN pattern, but realized that I'm pretty sure it's not legal to optimize `BinOp::Offset` to `PlaceElem::Index` without an extremely complicated condition.  Basically, the problem is that the `Index` projection on a dereferenced slice pointer *cares about the metadata*, since it's UB to `PlaceElem::Index` outside the range described by the metadata.  But then you cast the fat pointer to a thin pointer then offset it, that *ignores* the slice length metadata, so it's possible to write things that are legal with `Offset` but would be UB if translated in the obvious way to `Index`.  Checking (or even determining) the necessary conditions for that would be complicated and error-prone, whereas this intrinsic-based approach is quite straight-forward.

Zero backend changes, because it just lowers to MIR, so it's already supported naturally by CTFE/Miri/cg_llvm/cg_clif.

Author: bors

compiler/rustc_hir_analysis/src/check/intrinsic.rs

@@ -265,6 +265,7 @@ pub(crate) fn check_intrinsic_type(
                 vec![Ty::new_imm_ptr(tcx, param(0)), tcx.types.isize],
                 Ty::new_imm_ptr(tcx, param(0)),
             ),
+            sym::slice_get_unchecked => (3, 0, vec![param(1), tcx.types.usize], param(0)),
             sym::ptr_mask => (
                 1,
                 0,

compiler/rustc_mir_transform/src/lower_intrinsics.rs

@@ -262,6 +262,52 @@ impl<'tcx> crate::MirPass<'tcx> for LowerIntrinsics {
                         });
                         terminator.kind = TerminatorKind::Goto { target };
                     }
+                    sym::slice_get_unchecked => {
+                        let target = target.unwrap();
+                        let Ok([ptrish, index]) = take_array(args) else {
+                            span_bug!(
+                                terminator.source_info.span,
+                                "Wrong number of arguments for {intrinsic:?}",
+                            );
+                        };
+
+                        let place = ptrish.node.place().unwrap();
+                        assert!(!place.is_indirect());
+                        let updated_place = place.project_deeper(
+                            &[
+                                ProjectionElem::Deref,
+                                ProjectionElem::Index(
+                                    index.node.place().unwrap().as_local().unwrap(),
+                                ),
+                            ],
+                            tcx,
+                        );
+
+                        let ret_ty = generic_args.type_at(0);
+                        let rvalue = match *ret_ty.kind() {
+                            ty::RawPtr(_, Mutability::Not) => {
+                                Rvalue::RawPtr(RawPtrKind::Const, updated_place)
+                            }
+                            ty::RawPtr(_, Mutability::Mut) => {
+                                Rvalue::RawPtr(RawPtrKind::Mut, updated_place)
+                            }
+                            ty::Ref(region, _, Mutability::Not) => {
+                                Rvalue::Ref(region, BorrowKind::Shared, updated_place)
+                            }
+                            ty::Ref(region, _, Mutability::Mut) => Rvalue::Ref(
+                                region,
+                                BorrowKind::Mut { kind: MutBorrowKind::Default },
+                                updated_place,
+                            ),
+                            _ => bug!("Unknown return type {ret_ty:?}"),
+                        };
+
+                        block.statements.push(Statement {
+                            source_info: terminator.source_info,
+                            kind: StatementKind::Assign(Box::new((*destination, rvalue))),
+                        });
+                        terminator.kind = TerminatorKind::Goto { target };
+                    }
                     sym::transmute | sym::transmute_unchecked => {
                         let dst_ty = destination.ty(local_decls, tcx).ty;
                         let Ok([arg]) = take_array(args) else {

compiler/rustc_span/src/symbol.rs

@@ -2000,6 +2000,7 @@ symbols! {
         slice,
         slice_from_raw_parts,
         slice_from_raw_parts_mut,
+        slice_get_unchecked,
         slice_into_vec,
         slice_iter,
         slice_len_fn,

library/core/src/intrinsics/bounds.rs

@@ -0,0 +1,39 @@
+//! Various traits used to restrict intrinsics to not-completely-wrong types.
+
+/// Types with a built-in dereference operator in runtime MIR,
+/// aka references and raw pointers.
+///
+/// # Safety
+/// Must actually *be* such a type.
+pub unsafe trait BuiltinDeref: Sized {
+    type Pointee: ?Sized;
+}
+
+unsafe impl<T: ?Sized> BuiltinDeref for &mut T {
+    type Pointee = T;
+}
+unsafe impl<T: ?Sized> BuiltinDeref for &T {
+    type Pointee = T;
+}
+unsafe impl<T: ?Sized> BuiltinDeref for *mut T {
+    type Pointee = T;
+}
+unsafe impl<T: ?Sized> BuiltinDeref for *const T {
+    type Pointee = T;
+}
+
+pub trait ChangePointee<U: ?Sized>: BuiltinDeref {
+    type Output;
+}
+impl<'a, T: ?Sized + 'a, U: ?Sized + 'a> ChangePointee<U> for &'a mut T {
+    type Output = &'a mut U;
+}
+impl<'a, T: ?Sized + 'a, U: ?Sized + 'a> ChangePointee<U> for &'a T {
+    type Output = &'a U;
+}
+impl<T: ?Sized, U: ?Sized> ChangePointee<U> for *mut T {
+    type Output = *mut U;
+}
+impl<T: ?Sized, U: ?Sized> ChangePointee<U> for *const T {
+    type Output = *const U;
+}

library/core/src/intrinsics/mod.rs

@@ -53,6 +53,7 @@
 use crate::marker::{ConstParamTy, DiscriminantKind, Tuple};
 use crate::ptr;
 
+mod bounds;
 pub mod fallback;
 pub mod mir;
 pub mod simd;
@@ -1730,7 +1731,7 @@ pub const fn needs_drop<T: ?Sized>() -> bool;
 #[rustc_intrinsic_const_stable_indirect]
 #[rustc_nounwind]
 #[rustc_intrinsic]
-pub const unsafe fn offset<Ptr, Delta>(dst: Ptr, offset: Delta) -> Ptr;
+pub const unsafe fn offset<Ptr: bounds::BuiltinDeref, Delta>(dst: Ptr, offset: Delta) -> Ptr;
 
 /// Calculates the offset from a pointer, potentially wrapping.
 ///
@@ -1751,6 +1752,33 @@ pub const unsafe fn offset<Ptr, Delta>(dst: Ptr, offset: Delta) -> Ptr;
 #[rustc_intrinsic]
 pub const unsafe fn arith_offset<T>(dst: *const T, offset: isize) -> *const T;
 
+/// Projects to the `index`-th element of `slice_ptr`, as the same kind of pointer
+/// as the slice was provided -- so `&mut [T] → &mut T`, `&[T] → &T`,
+/// `*mut [T] → *mut T`, or `*const [T] → *const T` -- without a bounds check.
+///
+/// This is exposed via `<usize as SliceIndex>::get(_unchecked)(_mut)`,
+/// and isn't intended to be used elsewhere.
+///
+/// Expands in MIR to `{&, &mut, &raw const, &raw mut} (*slice_ptr)[index]`,
+/// depending on the types involved, so no backend support is needed.
+///
+/// # Safety
+///
+/// - `index < PtrMetadata(slice_ptr)`, so the indexing is in-bounds for the slice
+/// - the resulting offsetting is in-bounds of the allocated object, which is
+///   always the case for references, but needs to be upheld manually for pointers
+#[cfg(not(bootstrap))]
+#[rustc_nounwind]
+#[rustc_intrinsic]
+pub const unsafe fn slice_get_unchecked<
+    ItemPtr: bounds::ChangePointee<[T], Pointee = T, Output = SlicePtr>,
+    SlicePtr,
+    T,
+>(
+    slice_ptr: SlicePtr,
+    index: usize,
+) -> ItemPtr;
+
 /// Masks out bits of the pointer according to a mask.
 ///
 /// Note that, unlike most intrinsics, this is safe to call;
@@ -3575,18 +3603,9 @@ pub const fn type_id<T: ?Sized + 'static>() -> u128;
 #[unstable(feature = "core_intrinsics", issue = "none")]
 #[rustc_intrinsic_const_stable_indirect]
 #[rustc_intrinsic]
-pub const fn aggregate_raw_ptr<P: AggregateRawPtr<D, Metadata = M>, D, M>(data: D, meta: M) -> P;
-
-#[unstable(feature = "core_intrinsics", issue = "none")]
-pub trait AggregateRawPtr<D> {
-    type Metadata: Copy;
-}
-impl<P: ?Sized, T: ptr::Thin> AggregateRawPtr<*const T> for *const P {
-    type Metadata = <P as ptr::Pointee>::Metadata;
-}
-impl<P: ?Sized, T: ptr::Thin> AggregateRawPtr<*mut T> for *mut P {
-    type Metadata = <P as ptr::Pointee>::Metadata;
-}
+pub const fn aggregate_raw_ptr<P: bounds::BuiltinDeref, D, M>(data: D, meta: M) -> P
+where
+    <P as bounds::BuiltinDeref>::Pointee: ptr::Pointee<Metadata = M>;
 
 /// Lowers in MIR to `Rvalue::UnaryOp` with `UnOp::PtrMetadata`.
 ///

library/core/src/slice/index.rs

@@ -1,5 +1,7 @@
 //! Indexing implementations for `[T]`.
 
+#[cfg(not(bootstrap))]
+use crate::intrinsics::slice_get_unchecked;
 use crate::panic::const_panic;
 use crate::ub_checks::assert_unsafe_precondition;
 use crate::{ops, range};
@@ -83,13 +85,15 @@ const fn slice_end_index_overflow_fail() -> ! {
 // Both the safe and unsafe public methods share these helpers,
 // which use intrinsics directly to get *no* extra checks.
 
+#[cfg(bootstrap)]
 #[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) }
 }
 
+#[cfg(bootstrap)]
 #[inline(always)]
 const unsafe fn get_mut_noubcheck<T>(ptr: *mut [T], index: usize) -> *mut T {
     let ptr = ptr as *mut T;
@@ -103,8 +107,9 @@ const unsafe fn get_offset_len_noubcheck<T>(
     offset: usize,
     len: usize,
 ) -> *const [T] {
+    let ptr = ptr as *const T;
     // SAFETY: The caller already checked these preconditions
-    let ptr = unsafe { get_noubcheck(ptr, offset) };
+    let ptr = unsafe { crate::intrinsics::offset(ptr, offset) };
     crate::intrinsics::aggregate_raw_ptr(ptr, len)
 }
 
@@ -114,8 +119,9 @@ const unsafe fn get_offset_len_mut_noubcheck<T>(
     offset: usize,
     len: usize,
 ) -> *mut [T] {
+    let ptr = ptr as *mut T;
     // SAFETY: The caller already checked these preconditions
-    let ptr = unsafe { get_mut_noubcheck(ptr, offset) };
+    let ptr = unsafe { crate::intrinsics::offset(ptr, offset) };
     crate::intrinsics::aggregate_raw_ptr(ptr, len)
 }
 
@@ -224,15 +230,35 @@ 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(&*get_noubcheck(slice, self)) } } else { None }
+        if self < slice.len() {
+            #[cfg(bootstrap)]
+            // SAFETY: `self` is checked to be in bounds.
+            unsafe {
+                Some(&*get_noubcheck(slice, self))
+            }
+            #[cfg(not(bootstrap))]
+            // SAFETY: `self` is checked to be in bounds.
+            unsafe {
+                Some(slice_get_unchecked(slice, self))
+            }
+        } else {
+            None
+        }
     }
 
     #[inline]
     fn get_mut(self, slice: &mut [T]) -> Option<&mut T> {
         if self < slice.len() {
+            #[cfg(bootstrap)]
+            // SAFETY: `self` is checked to be in bounds.
+            unsafe {
+                Some(&mut *get_mut_noubcheck(slice, self))
+            }
+            #[cfg(not(bootstrap))]
             // SAFETY: `self` is checked to be in bounds.
-            unsafe { Some(&mut *get_mut_noubcheck(slice, self)) }
+            unsafe {
+                Some(slice_get_unchecked(slice, self))
+            }
         } else {
             None
         }
@@ -254,7 +280,14 @@ 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());
-            get_noubcheck(slice, self)
+            #[cfg(bootstrap)]
+            {
+                get_noubcheck(slice, self)
+            }
+            #[cfg(not(bootstrap))]
+            {
+                slice_get_unchecked(slice, self)
+            }
         }
     }
 
@@ -267,7 +300,16 @@ unsafe impl<T> SliceIndex<[T]> for usize {
             (this: usize = self, len: usize = slice.len()) => this < len
         );
         // SAFETY: see comments for `get_unchecked` above.
-        unsafe { get_mut_noubcheck(slice, self) }
+        unsafe {
+            #[cfg(bootstrap)]
+            {
+                get_mut_noubcheck(slice, self)
+            }
+            #[cfg(not(bootstrap))]
+            {
+                slice_get_unchecked(slice, self)
+            }
+        }
     }
 
     #[inline]

tests/mir-opt/lower_intrinsics.rs

@@ -263,3 +263,24 @@ pub fn get_metadata(a: *const i32, b: *const [u8], c: *const dyn std::fmt::Debug
     let _usize = ptr_metadata(b);
     let _vtable = ptr_metadata(c);
 }
+
+// EMIT_MIR lower_intrinsics.slice_get.LowerIntrinsics.diff
+pub unsafe fn slice_get<'a, 'b>(
+    r: &'a [i8],
+    rm: &'b mut [i16],
+    p: *const [i32],
+    pm: *mut [i64],
+    i: usize,
+) -> (&'a i8, &'b mut i16, *const i32, *mut i64) {
+    use std::intrinsics::slice_get_unchecked;
+    // CHECK: = &(*_{{[0-9]+}})[_{{[0-9]+}}]
+    // CHECK: = &mut (*_{{[0-9]+}})[_{{[0-9]+}}]
+    // CHECK: = &raw const (*_{{[0-9]+}})[_{{[0-9]+}}]
+    // CHECK: = &raw mut (*_{{[0-9]+}})[_{{[0-9]+}}]
+    (
+        slice_get_unchecked(r, i),
+        slice_get_unchecked(rm, i),
+        slice_get_unchecked(p, i),
+        slice_get_unchecked(pm, i),
+    )
+}