From 212417b87f8415d92ea0a365605011050a555e28 Mon Sep 17 00:00:00 2001
From: Ralf Jung <post@ralfj.de>
Date: Mon, 5 Aug 2024 18:23:14 +0200
Subject: [PATCH] custom MIR: add support for tail calls

---
 .../src/build/custom/parse/instruction.rs     | 22 +++++++++++++++++++
 compiler/rustc_span/src/symbol.rs             |  1 +
 library/core/src/intrinsics/mir.rs            |  8 +++++++
 tests/mir-opt/building/custom/terminators.rs  | 12 ++++++++++
 .../terminators.tail_call.built.after.mir     | 11 ++++++++++
 5 files changed, 54 insertions(+)
 create mode 100644 tests/mir-opt/building/custom/terminators.tail_call.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 3d4b706aa652b..56896d945e5f3 100644
--- a/compiler/rustc_mir_build/src/build/custom/parse/instruction.rs
+++ b/compiler/rustc_mir_build/src/build/custom/parse/instruction.rs
@@ -75,6 +75,9 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
             @call(mir_call, args) => {
                 self.parse_call(args)
             },
+            @call(mir_tail_call, args) => {
+                self.parse_tail_call(args)
+            },
             ExprKind::Match { scrutinee, arms, .. } => {
                 let discr = self.parse_operand(*scrutinee)?;
                 self.parse_match(arms, expr.span).map(|t| TerminatorKind::SwitchInt { discr, targets: t })
@@ -187,6 +190,25 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
         )
     }
 
+    fn parse_tail_call(&self, args: &[ExprId]) -> PResult<TerminatorKind<'tcx>> {
+        parse_by_kind!(self, args[0], _, "tail call",
+            ExprKind::Call { fun, args, fn_span, .. } => {
+                let fun = self.parse_operand(*fun)?;
+                let args = args
+                    .iter()
+                    .map(|arg|
+                        Ok(Spanned { node: self.parse_operand(*arg)?, span: self.thir.exprs[*arg].span  } )
+                    )
+                    .collect::<PResult<Box<[_]>>>()?;
+                Ok(TerminatorKind::TailCall {
+                    func: fun,
+                    args,
+                    fn_span: *fn_span,
+                })
+            },
+        )
+    }
+
     fn parse_rvalue(&self, expr_id: ExprId) -> PResult<Rvalue<'tcx>> {
         parse_by_kind!(self, expr_id, expr, "rvalue",
             @call(mir_discriminant, args) => self.parse_place(args[0]).map(Rvalue::Discriminant),
diff --git a/compiler/rustc_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs
index 9977fa7425a3e..94cf21da4efb8 100644
--- a/compiler/rustc_span/src/symbol.rs
+++ b/compiler/rustc_span/src/symbol.rs
@@ -1216,6 +1216,7 @@ symbols! {
         mir_static_mut,
         mir_storage_dead,
         mir_storage_live,
+        mir_tail_call,
         mir_unreachable,
         mir_unwind_cleanup,
         mir_unwind_continue,
diff --git a/library/core/src/intrinsics/mir.rs b/library/core/src/intrinsics/mir.rs
index fd49a96eaa049..c7cec396e1f2e 100644
--- a/library/core/src/intrinsics/mir.rs
+++ b/library/core/src/intrinsics/mir.rs
@@ -247,6 +247,8 @@
 //!       otherwise branch.
 //!  - [`Call`] has an associated function as well, with special syntax:
 //!    `Call(ret_val = function(arg1, arg2, ...), ReturnTo(next_block), UnwindContinue())`.
+//!  - [`TailCall`] does not have a return destination or next block, so its syntax is just
+//!    `TailCall(function(arg1, arg2, ...))`.
 
 #![unstable(
     feature = "custom_mir",
@@ -350,6 +352,12 @@ define!("mir_call",
     /// - [`UnwindCleanup`]
     fn Call(call: (), goto: ReturnToArg, unwind_action: UnwindActionArg)
 );
+define!("mir_tail_call",
+    /// Call a function.
+    ///
+    /// The argument must be of the form `fun(arg1, arg2, ...)`.
+    fn TailCall<T>(call: T)
+);
 define!("mir_unwind_resume",
     /// A terminator that resumes the unwinding.
     fn UnwindResume()
diff --git a/tests/mir-opt/building/custom/terminators.rs b/tests/mir-opt/building/custom/terminators.rs
index a8e0b4b35bf47..ed08040a2a58c 100644
--- a/tests/mir-opt/building/custom/terminators.rs
+++ b/tests/mir-opt/building/custom/terminators.rs
@@ -22,6 +22,18 @@ fn direct_call(x: i32) -> i32 {
     }
 }
 
+// EMIT_MIR terminators.tail_call.built.after.mir
+#[custom_mir(dialect = "built")]
+fn tail_call(x: i32) -> i32 {
+    mir! {
+        let y;
+        {
+            y = x + 42;
+            TailCall(ident(y))
+        }
+    }
+}
+
 // EMIT_MIR terminators.indirect_call.built.after.mir
 #[custom_mir(dialect = "built")]
 fn indirect_call(x: i32, f: fn(i32) -> i32) -> i32 {
diff --git a/tests/mir-opt/building/custom/terminators.tail_call.built.after.mir b/tests/mir-opt/building/custom/terminators.tail_call.built.after.mir
new file mode 100644
index 0000000000000..4cf6e459aa820
--- /dev/null
+++ b/tests/mir-opt/building/custom/terminators.tail_call.built.after.mir
@@ -0,0 +1,11 @@
+// MIR for `tail_call` after built
+
+fn tail_call(_1: i32) -> i32 {
+    let mut _0: i32;
+    let mut _2: i32;
+
+    bb0: {
+        _2 = Add(_1, const 42_i32);
+        tailcall ident::<i32>(Spanned { node: _2, span: $DIR/terminators.rs:32:28: 32:29 (#0) });
+    }
+}