From cecb901e683cdc1a2ffbb1e331ba4f262e2215c0 Mon Sep 17 00:00:00 2001
From: Andy Wang <cbeuw.andy@gmail.com>
Date: Tue, 11 Apr 2023 16:23:35 +0200
Subject: [PATCH] Add Offset binary op to custom mir

---
 .../src/build/custom/parse/instruction.rs             |  5 +++++
 library/core/src/intrinsics/mir.rs                    |  2 ++
 .../references.raw_pointer_offset.built.after.mir     | 10 ++++++++++
 tests/mir-opt/building/custom/references.rs           | 11 +++++++++++
 4 files changed, 28 insertions(+)
 create mode 100644 tests/mir-opt/building/custom/references.raw_pointer_offset.built.after.mir

diff --git a/compiler/rustc_mir_build/src/build/custom/parse/instruction.rs b/compiler/rustc_mir_build/src/build/custom/parse/instruction.rs
index 54028dfe87b73..931fe1b2433a0 100644
--- a/compiler/rustc_mir_build/src/build/custom/parse/instruction.rs
+++ b/compiler/rustc_mir_build/src/build/custom/parse/instruction.rs
@@ -148,6 +148,11 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
                     )),
                 )
             },
+            @call("mir_offset", args) => {
+                let ptr = self.parse_operand(args[0])?;
+                let offset = self.parse_operand(args[1])?;
+                Ok(Rvalue::BinaryOp(BinOp::Offset, Box::new((ptr, offset))))
+            },
             @call("mir_len", args) => Ok(Rvalue::Len(self.parse_place(args[0])?)),
             ExprKind::Borrow { borrow_kind, arg } => Ok(
                 Rvalue::Ref(self.tcx.lifetimes.re_erased, *borrow_kind, self.parse_place(*arg)?)
diff --git a/library/core/src/intrinsics/mir.rs b/library/core/src/intrinsics/mir.rs
index 6690c1a76d5f0..d9d62eb759e69 100644
--- a/library/core/src/intrinsics/mir.rs
+++ b/library/core/src/intrinsics/mir.rs
@@ -232,6 +232,7 @@
 //!  - `&`, `&mut`, `addr_of!`, and `addr_of_mut!` all work to create their associated rvalue.
 //!  - [`Discriminant`] and [`Len`] have associated functions.
 //!  - Unary and binary operations use their normal Rust syntax - `a * b`, `!c`, etc.
+//!  - The binary operation `Offset` can be created via [`Offset`].
 //!  - Checked binary operations are represented by wrapping the associated binop in [`Checked`].
 //!  - Array repetition syntax (`[foo; 10]`) creates the associated rvalue.
 //!
@@ -289,6 +290,7 @@ define!(
     fn Discriminant<T>(place: T) -> <T as ::core::marker::DiscriminantKind>::Discriminant
 );
 define!("mir_set_discriminant", fn SetDiscriminant<T>(place: T, index: u32));
+define!("mir_offset", fn Offset<T, U>(ptr: T, count: U) -> T);
 define!(
     "mir_field",
     /// Access the field with the given index of some place.
diff --git a/tests/mir-opt/building/custom/references.raw_pointer_offset.built.after.mir b/tests/mir-opt/building/custom/references.raw_pointer_offset.built.after.mir
new file mode 100644
index 0000000000000..f614aef4029d0
--- /dev/null
+++ b/tests/mir-opt/building/custom/references.raw_pointer_offset.built.after.mir
@@ -0,0 +1,10 @@
+// MIR for `raw_pointer_offset` after built
+
+fn raw_pointer_offset(_1: *const i32) -> *const i32 {
+    let mut _0: *const i32;              // return place in scope 0 at $DIR/references.rs:+0:45: +0:55
+
+    bb0: {
+        _0 = Offset(_1, const 1_isize);  // scope 0 at $DIR/references.rs:+2:9: +2:33
+        return;                          // scope 0 at $DIR/references.rs:+3:9: +3:17
+    }
+}
diff --git a/tests/mir-opt/building/custom/references.rs b/tests/mir-opt/building/custom/references.rs
index a1c896de04cf0..f87f6664c7a5c 100644
--- a/tests/mir-opt/building/custom/references.rs
+++ b/tests/mir-opt/building/custom/references.rs
@@ -45,11 +45,22 @@ pub fn raw_pointer(x: *const i32) -> *const i32 {
     })
 }
 
+// EMIT_MIR references.raw_pointer_offset.built.after.mir
+#[custom_mir(dialect = "built")]
+pub fn raw_pointer_offset(x: *const i32) -> *const i32 {
+    mir!({
+        RET = Offset(x, 1_isize);
+        Return()
+    })
+}
+
 fn main() {
     let mut x = 5;
+    let arr = [1, 2];
     assert_eq!(*mut_ref(&mut x), 5);
     assert_eq!(*immut_ref(&x), 5);
     unsafe {
         assert_eq!(*raw_pointer(addr_of!(x)), 5);
+        assert_eq!(*raw_pointer_offset(addr_of!(arr[0])), 2);
     }
 }