Support more projections in custom mir

9 files changed, 210 insertions, 6 deletions

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 bc822113ce1..2f26499a3b6 100644
--- a/compiler/rustc_mir_build/src/build/custom/parse/instruction.rs
+++ b/compiler/rustc_mir_build/src/build/custom/parse/instruction.rs
@@ -1,4 +1,5 @@
 use rustc_middle::mir::interpret::{ConstValue, Scalar};
+use rustc_middle::mir::tcx::PlaceTy;
 use rustc_middle::{mir::*, thir::*, ty};
 use rustc_span::Span;
 use rustc_target::abi::VariantIdx;
@@ -118,12 +119,42 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
     }
 
     fn parse_place(&self, expr_id: ExprId) -> PResult<Place<'tcx>> {
-        parse_by_kind!(self, expr_id, _, "place",
-            ExprKind::Deref { arg } => Ok(
-                self.parse_place(*arg)?.project_deeper(&[PlaceElem::Deref], self.tcx)
-            ),
-            _ => self.parse_local(expr_id).map(Place::from),
-        )
+        self.parse_place_inner(expr_id).map(|(x, _)| x)
+    }
+
+    fn parse_place_inner(&self, expr_id: ExprId) -> PResult<(Place<'tcx>, PlaceTy<'tcx>)> {
+        let (parent, proj) = parse_by_kind!(self, expr_id, expr, "place",
+            @call("mir_field", args) => {
+                let (parent, ty) = self.parse_place_inner(args[0])?;
+                let field = Field::from_u32(self.parse_integer_literal(args[1])? as u32);
+                let field_ty = ty.field_ty(self.tcx, field);
+                let proj = PlaceElem::Field(field, field_ty);
+                let place = parent.project_deeper(&[proj], self.tcx);
+                return Ok((place, PlaceTy::from_ty(field_ty)));
+            },
+            @call("mir_variant", args) => {
+                (args[0], PlaceElem::Downcast(
+                    None,
+                    VariantIdx::from_u32(self.parse_integer_literal(args[1])? as u32)
+                ))
+            },
+            ExprKind::Deref { arg } => {
+                parse_by_kind!(self, *arg, _, "does not matter",
+                    @call("mir_make_place", args) => return self.parse_place_inner(args[0]),
+                    _ => (*arg, PlaceElem::Deref),
+                )
+            },
+            ExprKind::Index { lhs, index } => (*lhs, PlaceElem::Index(self.parse_local(*index)?)),
+            ExprKind::Field { lhs, name: field, .. } => (*lhs, PlaceElem::Field(*field, expr.ty)),
+            _ => {
+                let place = self.parse_local(expr_id).map(Place::from)?;
+                return Ok((place, PlaceTy::from_ty(expr.ty)))
+            },
+        );
+        let (parent, ty) = self.parse_place_inner(parent)?;
+        let place = parent.project_deeper(&[proj], self.tcx);
+        let ty = ty.projection_ty(self.tcx, proj);
+        Ok((place, ty))
     }
 
     fn parse_local(&self, expr_id: ExprId) -> PResult<Local> {
diff --git a/library/core/src/intrinsics/mir.rs b/library/core/src/intrinsics/mir.rs
index 4759f128bbb..7d7c61b1180 100644
--- a/library/core/src/intrinsics/mir.rs
+++ b/library/core/src/intrinsics/mir.rs
@@ -88,6 +88,9 @@ 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_field", fn Field<F>(place: (), field: u32) -> F);
+define!("mir_variant", fn Variant<T>(place: T, index: u32) -> ());
+define!("mir_make_place", fn __internal_make_place<T>(place: T) -> *mut T);
 
 /// Convenience macro for generating custom MIR.
 ///
@@ -145,6 +148,25 @@ pub macro mir {
     }}
 }
 
+/// Helper macro that allows you to treat a value expression like a place expression.
+///
+/// This is necessary in combination with the [`Field`] and [`Variant`] methods. Specifically,
+/// something like this won't compile on its own, reporting an error about not being able to assign
+/// to such an expression:
+///
+/// ```rust,ignore(syntax-highlighting-only)
+/// Field(something, 0) = 5;
+/// ```
+///
+/// Instead, you'll need to write
+///
+/// ```rust,ignore(syntax-highlighting-only)
+/// place!(Field(something, 0)) = 5;
+/// ```
+pub macro place($e:expr) {
+    (*::core::intrinsics::mir::__internal_make_place($e))
+}
+
 /// Helper macro that extracts the `let` declarations out of a bunch of statements.
 ///
 /// This macro is written using the "statement muncher" strategy. Each invocation parses the first
