Add slice_index intrinsic

scottmcm · scottmcm · commit 2d383adcbe4e · 2023-10-21T18:29:46.000-07:00
diff --git a/compiler/rustc_hir_analysis/src/check/intrinsic.rs b/compiler/rustc_hir_analysis/src/check/intrinsic.rs
@@ -216,6 +216,7 @@ pub fn check_intrinsic_type(tcx: TyCtxt<'_>, it: &hir::ForeignItem<'_>) {
 
             sym::type_name => (1, Vec::new(), Ty::new_static_str(tcx)),
             sym::type_id => (1, Vec::new(), tcx.types.u128),
+            sym::slice_index => (2, vec![param(0), tcx.types.usize], param(1)),
             sym::offset => (2, vec![param(0), param(1)], param(0)),
             sym::arith_offset => (
                 1,
diff --git a/compiler/rustc_mir_transform/src/lower_intrinsics.rs b/compiler/rustc_mir_transform/src/lower_intrinsics.rs
@@ -234,6 +234,43 @@ impl<'tcx> MirPass<'tcx> for LowerIntrinsics {
                             terminator.kind = TerminatorKind::Goto { target };
                         }
                     }
+                    sym::slice_index => {
+                        let target = target.unwrap();
+
+                        let elem =
+                            if let Some(slice_ptr) = args[0].place()
+                            && let Some(index_place) = args[1].place()
+                            && let Some(index) = index_place.as_local()
+                            {
+                                let slice = tcx.mk_place_deref(slice_ptr);
+                                tcx.mk_place_index(slice, index)
+                            } else {
+                                span_bug!(
+                                    terminator.source_info.span,
+                                    "Invalid call to slice_index"
+                                );
+                            };
+
+                        let elem_ptr =
+                            match generic_args.type_at(1).kind() {
+                                ty::Ref(region, _, Mutability::Mut) =>
+                                    Rvalue::Ref(*region, BorrowKind::Mut { kind: MutBorrowKind::Default }, elem),
+                                ty::Ref(region, _, Mutability::Not) =>
+                                    Rvalue::Ref(*region, BorrowKind::Shared, elem),
+                                ty::RawPtr(ty_and_mut) =>
+                                    Rvalue::AddressOf(ty_and_mut.mutbl, elem),
+                                x => bug!("Invalid return type for slice_index: {x:?}"),
+                            };
+
+                        block.statements.push(Statement {
+                            source_info: terminator.source_info,
+                            kind: StatementKind::Assign(Box::new((
+                                *destination,
+                                elem_ptr,
+                            ))),
+                        });
+                        terminator.kind = TerminatorKind::Goto { target };
+                    }
                     sym::offset => {
                         let target = target.unwrap();
                         let Ok([ptr, delta]) = <[_; 2]>::try_from(std::mem::take(args)) else {
diff --git a/compiler/rustc_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs
@@ -1528,6 +1528,7 @@ symbols! {
         slice,
         slice_from_raw_parts,
         slice_from_raw_parts_mut,
+        slice_index,
         slice_len_fn,
         slice_patterns,
         slicing_syntax,
diff --git a/library/core/src/intrinsics.rs b/library/core/src/intrinsics.rs
@@ -74,6 +74,7 @@ pub unsafe fn drop_in_place<T: ?Sized>(to_drop: *mut T) {
     unsafe { crate::ptr::drop_in_place(to_drop) }
 }
 
+#[allow(private_bounds)]
 extern "rust-intrinsic" {
     // N.B., these intrinsics take raw pointers because they mutate aliased
     // memory, which is not valid for either `&` or `&mut`.
@@ -1410,6 +1411,39 @@ extern "rust-intrinsic" {
     #[rustc_nounwind]
     pub fn needs_drop<T: ?Sized>() -> bool;
 
+    /// Indexes into a slice using `ProjectionElem::Index`
+    ///
+    /// Depending on the return type, this lowers to one of the following in MIR:
+    /// - `&raw const (*slice)[index]`, for `*const _`
+    /// - `&raw mut (*slice)[index]`, for `*mut _`
+    /// - `&(*slice)[index]`, for `&_`
+    /// - `&mut (*slice)[index]`, for `&mut _`
+    ///
+    /// This is intended for use in implementing `[T]::get_unchecked_mut`,
+    /// `[T]::get`, and friends. It should not be used elsewhere; ICEs from
+    /// misuse in callers of intrinsics are *not* considered compiler bugs.
+    ///
+    /// While semantically unnecessary, as it could be done via `offset`, this
+    /// reduces and simplifies the MIR produced for common operations.
+    ///
+    /// # Safety
+    ///
+    /// Requires `index < slice.len()` and that `slice` points to a valid slice.
+    ///
+    /// If the return type is a reference (mutable or shared), then you must
+    /// ensure that its lifetime is no longer than the input's.
+    #[cfg(not(bootstrap))]
+    #[must_use = "returns a new pointer rather than modifying its argument"]
+    #[rustc_const_stable(feature = "const_ptr_offset", since = "1.61.0")]
+    #[rustc_nounwind]
+    pub fn slice_index<
+        Slice: PointerOrReference<To = [Item::To]>,
+        Item: PointerOrReference<To: Sized>,
+    >(
+        slice: Slice,
+        index: usize,
+    ) -> Item;
+
     /// Calculates the offset from a pointer.
     ///
     /// This is implemented as an intrinsic to avoid converting to and from an
@@ -1430,7 +1464,7 @@ extern "rust-intrinsic" {
     #[must_use = "returns a new pointer rather than modifying its argument"]
     #[rustc_const_stable(feature = "const_ptr_offset", since = "1.61.0")]
     #[rustc_nounwind]
-    pub fn offset<Ptr, Delta>(dst: Ptr, offset: Delta) -> Ptr;
+    pub fn offset<Ptr: PointerOrReference<To: Sized>, Delta>(dst: Ptr, offset: Delta) -> Ptr;
 
     /// Calculates the offset from a pointer, potentially wrapping.
     ///
@@ -2855,3 +2889,20 @@ pub const unsafe fn write_bytes<T>(dst: *mut T, val: u8, count: usize) {
         write_bytes(dst, val, count)
     }
 }
+
+/// No semantic meaning; just to help reduce mistakes with generic intrinsics
+trait PointerOrReference {
+    type To: ?Sized;
+}
+impl<T: ?Sized> PointerOrReference for &T {
+    type To = T;
+}
+impl<T: ?Sized> PointerOrReference for &mut T {
+    type To = T;
+}
+impl<T: ?Sized> PointerOrReference for *const T {
+    type To = T;
+}
+impl<T: ?Sized> PointerOrReference for *mut T {
+    type To = T;
+}
diff --git a/tests/mir-opt/lower_intrinsics.index_project.LowerIntrinsics.panic-abort.diff b/tests/mir-opt/lower_intrinsics.index_project.LowerIntrinsics.panic-abort.diff
@@ -0,0 +1,97 @@
+- // MIR for `index_project` before LowerIntrinsics
++ // MIR for `index_project` after LowerIntrinsics
+  
+  fn index_project(_1: &mut [i32], _2: &[i32], _3: *mut [i32], _4: *const [i32], _5: usize) -> () {
+      debug a => _1;
+      debug b => _2;
+      debug c => _3;
+      debug d => _4;
+      debug i => _5;
+      let mut _0: ();
+      let _6: &mut i32;
+      let mut _7: &mut [i32];
+      let mut _8: usize;
+      let mut _10: &[i32];
+      let mut _11: usize;
+      let mut _13: *mut [i32];
+      let mut _14: usize;
+      let mut _16: *const [i32];
+      let mut _17: usize;
+      scope 1 {
+          debug _a => _6;
+          let _9: &i32;
+          scope 2 {
+              debug _b => _9;
+              let _12: *mut i32;
+              scope 3 {
+                  debug _c => _12;
+                  let _15: *const i32;
+                  scope 4 {
+                      debug _d => _15;
+                  }
+              }
+          }
+      }
+  
+      bb0: {
+          StorageLive(_6);
+          StorageLive(_7);
+          _7 = move _1;
+          StorageLive(_8);
+          _8 = _5;
+-         _6 = slice_index::<&mut [i32], &mut i32>(move _7, move _8) -> [return: bb1, unwind unreachable];
++         _6 = &mut (*_7)[_8];
++         goto -> bb1;
+      }
+  
+      bb1: {
+          StorageDead(_8);
+          StorageDead(_7);
+          StorageLive(_9);
+          StorageLive(_10);
+          _10 = _2;
+          StorageLive(_11);
+          _11 = _5;
+-         _9 = slice_index::<&[i32], &i32>(move _10, move _11) -> [return: bb2, unwind unreachable];
++         _9 = &(*_10)[_11];
++         goto -> bb2;
+      }
+  
+      bb2: {
+          StorageDead(_11);
+          StorageDead(_10);
+          StorageLive(_12);
+          StorageLive(_13);
+          _13 = _3;
+          StorageLive(_14);
+          _14 = _5;
+-         _12 = slice_index::<*mut [i32], *mut i32>(move _13, move _14) -> [return: bb3, unwind unreachable];
++         _12 = &raw mut (*_13)[_14];
++         goto -> bb3;
+      }
+  
+      bb3: {
+          StorageDead(_14);
+          StorageDead(_13);
+          StorageLive(_15);
+          StorageLive(_16);
+          _16 = _4;
+          StorageLive(_17);
+          _17 = _5;
+-         _15 = slice_index::<*const [i32], *const i32>(move _16, move _17) -> [return: bb4, unwind unreachable];
++         _15 = &raw const (*_16)[_17];
++         goto -> bb4;
+      }
+  
+      bb4: {
+          StorageDead(_17);
+          StorageDead(_16);
+          _0 = const ();
+          StorageDead(_15);
+          StorageDead(_12);
+          StorageDead(_9);
+          StorageDead(_6);
+          return;
+      }
+  }
+  
diff --git a/tests/mir-opt/lower_intrinsics.index_project.LowerIntrinsics.panic-unwind.diff b/tests/mir-opt/lower_intrinsics.index_project.LowerIntrinsics.panic-unwind.diff
@@ -0,0 +1,97 @@
+- // MIR for `index_project` before LowerIntrinsics
++ // MIR for `index_project` after LowerIntrinsics
+  
+  fn index_project(_1: &mut [i32], _2: &[i32], _3: *mut [i32], _4: *const [i32], _5: usize) -> () {
+      debug a => _1;
+      debug b => _2;
+      debug c => _3;
+      debug d => _4;
+      debug i => _5;
+      let mut _0: ();
+      let _6: &mut i32;
+      let mut _7: &mut [i32];
+      let mut _8: usize;
+      let mut _10: &[i32];
+      let mut _11: usize;
+      let mut _13: *mut [i32];
+      let mut _14: usize;
+      let mut _16: *const [i32];
+      let mut _17: usize;
+      scope 1 {
+          debug _a => _6;
+          let _9: &i32;
+          scope 2 {
+              debug _b => _9;
+              let _12: *mut i32;
+              scope 3 {
+                  debug _c => _12;
+                  let _15: *const i32;
+                  scope 4 {
+                      debug _d => _15;
+                  }
+              }
+          }
+      }
+  
+      bb0: {
+          StorageLive(_6);
+          StorageLive(_7);
+          _7 = move _1;
+          StorageLive(_8);
+          _8 = _5;
+-         _6 = slice_index::<&mut [i32], &mut i32>(move _7, move _8) -> [return: bb1, unwind unreachable];
++         _6 = &mut (*_7)[_8];
++         goto -> bb1;
+      }
+  
+      bb1: {
+          StorageDead(_8);
+          StorageDead(_7);
+          StorageLive(_9);
+          StorageLive(_10);
+          _10 = _2;
+          StorageLive(_11);
+          _11 = _5;
+-         _9 = slice_index::<&[i32], &i32>(move _10, move _11) -> [return: bb2, unwind unreachable];
++         _9 = &(*_10)[_11];
++         goto -> bb2;
+      }
+  
+      bb2: {
+          StorageDead(_11);
+          StorageDead(_10);
+          StorageLive(_12);
+          StorageLive(_13);
+          _13 = _3;
+          StorageLive(_14);
+          _14 = _5;
+-         _12 = slice_index::<*mut [i32], *mut i32>(move _13, move _14) -> [return: bb3, unwind unreachable];
++         _12 = &raw mut (*_13)[_14];
++         goto -> bb3;
+      }
+  
+      bb3: {
+          StorageDead(_14);
+          StorageDead(_13);
+          StorageLive(_15);
+          StorageLive(_16);
+          _16 = _4;
+          StorageLive(_17);
+          _17 = _5;
+-         _15 = slice_index::<*const [i32], *const i32>(move _16, move _17) -> [return: bb4, unwind unreachable];
++         _15 = &raw const (*_16)[_17];
++         goto -> bb4;
+      }
+  
+      bb4: {
+          StorageDead(_17);
+          StorageDead(_16);
+          _0 = const ();
+          StorageDead(_15);
+          StorageDead(_12);
+          StorageDead(_9);
+          StorageDead(_6);
+          return;
+      }
+  }
+  
diff --git a/tests/mir-opt/lower_intrinsics.rs b/tests/mir-opt/lower_intrinsics.rs
@@ -241,3 +241,11 @@ pub unsafe fn ptr_offset(p: *const i32, d: isize) -> *const i32 {
 
     core::intrinsics::offset(p, d)
 }
+
+// EMIT_MIR lower_intrinsics.index_project.LowerIntrinsics.diff
+pub unsafe fn index_project(a: &mut [i32], b: &[i32], c: *mut [i32], d: *const [i32], i: usize) {
+    let _a: &mut i32 = core::intrinsics::slice_index(a, i);
+    let _b: &i32 = core::intrinsics::slice_index(b, i);
+    let _c: *mut i32 = core::intrinsics::slice_index(c, i);
+    let _d: *const i32 = core::intrinsics::slice_index(d, i);
+}
