Custom MIR: Support cleanup blocks

Cleanup blocks are declared with `bb (cleanup) = { ... }`.

`Call` and `Drop` terminators take an additional argument describing the
unwind action, which is one of the following:

* `Continue()`
* `Unreachable()`
* `Termiante(reason)`, where reason is `Abi` or `InCleanup`
* `Cleanup(block)`

Also support resume and terminate terminators:

* `Resume()`
* `Terminate(reason)`

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

@@ -30,7 +30,10 @@ macro_rules! parse_by_kind {
             @call($name:literal, $args:ident) => $call_expr:expr,
         )*
         $(
-            $pat:pat => $expr:expr,
+            @variant($adt:literal, $variant:literal) => $variant_expr:expr,
+        )*
+        $(
+            $pat:pat $(if $guard:expr)? => $expr:expr,
         )*
     ) => {{
         let expr_id = $self.preparse($expr_id);
@@ -49,7 +52,13 @@ macro_rules! parse_by_kind {
                 } => $call_expr,
             )*
             $(
-                $pat => $expr,
+                ExprKind::Adt(box AdtExpr { adt_def, variant_index, .. }) if {
+                    $self.tcx.is_diagnostic_item(rustc_span::Symbol::intern($adt), adt_def.did()) &&
+                    adt_def.variants()[*variant_index].name == rustc_span::Symbol::intern($variant)
+                } => $variant_expr,
+            )*
+            $(
+                $pat $(if $guard)? => $expr,
             )*
             #[allow(unreachable_patterns)]
             _ => return Err($self.expr_error(expr_id, $expected))
@@ -172,7 +181,8 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
             ExprKind::Block { block } => &self.thir[*block].stmts,
         );
         for (i, block_def) in block_defs.iter().enumerate() {
-            let block = self.parse_block_def(self.statement_as_expr(*block_def)?)?;
+            let is_cleanup = self.body.basic_blocks_mut()[BasicBlock::from_usize(i)].is_cleanup;
+            let block = self.parse_block_def(self.statement_as_expr(*block_def)?, is_cleanup)?;
             self.body.basic_blocks_mut()[BasicBlock::from_usize(i)] = block;
         }
 
