Rollup merge of rust-lang#144232 - xacrimon:explicit-tail-call, r=WaffleLapkin

Implement support for `become` and explicit tail call codegen for the LLVM backend

This PR implements codegen of explicit tail calls via `become` in `rustc_codegen_ssa` and support within the LLVM backend. Completes a task on (rust-lang#112788). This PR implements all the necessary bits to make explicit tail calls usable, other backends have received stubs for now and will ICE if you use `become` on them. I suspect there is some bikeshedding to be done on how we should go about implementing this for other backends, but it should be relatively straightforward for GCC after this is merged.

During development I also put together a POC bytecode VM based on tail call dispatch to test these changes out and analyze the codegen to make sure it generates expected assembly. That is available [here](https://github.com/xacrimon/tcvm).

Author: Zalathar

compiler/rustc_codegen_gcc/messages.ftl

@@ -4,3 +4,5 @@ codegen_gcc_unwinding_inline_asm =
 codegen_gcc_copy_bitcode = failed to copy bitcode to object file: {$err}
 
 codegen_gcc_lto_bitcode_from_rlib = failed to get bitcode from object file for LTO ({$gcc_err})
+
+codegen_gcc_explicit_tail_calls_unsupported = explicit tail calls with the 'become' keyword are not implemented in the GCC backend

compiler/rustc_codegen_gcc/src/builder.rs

@@ -34,6 +34,7 @@ use rustc_target::spec::{HasTargetSpec, HasX86AbiOpt, Target, X86Abi};
 
 use crate::common::{SignType, TypeReflection, type_is_pointer};
 use crate::context::CodegenCx;
+use crate::errors;
 use crate::intrinsic::llvm;
 use crate::type_of::LayoutGccExt;
 
@@ -1742,6 +1743,20 @@ impl<'a, 'gcc, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'gcc, 'tcx> {
         call
     }
 
+    fn tail_call(
+        &mut self,
+        _llty: Self::Type,
+        _fn_attrs: Option<&CodegenFnAttrs>,
+        _fn_abi: &FnAbi<'tcx, Ty<'tcx>>,
+        _llfn: Self::Value,
+        _args: &[Self::Value],
+        _funclet: Option<&Self::Funclet>,
+        _instance: Option<Instance<'tcx>>,
+    ) {
+        // FIXME: implement support for explicit tail calls like rustc_codegen_llvm.
+        self.tcx.dcx().emit_fatal(errors::ExplicitTailCallsUnsupported);
+    }
+
     fn zext(&mut self, value: RValue<'gcc>, dest_typ: Type<'gcc>) -> RValue<'gcc> {
         // FIXME(antoyo): this does not zero-extend.
         self.gcc_int_cast(value, dest_typ)

compiler/rustc_codegen_gcc/src/errors.rs

@@ -19,3 +19,7 @@ pub(crate) struct CopyBitcode {
 pub(crate) struct LtoBitcodeFromRlib {
     pub gcc_err: String,
 }
+
+#[derive(Diagnostic)]
+#[diag(codegen_gcc_explicit_tail_calls_unsupported)]
+pub(crate) struct ExplicitTailCallsUnsupported;

compiler/rustc_codegen_llvm/src/builder.rs

@@ -15,6 +15,7 @@ use rustc_codegen_ssa::mir::place::PlaceRef;
 use rustc_codegen_ssa::traits::*;
 use rustc_data_structures::small_c_str::SmallCStr;
 use rustc_hir::def_id::DefId;
+use rustc_middle::bug;
 use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrs;
 use rustc_middle::ty::layout::{
     FnAbiError, FnAbiOfHelpers, FnAbiRequest, HasTypingEnv, LayoutError, LayoutOfHelpers,
@@ -24,7 +25,7 @@ use rustc_middle::ty::{self, Instance, Ty, TyCtxt};
 use rustc_sanitizers::{cfi, kcfi};
 use rustc_session::config::OptLevel;
 use rustc_span::Span;
-use rustc_target::callconv::FnAbi;
+use rustc_target::callconv::{FnAbi, PassMode};
 use rustc_target::spec::{HasTargetSpec, SanitizerSet, Target};
 use smallvec::SmallVec;
 use tracing::{debug, instrument};
@@ -1431,6 +1432,28 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         call
     }
 
+    fn tail_call(
+        &mut self,
+        llty: Self::Type,
+        fn_attrs: Option<&CodegenFnAttrs>,
+        fn_abi: &FnAbi<'tcx, Ty<'tcx>>,
+        llfn: Self::Value,
+        args: &[Self::Value],
+        funclet: Option<&Self::Funclet>,
+        instance: Option<Instance<'tcx>>,
+    ) {
+        let call = self.call(llty, fn_attrs, Some(fn_abi), llfn, args, funclet, instance);
+        llvm::LLVMRustSetTailCallKind(call, llvm::TailCallKind::MustTail);
+
+        match &fn_abi.ret.mode {
+            PassMode::Ignore | PassMode::Indirect { .. } => self.ret_void(),
+            PassMode::Direct(_) | PassMode::Pair { .. } => self.ret(call),
+            mode @ PassMode::Cast { .. } => {
+                bug!("Encountered `PassMode::{mode:?}` during codegen")
+            }
+        }
+    }
+
     fn zext(&mut self, val: &'ll Value, dest_ty: &'ll Type) -> &'ll Value {
         unsafe { llvm::LLVMBuildZExt(self.llbuilder, val, dest_ty, UNNAMED) }
     }

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -97,6 +97,16 @@ pub(crate) enum ModuleFlagMergeBehavior {
 
 // Consts for the LLVM CallConv type, pre-cast to usize.
 
+#[derive(Copy, Clone, PartialEq, Debug)]
+#[repr(C)]
+#[allow(dead_code)]
+pub(crate) enum TailCallKind {
+    None = 0,
+    Tail = 1,
+    MustTail = 2,
+    NoTail = 3,
+}
+
 /// LLVM CallingConv::ID. Should we wrap this?
 ///
 /// See <https://github.com/llvm/llvm-project/blob/main/llvm/include/llvm/IR/CallingConv.h>
@@ -1186,6 +1196,7 @@ unsafe extern "C" {
     pub(crate) safe fn LLVMIsGlobalConstant(GlobalVar: &Value) -> Bool;
     pub(crate) safe fn LLVMSetGlobalConstant(GlobalVar: &Value, IsConstant: Bool);
     pub(crate) safe fn LLVMSetTailCall(CallInst: &Value, IsTailCall: Bool);
+    pub(crate) safe fn LLVMRustSetTailCallKind(CallInst: &Value, Kind: TailCallKind);
 
     // Operations on attributes
     pub(crate) fn LLVMCreateStringAttribute(

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -35,6 +35,14 @@ enum MergingSucc {
     True,
 }
 
+/// Indicates to the call terminator codegen whether a cal
+/// is a normal call or an explicit tail call.
+#[derive(Debug, PartialEq)]
+enum CallKind {
+    Normal,
+    Tail,
+}
+
 /// Used by `FunctionCx::codegen_terminator` for emitting common patterns
 /// e.g., creating a basic block, calling a function, etc.
 struct TerminatorCodegenHelper<'tcx> {
@@ -160,6 +168,7 @@ impl<'a, 'tcx> TerminatorCodegenHelper<'tcx> {
         mut unwind: mir::UnwindAction,
         lifetime_ends_after_call: &[(Bx::Value, Size)],
         instance: Option<Instance<'tcx>>,
+        kind: CallKind,
         mergeable_succ: bool,
     ) -> MergingSucc {
         let tcx = bx.tcx();
@@ -221,6 +230,11 @@ impl<'a, 'tcx> TerminatorCodegenHelper<'tcx> {
             }
         };
 
+        if kind == CallKind::Tail {
+            bx.tail_call(fn_ty, fn_attrs, fn_abi, fn_ptr, llargs, self.funclet(fx), instance);
+            return MergingSucc::False;
+        }
+
         if let Some(unwind_block) = unwind_block {
             let ret_llbb = if let Some((_, target)) = destination {
                 fx.llbb(target)
@@ -659,6 +673,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             unwind,
             &[],
             Some(drop_instance),
+            CallKind::Normal,
             !maybe_null && mergeable_succ,
         )
     }
@@ -747,8 +762,19 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         let (fn_abi, llfn, instance) = common::build_langcall(bx, span, lang_item);
 
         // Codegen the actual panic invoke/call.
-        let merging_succ =
-            helper.do_call(self, bx, fn_abi, llfn, &args, None, unwind, &[], Some(instance), false);
+        let merging_succ = helper.do_call(
+            self,
+            bx,
+            fn_abi,
+            llfn,
+            &args,
+            None,
+            unwind,
+            &[],
+            Some(instance),
+            CallKind::Normal,
+            false,
+        );
         assert_eq!(merging_succ, MergingSucc::False);
         MergingSucc::False
     }
@@ -777,6 +803,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             mir::UnwindAction::Unreachable,
             &[],
             Some(instance),
+            CallKind::Normal,
             false,
         );
         assert_eq!(merging_succ, MergingSucc::False);
@@ -845,6 +872,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             unwind,
             &[],
             Some(instance),
+            CallKind::Normal,
             mergeable_succ,
         ))
     }