diff --git a/src/test/mir-opt/building/custom/projections.rs b/src/test/mir-opt/building/custom/projections.rs
new file mode 100644
index 00000000000..5e472e531c7
--- /dev/null
+++ b/src/test/mir-opt/building/custom/projections.rs
@@ -0,0 +1,85 @@
+#![feature(custom_mir, core_intrinsics)]
+
+extern crate core;
+use core::intrinsics::mir::*;
+
+union U {
+    a: i32,
+    b: u32,
+}
+
+// EMIT_MIR projections.unions.built.after.mir
+#[custom_mir(dialect = "built")]
+fn unions(u: U) -> i32 {
+    mir!({
+        RET = u.a;
+        Return()
+    })
+}
+
+// EMIT_MIR projections.tuples.built.after.mir
+#[custom_mir(dialect = "analysis", phase = "post-cleanup")]
+fn tuples(i: (u32, i32)) -> (u32, i32) {
+    mir!(
+        // FIXME(JakobDegen): This is necessary because we can't give type hints for `RET`
+        let temp: (u32, i32);
+        {
+            temp.0 = i.0;
+            temp.1 = i.1;
+
+            RET = temp;
+            Return()
+        }
+    )
+}
+
+// EMIT_MIR projections.unwrap.built.after.mir
+#[custom_mir(dialect = "built")]
+fn unwrap(opt: Option<i32>) -> i32 {
+    mir!({
+        RET = Field(Variant(opt, 1), 0);
+        Return()
+    })
+}
+
+// EMIT_MIR projections.unwrap_deref.built.after.mir
+#[custom_mir(dialect = "built")]
+fn unwrap_deref(opt: Option<&i32>) -> i32 {
+    mir!({
+        RET = *Field::<&i32>(Variant(opt, 1), 0);
+        Return()
+    })
+}
+
+// EMIT_MIR projections.set.built.after.mir
+#[custom_mir(dialect = "built")]
+fn set(opt: &mut Option<i32>) {
+    mir!({
+        place!(Field(Variant(*opt, 1), 0)) = 10;
+        Return()
+    })
+}
+
+// EMIT_MIR projections.simple_index.built.after.mir
+#[custom_mir(dialect = "built")]
+fn simple_index(a: [i32; 10], b: &[i32]) -> i32 {
+    mir!({
+        let temp = 3;
+        RET = a[temp];
+        RET = (*b)[temp];
+        Return()
+    })
+}
+
+fn main() {
+    assert_eq!(unions(U { a: 5 }), 5);
+    assert_eq!(tuples((5, 6)), (5, 6));
+
+    assert_eq!(unwrap(Some(5)), 5);
+    assert_eq!(unwrap_deref(Some(&5)), 5);
+    let mut o = Some(5);
+    set(&mut o);
+    assert_eq!(o, Some(10));
+
+    assert_eq!(simple_index([0; 10], &[0; 10]), 0);
+}
diff --git a/src/test/mir-opt/building/custom/projections.set.built.after.mir b/src/test/mir-opt/building/custom/projections.set.built.after.mir
new file mode 100644
index 00000000000..2f15176a61f
--- /dev/null
+++ b/src/test/mir-opt/building/custom/projections.set.built.after.mir
@@ -0,0 +1,10 @@
+// MIR for `set` after built
+
+fn set(_1: &mut Option<i32>) -> () {
+    let mut _0: ();                      // return place in scope 0 at $DIR/projections.rs:+0:31: +0:31
+
+    bb0: {
+        (((*_1) as variant#1).0: i32) = const 10_i32; // scope 0 at $DIR/projections.rs:+2:9: +2:48
+        return;                          // scope 0 at $DIR/projections.rs:+3:9: +3:17
+    }
+}
diff --git a/src/test/mir-opt/building/custom/projections.simple_index.built.after.mir b/src/test/mir-opt/building/custom/projections.simple_index.built.after.mir
new file mode 100644
index 00000000000..fc422e4b3f0
--- /dev/null
+++ b/src/test/mir-opt/building/custom/projections.simple_index.built.after.mir
@@ -0,0 +1,13 @@
+// MIR for `simple_index` after built
+
+fn simple_index(_1: [i32; 10], _2: &[i32]) -> i32 {
+    let mut _0: i32;                     // return place in scope 0 at $DIR/projections.rs:+0:45: +0:48
+    let mut _3: usize;                   // in scope 0 at $SRC_DIR/core/src/intrinsics/mir.rs:LL:COL
+
+    bb0: {
+        _3 = const 3_usize;              // scope 0 at $SRC_DIR/core/src/intrinsics/mir.rs:LL:COL
+        _0 = _1[_3];                     // scope 0 at $DIR/projections.rs:+3:9: +3:22
+        _0 = (*_2)[_3];                  // scope 0 at $DIR/projections.rs:+4:9: +4:25
+        return;                          // scope 0 at $DIR/projections.rs:+5:9: +5:17
+    }
+}
diff --git a/src/test/mir-opt/building/custom/projections.tuples.built.after.mir b/src/test/mir-opt/building/custom/projections.tuples.built.after.mir
new file mode 100644
index 00000000000..65487d3c9ed
--- /dev/null
+++ b/src/test/mir-opt/building/custom/projections.tuples.built.after.mir
@@ -0,0 +1,13 @@
+// MIR for `tuples` after built
+
+fn tuples(_1: (u32, i32)) -> (u32, i32) {
+    let mut _0: (u32, i32);              // return place in scope 0 at $DIR/projections.rs:+0:29: +0:39
+    let mut _2: (u32, i32);              // in scope 0 at $SRC_DIR/core/src/intrinsics/mir.rs:LL:COL
+
+    bb0: {
+        (_2.0: u32) = (_1.0: u32);       // scope 0 at $DIR/projections.rs:+5:13: +5:25
+        (_2.1: i32) = (_1.1: i32);       // scope 0 at $DIR/projections.rs:+6:13: +6:25
+        _0 = _2;                         // scope 0 at $DIR/projections.rs:+8:13: +8:23
+        return;                          // scope 0 at $DIR/projections.rs:+9:13: +9:21
+    }
+}
diff --git a/src/test/mir-opt/building/custom/projections.unions.built.after.mir b/src/test/mir-opt/building/custom/projections.unions.built.after.mir
new file mode 100644
index 00000000000..922538a5f17
--- /dev/null
+++ b/src/test/mir-opt/building/custom/projections.unions.built.after.mir
@@ -0,0 +1,10 @@
+// MIR for `unions` after built
+
+fn unions(_1: U) -> i32 {
+    let mut _0: i32;                     // return place in scope 0 at $DIR/projections.rs:+0:20: +0:23
+
+    bb0: {
+        _0 = (_1.0: i32);                // scope 0 at $DIR/projections.rs:+2:9: +2:18
+        return;                          // scope 0 at $DIR/projections.rs:+3:9: +3:17
+    }
+}
diff --git a/src/test/mir-opt/building/custom/projections.unwrap.built.after.mir b/src/test/mir-opt/building/custom/projections.unwrap.built.after.mir
new file mode 100644
index 00000000000..75b03a3c393
--- /dev/null
+++ b/src/test/mir-opt/building/custom/projections.unwrap.built.after.mir
@@ -0,0 +1,10 @@
+// MIR for `unwrap` after built
+
+fn unwrap(_1: Option<i32>) -> i32 {
+    let mut _0: i32;                     // return place in scope 0 at $DIR/projections.rs:+0:32: +0:35
+
+    bb0: {
+        _0 = ((_1 as variant#1).0: i32); // scope 0 at $DIR/projections.rs:+2:9: +2:40
+        return;                          // scope 0 at $DIR/projections.rs:+3:9: +3:17
+    }
+}
diff --git a/src/test/mir-opt/building/custom/projections.unwrap_deref.built.after.mir b/src/test/mir-opt/building/custom/projections.unwrap_deref.built.after.mir
new file mode 100644
index 00000000000..c6b0f7efa9b
--- /dev/null
+++ b/src/test/mir-opt/building/custom/projections.unwrap_deref.built.after.mir
@@ -0,0 +1,10 @@
+// MIR for `unwrap_deref` after built
+
+fn unwrap_deref(_1: Option<&i32>) -> i32 {
+    let mut _0: i32;                     // return place in scope 0 at $DIR/projections.rs:+0:39: +0:42
+
+    bb0: {
+        _0 = (*((_1 as variant#1).0: &i32)); // scope 0 at $DIR/projections.rs:+2:9: +2:49
+        return;                          // scope 0 at $DIR/projections.rs:+3:9: +3:17
+    }
+}