@@ -181,8 +191,9 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
 
     fn parse_block_decls(&mut self, stmts: impl Iterator<Item = StmtId>) -> PResult<()> {
         for stmt in stmts {
-            let (var, _, _) = self.parse_let_statement(stmt)?;
-            let data = BasicBlockData::new(None);
+            let (var, _, _, initializer) = self.parse_let_statement(stmt)?;
+            let mut data = BasicBlockData::new(None);
+            data.is_cleanup = initializer.is_some();
             let block = self.body.basic_blocks_mut().push(data);
             self.block_map.insert(var, block);
         }
@@ -195,7 +206,7 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
         self.local_map.insert(ret_var, Local::from_u32(0));
 
         for stmt in stmts {
-            let (var, ty, span) = self.parse_let_statement(stmt)?;
+            let (var, ty, span, _) = self.parse_let_statement(stmt)?;
             let decl = LocalDecl::new(ty, span);
             let local = self.body.local_decls.push(decl);
             self.local_map.insert(var, local);
@@ -251,15 +262,17 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
         Ok(())
     }
 
-    fn parse_let_statement(&mut self, stmt_id: StmtId) -> PResult<(LocalVarId, Ty<'tcx>, Span)> {
-        let pattern = match &self.thir[stmt_id].kind {
-            StmtKind::Let { pattern, .. } => pattern,
-            StmtKind::Expr { expr, .. } => {
-                return Err(self.expr_error(*expr, "let statement"));
+    fn parse_let_statement(
+        &mut self,
+        stmt_id: StmtId,
+    ) -> PResult<(LocalVarId, Ty<'tcx>, Span, Option<ExprId>)> {
+        match &self.thir[stmt_id].kind {
+            StmtKind::Let { pattern, initializer, .. } => {
+                let (var, ty, span) = self.parse_var(pattern)?;
+                Ok((var, ty, span, *initializer))
             }
-        };
-
-        self.parse_var(pattern)
+            StmtKind::Expr { expr, .. } => Err(self.expr_error(*expr, "let statement")),
+        }
     }
 
     fn parse_var(&mut self, mut pat: &Pat<'tcx>) -> PResult<(LocalVarId, Ty<'tcx>, Span)> {
@@ -281,12 +294,13 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
         }
     }
 
-    fn parse_block_def(&self, expr_id: ExprId) -> PResult<BasicBlockData<'tcx>> {
+    fn parse_block_def(&self, expr_id: ExprId, is_cleanup: bool) -> PResult<BasicBlockData<'tcx>> {
         let block = parse_by_kind!(self, expr_id, _, "basic block",
             ExprKind::Block { block } => &self.thir[*block],
         );
 
         let mut data = BasicBlockData::new(None);
+        data.is_cleanup = is_cleanup;
         for stmt_id in &*block.stmts {
             let stmt = self.statement_as_expr(*stmt_id)?;
             let span = self.thir[stmt].span;

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

@@ -52,11 +52,17 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
             @call("mir_unreachable", _args) => {
                 Ok(TerminatorKind::Unreachable)
             },
+            @call("mir_resume", _args) => {
+                Ok(TerminatorKind::UnwindResume)
+            },
+            @call("mir_terminate", args) => {
+                Ok(TerminatorKind::UnwindTerminate(self.parse_unwind_terminate_reason(args[0])?))
+            },
             @call("mir_drop", args) => {
                 Ok(TerminatorKind::Drop {
                     place: self.parse_place(args[0])?,
                     target: self.parse_block(args[1])?,
-                    unwind: UnwindAction::Continue,
+                    unwind: self.parse_unwind_action(args[2])?,
                     replace: false,
                 })
             },
@@ -70,6 +76,34 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
         )
     }
 
+    fn parse_unwind_terminate_reason(&self, expr_id: ExprId) -> PResult<UnwindTerminateReason> {
+        parse_by_kind!(self, expr_id, _, "unwind terminate reason",
+            @variant("mir_unwind_terminate_reason", "Abi") => {
+                Ok(UnwindTerminateReason::Abi)
+            },
+            @variant("mir_unwind_terminate_reason", "InCleanup") => {
+                Ok(UnwindTerminateReason::InCleanup)
+            },
+        )
+    }
+
+    fn parse_unwind_action(&self, expr_id: ExprId) -> PResult<UnwindAction> {
+        parse_by_kind!(self, expr_id, _, "unwind action",
+            @call("mir_continue", _args) => {
+                Ok(UnwindAction::Continue)
+            },
+            @call("mir_unreachable", _args) => {
+                Ok(UnwindAction::Unreachable)
+            },
+            @call("mir_terminate", args) => {
+                Ok(UnwindAction::Terminate(self.parse_unwind_terminate_reason(args[0])?))
+            },
+            @call("mir_cleanup", args) => {
+                Ok(UnwindAction::Cleanup(self.parse_block(args[0])?))
+            },
+        )
+    }
+
     fn parse_match(&self, arms: &[ArmId], span: Span) -> PResult<SwitchTargets> {
         let Some((otherwise, rest)) = arms.split_last() else {
             return Err(ParseError {
@@ -113,6 +147,7 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
         );
         let destination = self.parse_place(destination)?;
         let target = self.parse_block(args[1])?;
+        let unwind = self.parse_unwind_action(args[2])?;
 
         parse_by_kind!(self, call, _, "function call",
             ExprKind::Call { fun, args, from_hir_call, fn_span, .. } => {
@@ -126,7 +161,7 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
                     args,
                     destination,
                     target: Some(target),
-                    unwind: UnwindAction::Continue,
+                    unwind,
                     call_source: if *from_hir_call { CallSource::Normal } else {
                         CallSource::OverloadedOperator
                     },

library/core/src/intrinsics/mir.rs

@@ -110,15 +110,15 @@
 //!         let popped;
 //!
 //!         {
-//!             Call(_unused = Vec::push(v, value), pop)
+//!             Call(_unused = Vec::push(v, value), pop, Continue())
 //!         }
 //!
 //!         pop = {
-//!             Call(popped = Vec::pop(v), drop)
+//!             Call(popped = Vec::pop(v), drop, Continue())
 //!         }
 //!
 //!         drop = {
-//!             Drop(popped, ret)
+//!             Drop(popped, ret, Continue())
 //!         }
 //!
 //!         ret = {
@@ -238,10 +238,6 @@
 //!
 //! #### Terminators
 //!
-//! 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`] and [`Drop`](Drop()) 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
@@ -262,6 +258,18 @@
 /// All terminators will have this type as a return type. It helps achieve some type safety.
 pub struct BasicBlock;
 
+/// The reason we are terminating the process during unwinding.
+#[rustc_diagnostic_item = "mir_unwind_terminate_reason"]
+pub enum UnwindTerminateReason {
+    /// Unwinding is just not possible given the ABI of this function.
+    Abi,
+    /// We were already cleaning up for an ongoing unwind, and a *second*, *nested* unwind was
+    /// triggered by the drop glue.
+    InCleanup,
+}
+
+pub use UnwindTerminateReason::*;
+
 macro_rules! define {
     ($name:literal, $( #[ $meta:meta ] )* fn $($sig:tt)*) => {
         #[rustc_diagnostic_item = $name]
@@ -270,12 +278,15 @@ macro_rules! define {
         pub fn $($sig)* { panic!() }
     }
 }
-
 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_call", fn Call(call: (), goto: BasicBlock));
+define!("mir_continue", fn Continue());
+define!("mir_resume", fn Resume());
+define!("mir_terminate", fn Terminate(reason: UnwindTerminateReason));
+define!("mir_cleanup", fn Cleanup(goto: BasicBlock));
+define!("mir_drop", fn Drop<T, U>(place: T, goto: BasicBlock, unwind_action: U));
+define!("mir_call", fn Call<U>(call: (), goto: BasicBlock, unwind_action: U));
 define!("mir_storage_live", fn StorageLive<T>(local: T));
 define!("mir_storage_dead", fn StorageDead<T>(local: T));
 define!("mir_deinit", fn Deinit<T>(place: T));
@@ -382,16 +393,15 @@ pub macro mir {
         }
 
         $(
-            $block_name:ident = {
+            $block_name:ident $(($block_cleanup:ident))? = {
                 $($block:tt)*
             }
         )*
     ) => {{
         // First, we declare all basic blocks.
-        $(
-            let $block_name: ::core::intrinsics::mir::BasicBlock;
-        )*
-
+        __internal_declare_basic_blocks!($(
+            $block_name $(($block_cleanup))?
+        )*);
         {
             // Now all locals
             #[allow(non_snake_case)]
@@ -585,3 +595,19 @@ pub macro __internal_remove_let {
         }
     },
 }
+
+/// Helper macro that declares the basic blocks.
+///
+/// The cleanup block are distinguished by an initializer.
+#[doc(hidden)]
+pub macro __internal_declare_basic_blocks {
+    () => {},
+    ($name:ident (cleanup) $($rest:tt)*) => {
+        let $name = ::core::intrinsics::mir::BasicBlock;
+        __internal_declare_basic_blocks!($($rest)*)
+    },
+    ($name:ident $($rest:tt)*) => {
+        let $name : ::core::intrinsics::mir::BasicBlock;
+        __internal_declare_basic_blocks!($($rest)*)
+    },
+}

tests/mir-opt/building/custom/terminate.rs

@@ -0,0 +1,34 @@
+// compile-flags: --crate-type=lib
+// edition:2021
+#![feature(custom_mir, core_intrinsics)]
+use core::intrinsics::mir::*;
+
+// CHECK-LABEL: fn f()
+// CHECK:       bb1 (cleanup): {
+// CHECK-NEXT:  abort(abi);
+#[custom_mir(dialect = "runtime", phase = "optimized")]
+pub fn f() {
+    mir!(
+        {
+            Return()
+        }
+        bb1(cleanup) = {
+            Terminate(Abi)
+        }
+    )
+}
+
+// CHECK-LABEL: fn g()
+// CHECK:       bb1 (cleanup): {
+// CHECK-NEXT:  abort(cleanup);
+#[custom_mir(dialect = "runtime", phase = "optimized")]
+pub fn g() {
+    mir!(
+        {
+            Return()
+        }
+        bb1(cleanup) = {
+            Terminate(InCleanup)
+        }
+    )
+}

tests/mir-opt/building/custom/terminators.rs

@@ -13,7 +13,7 @@ fn ident<T>(t: T) -> T {
 fn direct_call(x: i32) -> i32 {
     mir!(
         {
-            Call(RET = ident(x), retblock)
+            Call(RET = ident(x), retblock, Continue())
         }
 
         retblock = {
@@ -27,7 +27,7 @@ fn direct_call(x: i32) -> i32 {
 fn indirect_call(x: i32, f: fn(i32) -> i32) -> i32 {
     mir!(
         {
-            Call(RET = f(x), retblock)
+            Call(RET = f(x), retblock, Continue())
         }
 
         retblock = {
@@ -49,7 +49,7 @@ impl<'a> Drop for WriteOnDrop<'a> {
 fn drop_first<'a>(a: WriteOnDrop<'a>, b: WriteOnDrop<'a>) {
     mir!(
         {
-            Drop(a, retblock)
+            Drop(a, retblock, Continue())
         }
 
         retblock = {
@@ -64,7 +64,7 @@ fn drop_first<'a>(a: WriteOnDrop<'a>, b: WriteOnDrop<'a>) {
 fn drop_second<'a>(a: WriteOnDrop<'a>, b: WriteOnDrop<'a>) {
     mir!(
         {
-            Drop(b, retblock)
+            Drop(b, retblock, Continue())
         }
 
         retblock = {