@@ -860,6 +888,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         target: Option<mir::BasicBlock>,
         unwind: mir::UnwindAction,
         fn_span: Span,
+        kind: CallKind,
         mergeable_succ: bool,
     ) -> MergingSucc {
         let source_info = mir::SourceInfo { span: fn_span, ..terminator.source_info };
@@ -1003,8 +1032,13 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         // We still need to call `make_return_dest` even if there's no `target`, since
         // `fn_abi.ret` could be `PassMode::Indirect`, even if it is uninhabited,
         // and `make_return_dest` adds the return-place indirect pointer to `llargs`.
-        let return_dest = self.make_return_dest(bx, destination, &fn_abi.ret, &mut llargs);
-        let destination = target.map(|target| (return_dest, target));
+        let destination = match kind {
+            CallKind::Normal => {
+                let return_dest = self.make_return_dest(bx, destination, &fn_abi.ret, &mut llargs);
+                target.map(|target| (return_dest, target))
+            }
+            CallKind::Tail => None,
+        };
 
         // Split the rust-call tupled arguments off.
         let (first_args, untuple) = if sig.abi() == ExternAbi::RustCall
@@ -1020,6 +1054,14 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         // to generate `lifetime_end` when the call returns.
         let mut lifetime_ends_after_call: Vec<(Bx::Value, Size)> = Vec::new();
         'make_args: for (i, arg) in first_args.iter().enumerate() {
+            if kind == CallKind::Tail && matches!(fn_abi.args[i].mode, PassMode::Indirect { .. }) {
+                // FIXME: https://github.com/rust-lang/rust/pull/144232#discussion_r2218543841
+                span_bug!(
+                    fn_span,
+                    "arguments using PassMode::Indirect are currently not supported for tail calls"
+                );
+            }
+
             let mut op = self.codegen_operand(bx, &arg.node);
 
             if let (0, Some(ty::InstanceKind::Virtual(_, idx))) = (i, instance.map(|i| i.def)) {
@@ -1147,6 +1189,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             unwind,
             &lifetime_ends_after_call,
             instance,
+            kind,
             mergeable_succ,
         )
     }
