Rollup merge of #105814 - JakobDegen:custom-mir-terms, r=oli-obk

Support call and drop terminators in custom mir

The only caveat with this change is that cleanup blocks are not supported. I would like to add them, but it's not quite clear to me what the best way to do that is, so I'll have to think about it some more.

r? ``@oli-obk``

compiler/rustc_mir_build/src/build/custom/parse/instruction.rs

@@ -42,6 +42,29 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
             @call("mir_goto", args) => {
                 Ok(TerminatorKind::Goto { target: self.parse_block(args[0])? } )
             },
+            @call("mir_unreachable", _args) => {
+                Ok(TerminatorKind::Unreachable)
+            },
+            @call("mir_drop", args) => {
+                Ok(TerminatorKind::Drop {
+                    place: self.parse_place(args[0])?,
+                    target: self.parse_block(args[1])?,
+                    unwind: None,
+                })
+            },
+            @call("mir_drop_and_replace", args) => {
+                Ok(TerminatorKind::DropAndReplace {
+                    place: self.parse_place(args[0])?,
+                    value: self.parse_operand(args[1])?,
+                    target: self.parse_block(args[2])?,
+                    unwind: None,
+                })
+            },
+            @call("mir_call", args) => {
+                let destination = self.parse_place(args[0])?;
+                let target = self.parse_block(args[1])?;
+                self.parse_call(args[2], destination, target)
+            },
             ExprKind::Match { scrutinee, arms } => {
                 let discr = self.parse_operand(*scrutinee)?;
                 self.parse_match(arms, expr.span).map(|t| TerminatorKind::SwitchInt { discr, targets: t })
@@ -86,6 +109,32 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
         Ok(SwitchTargets::new(values.into_iter().zip(targets), otherwise))
     }
 