@@ -1388,15 +1431,23 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 target,
                 unwind,
                 fn_span,
+                CallKind::Normal,
                 mergeable_succ(),
             ),
-            mir::TerminatorKind::TailCall { .. } => {
-                // FIXME(explicit_tail_calls): implement tail calls in ssa backend
-                span_bug!(
-                    terminator.source_info.span,
-                    "`TailCall` terminator is not yet supported by `rustc_codegen_ssa`"
-                )
-            }
+            mir::TerminatorKind::TailCall { ref func, ref args, fn_span } => self
+                .codegen_call_terminator(
+                    helper,
+                    bx,
+                    terminator,
+                    func,
+                    args,
+                    mir::Place::from(mir::RETURN_PLACE),
+                    None,
+                    mir::UnwindAction::Unreachable,
+                    fn_span,
+                    CallKind::Tail,
+                    mergeable_succ(),
+                ),
             mir::TerminatorKind::CoroutineDrop | mir::TerminatorKind::Yield { .. } => {
                 bug!("coroutine ops in codegen")
             }

compiler/rustc_codegen_ssa/src/traits/builder.rs

@@ -595,6 +595,18 @@ pub trait BuilderMethods<'a, 'tcx>:
         funclet: Option<&Self::Funclet>,
         instance: Option<Instance<'tcx>>,
     ) -> Self::Value;