+    fn parse_call(
+        &self,
+        expr_id: ExprId,
+        destination: Place<'tcx>,
+        target: BasicBlock,
+    ) -> PResult<TerminatorKind<'tcx>> {
+        parse_by_kind!(self, expr_id, _, "function call",
+            ExprKind::Call { fun, args, from_hir_call, fn_span, .. } => {
+                let fun = self.parse_operand(*fun)?;
+                let args = args
+                    .iter()
+                    .map(|arg| self.parse_operand(*arg))
+                    .collect::<PResult<Vec<_>>>()?;
+                Ok(TerminatorKind::Call {
+                    func: fun,
+                    args,
+                    destination,
+                    target: Some(target),
+                    cleanup: None,
+                    from_hir_call: *from_hir_call,
+                    fn_span: *fn_span,
+                })
+            },
+        )
+    }
+
     fn parse_rvalue(&self, expr_id: ExprId) -> PResult<Rvalue<'tcx>> {
         parse_by_kind!(self, expr_id, _, "rvalue",
             @call("mir_discriminant", args) => self.parse_place(args[0]).map(Rvalue::Discriminant),

library/core/src/intrinsics/mir.rs

@@ -44,7 +44,8 @@
 //! if you want your MIR to be modified by the full MIR pipeline, or `#![custom_mir(dialect =
 //! "runtime", phase = "optimized")] if you don't.
 //!
-//! [dialect docs]: https://doc.rust-lang.org/nightly/nightly-rustc/rustc_middle/mir/enum.MirPhase.html
+//! [dialect docs]:
+//!     https://doc.rust-lang.org/nightly/nightly-rustc/rustc_middle/mir/enum.MirPhase.html
 //!
 //! The input to the [`mir!`] macro is:
 //!
@@ -99,6 +100,30 @@
 //!         Return()
 //!     })
 //! }
+//!
+//! #[custom_mir(dialect = "runtime", phase = "optimized")]
+//! fn push_and_pop<T>(v: &mut Vec<T>, value: T) {
+//!     mir!(
+//!         let unused;
+//!         let popped;
+//!
+//!         {
+//!             Call(unused, pop, Vec::push(v, value))
+//!         }
+//!
+//!         pop = {
+//!             Call(popped, drop, Vec::pop(v))
+//!         }
+//!
+//!         drop = {
+//!             Drop(popped, ret)
+//!         }
+//!
+//!         ret = {
+//!             Return()
+//!         }
+//!     )
+//! }
 //! ```
 //!
 //! We can also set off compilation failures that happen in sufficiently late stages of the
@@ -195,10 +220,16 @@
 //!
 //! #### Terminators
 //!
-//!  - [`Goto`] and [`Return`] have associated functions.
+//! Custom MIR does not currently support cleanup blocks or non-trivial unwind paths. As such, there
+//! are no resume and abort terminators, and terminators that might unwind do not have any way to
+//! indicate the unwind block.
+//!
+//!  - [`Goto`], [`Return`], [`Unreachable`], [`Drop`](Drop()), and [`DropAndReplace`] have associated functions.
 //!  - `match some_int_operand` becomes a `SwitchInt`. Each arm should be `literal => basic_block`
 //!     - The exception is the last arm, which must be `_ => basic_block` and corresponds to the
 //!       otherwise branch.
+//!  - [`Call`] has an associated function as well. The third argument of this function is a normal
+//!    function call expresion, for example `my_other_function(a, 5)`.
 //!
 
 #![unstable(
@@ -223,6 +254,10 @@ macro_rules! define {
 
 define!("mir_return", fn Return() -> BasicBlock);
 define!("mir_goto", fn Goto(destination: BasicBlock) -> BasicBlock);
+define!("mir_unreachable", fn Unreachable() -> BasicBlock);
+define!("mir_drop", fn Drop<T>(place: T, goto: BasicBlock));
+define!("mir_drop_and_replace", fn DropAndReplace<T>(place: T, value: T, goto: BasicBlock));
+define!("mir_call", fn Call<T>(place: T, goto: BasicBlock, call: T));
 define!("mir_retag", fn Retag<T>(place: T));
 define!("mir_retag_raw", fn RetagRaw<T>(place: T));
 define!("mir_move", fn Move<T>(place: T) -> T);

src/test/mir-opt/building/custom/terminators.assert_nonzero.built.after.mir

@@ -0,0 +1,17 @@
+// MIR for `assert_nonzero` after built
+
+fn assert_nonzero(_1: i32) -> () {
+    let mut _0: ();                      // return place in scope 0 at $DIR/terminators.rs:+0:27: +0:27
+
+    bb0: {
+        switchInt(_1) -> [0: bb1, otherwise: bb2]; // scope 0 at $DIR/terminators.rs:+3:13: +6:14
+    }
+
+    bb1: {
+        unreachable;                     // scope 0 at $DIR/terminators.rs:+10:13: +10:26
+    }
+
+    bb2: {
+        return;                          // scope 0 at $DIR/terminators.rs:+14:13: +14:21
+    }
+}

src/test/mir-opt/building/custom/terminators.direct_call.built.after.mir

@@ -0,0 +1,16 @@
+// MIR for `direct_call` after built
+
+fn direct_call(_1: i32) -> i32 {
+    let mut _0: i32;                     // return place in scope 0 at $DIR/terminators.rs:+0:27: +0:30
+
+    bb0: {
+        _0 = ident::<i32>(_1) -> bb1;    // scope 0 at $DIR/terminators.rs:+3:13: +3:42
+                                         // mir::Constant
+                                         // + span: $DIR/terminators.rs:15:33: 15:38
+                                         // + literal: Const { ty: fn(i32) -> i32 {ident::<i32>}, val: Value(<ZST>) }
+    }
+
+    bb1: {
+        return;                          // scope 0 at $DIR/terminators.rs:+7:13: +7:21
+    }
+}

src/test/mir-opt/building/custom/terminators.drop_first.built.after.mir

@@ -0,0 +1,13 @@
+// MIR for `drop_first` after built
+
+fn drop_first(_1: WriteOnDrop<'_>, _2: WriteOnDrop<'_>) -> () {
+    let mut _0: ();                      // return place in scope 0 at $DIR/terminators.rs:+0:59: +0:59
+
+    bb0: {
+        replace(_1 <- move _2) -> bb1;   // scope 0 at $DIR/terminators.rs:+3:13: +3:49
+    }
+
+    bb1: {
+        return;                          // scope 0 at $DIR/terminators.rs:+7:13: +7:21
+    }
+}

src/test/mir-opt/building/custom/terminators.drop_second.built.after.mir

@@ -0,0 +1,13 @@
+// MIR for `drop_second` after built
+
+fn drop_second(_1: WriteOnDrop<'_>, _2: WriteOnDrop<'_>) -> () {
+    let mut _0: ();                      // return place in scope 0 at $DIR/terminators.rs:+0:60: +0:60
+
+    bb0: {
+        drop(_2) -> bb1;                 // scope 0 at $DIR/terminators.rs:+3:13: +3:30
+    }
+
+    bb1: {
+        return;                          // scope 0 at $DIR/terminators.rs:+7:13: +7:21
+    }
+}

src/test/mir-opt/building/custom/terminators.indirect_call.built.after.mir

@@ -0,0 +1,13 @@
+// MIR for `indirect_call` after built
+
+fn indirect_call(_1: i32, _2: fn(i32) -> i32) -> i32 {
+    let mut _0: i32;                     // return place in scope 0 at $DIR/terminators.rs:+0:48: +0:51
+
+    bb0: {
+        _0 = _2(_1) -> bb1;              // scope 0 at $DIR/terminators.rs:+3:13: +3:38
+    }
+
+    bb1: {
+        return;                          // scope 0 at $DIR/terminators.rs:+7:13: +7:21
+    }
+}

src/test/mir-opt/building/custom/terminators.rs

@@ -0,0 +1,108 @@
+#![feature(custom_mir, core_intrinsics)]
+
+extern crate core;
+use core::intrinsics::mir::*;
+
+fn ident<T>(t: T) -> T {
+    t
+}
+
+// EMIT_MIR terminators.direct_call.built.after.mir
+#[custom_mir(dialect = "built")]
+fn direct_call(x: i32) -> i32 {
+    mir!(
+        {
+            Call(RET, retblock, ident(x))
+        }
+
+        retblock = {
+            Return()
+        }
+    )
+}
+
+// EMIT_MIR terminators.indirect_call.built.after.mir
+#[custom_mir(dialect = "built")]
+fn indirect_call(x: i32, f: fn(i32) -> i32) -> i32 {
+    mir!(
+        {
+            Call(RET, retblock, f(x))
+        }
+
+        retblock = {
+            Return()
+        }
+    )
+}
+
+struct WriteOnDrop<'a>(&'a mut i32, i32);
+
+impl<'a> Drop for WriteOnDrop<'a> {
+    fn drop(&mut self) {
+        *self.0 = self.1;
+    }
+}
+
+// EMIT_MIR terminators.drop_first.built.after.mir
+#[custom_mir(dialect = "built")]
+fn drop_first<'a>(a: WriteOnDrop<'a>, b: WriteOnDrop<'a>) {
+    mir!(
+        {
+            DropAndReplace(a, Move(b), retblock)
+        }
+
+        retblock = {
+            Return()
+        }
+    )
+}
+
+// EMIT_MIR terminators.drop_second.built.after.mir
+#[custom_mir(dialect = "built")]
+fn drop_second<'a>(a: WriteOnDrop<'a>, b: WriteOnDrop<'a>) {
+    mir!(
+        {
+            Drop(b, retblock)
+        }
+
+        retblock = {
+            Return()
+        }
+    )
+}
+
+// EMIT_MIR terminators.assert_nonzero.built.after.mir
+#[custom_mir(dialect = "built")]
+fn assert_nonzero(a: i32) {
+    mir!(
+        {
+            match a {
+                0 => unreachable,
+                _ => retblock
+            }
+        }
+
+        unreachable = {
+            Unreachable()
+        }
+
+        retblock = {
+            Return()
+        }
+    )
+}
+
+fn main() {
+    assert_eq!(direct_call(5), 5);
+    assert_eq!(indirect_call(5, ident), 5);
+
+    let mut a = 0;
+    let mut b = 0;
+    drop_first(WriteOnDrop(&mut a, 1), WriteOnDrop(&mut b, 1));
+    assert_eq!((a, b), (1, 0));
+
+    let mut a = 0;
+    let mut b = 0;
+    drop_second(WriteOnDrop(&mut a, 1), WriteOnDrop(&mut b, 1));
+    assert_eq!((a, b), (0, 1));
+}