+
+    fn tail_call(
+        &mut self,
+        llty: Self::Type,
+        fn_attrs: Option<&CodegenFnAttrs>,
+        fn_abi: &FnAbi<'tcx, Ty<'tcx>>,
+        llfn: Self::Value,
+        args: &[Self::Value],
+        funclet: Option<&Self::Funclet>,
+        instance: Option<Instance<'tcx>>,
+    );
+
     fn zext(&mut self, val: Self::Value, dest_ty: Self::Type) -> Self::Value;
 
     fn apply_attrs_to_cleanup_callsite(&mut self, llret: Self::Value);

compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp

@@ -1986,3 +1986,29 @@ extern "C" void LLVMRustSetNoSanitizeHWAddress(LLVMValueRef Global) {
   MD.NoHWAddress = true;
   GV.setSanitizerMetadata(MD);
 }
+
+enum class LLVMRustTailCallKind {
+  None = 0,
+  Tail = 1,
+  MustTail = 2,
+  NoTail = 3
+};
+
+extern "C" void LLVMRustSetTailCallKind(LLVMValueRef Call,
+                                        LLVMRustTailCallKind Kind) {
+  CallInst *CI = unwrap<CallInst>(Call);
+  switch (Kind) {
+  case LLVMRustTailCallKind::None:
+    CI->setTailCallKind(CallInst::TCK_None);
+    break;
+  case LLVMRustTailCallKind::Tail:
+    CI->setTailCallKind(CallInst::TCK_Tail);
+    break;
+  case LLVMRustTailCallKind::MustTail:
+    CI->setTailCallKind(CallInst::TCK_MustTail);
+    break;
+  case LLVMRustTailCallKind::NoTail:
+    CI->setTailCallKind(CallInst::TCK_NoTail);
+    break;
+  }
+}

tests/codegen-llvm/become-musttail.rs

@@ -0,0 +1,18 @@
+//@ compile-flags: -C opt-level=0 -Cpanic=abort -C no-prepopulate-passes
+//@ needs-unwind
+
+#![crate_type = "lib"]
+#![feature(explicit_tail_calls)]
+
+// CHECK-LABEL: define {{.*}}@fibonacci(
+#[no_mangle]
+#[inline(never)]
+pub fn fibonacci(n: u64, a: u64, b: u64) -> u64 {
+    // CHECK: musttail call {{.*}}@fibonacci(
+    // CHECK-NEXT: ret i64
+    match n {
+        0 => a,
+        1 => b,
+        _ => become fibonacci(n - 1, b, a + b),
+    }
+}

tests/ui/explicit-tail-calls/recursion-etc.rs

@@ -0,0 +1,17 @@
+//@ run-pass
+#![expect(incomplete_features)]
+#![feature(explicit_tail_calls)]
+
+use std::hint::black_box;
+
+pub fn count(curr: u64, top: u64) -> u64 {
+   if black_box(curr) >= top {
+        curr
+   } else {
+        become count(curr + 1, top)
+   }
+}
+
+fn main() {
+    println!("{}", count(0, black_box(1000000)));
+}