From 73f8309300bce03e62c8a49e8a06b884175b5f88 Mon Sep 17 00:00:00 2001
From: Gary Guo <gary@garyguo.net>
Date: Fri, 11 Oct 2024 00:26:21 +0100
Subject: [PATCH] Support use of asm goto with outputs and `options(noreturn)`

When labels are present, the `noreturn` option really means that asm block
won't fallthrough -- if labels are present, then outputs can still be
meaningfully used.
---
 compiler/rustc_builtin_macros/src/asm.rs |  5 ++++-
 compiler/rustc_codegen_llvm/src/asm.rs   |  9 ++++++++-
 tests/codegen/asm/goto.rs                | 12 +++++++++++-
 tests/ui/asm/x86_64/goto.rs              | 20 ++++++++++++++++++++
 tests/ui/asm/x86_64/goto.stderr          |  4 ++--
 5 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/compiler/rustc_builtin_macros/src/asm.rs b/compiler/rustc_builtin_macros/src/asm.rs
index 9ae48024f44..14ac3cd74e8 100644
--- a/compiler/rustc_builtin_macros/src/asm.rs
+++ b/compiler/rustc_builtin_macros/src/asm.rs
@@ -300,7 +300,10 @@ pub fn parse_asm_args<'a>(
     if args.options.contains(ast::InlineAsmOptions::PURE) && !have_real_output {
         dcx.emit_err(errors::AsmPureNoOutput { spans: args.options_spans.clone() });
     }
-    if args.options.contains(ast::InlineAsmOptions::NORETURN) && !outputs_sp.is_empty() {
+    if args.options.contains(ast::InlineAsmOptions::NORETURN)
+        && !outputs_sp.is_empty()
+        && labels_sp.is_empty()
+    {
         let err = dcx.create_err(errors::AsmNoReturn { outputs_sp });
         // Bail out now since this is likely to confuse MIR
         return Err(err);
diff --git a/compiler/rustc_codegen_llvm/src/asm.rs b/compiler/rustc_codegen_llvm/src/asm.rs
index 2455d13df8d..07473190d6f 100644
--- a/compiler/rustc_codegen_llvm/src/asm.rs
+++ b/compiler/rustc_codegen_llvm/src/asm.rs
@@ -343,7 +343,14 @@ impl<'ll, 'tcx> AsmBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
         attributes::apply_to_callsite(result, llvm::AttributePlace::Function, &{ attrs });
 
         // Write results to outputs. We need to do this for all possible control flow.
-        for block in Some(dest).into_iter().chain(labels.iter().copied().map(Some)) {
+        //
+        // Note that `dest` maybe populated with unreachable_block when asm goto with outputs
+        // is used (because we need to codegen callbr which always needs a destination), so
+        // here we use the NORETURN option to determine if `dest` should be used.
+        for block in (if options.contains(InlineAsmOptions::NORETURN) { None } else { Some(dest) })
+            .into_iter()
+            .chain(labels.iter().copied().map(Some))
+        {
             if let Some(block) = block {
                 self.switch_to_block(block);
             }
diff --git a/tests/codegen/asm/goto.rs b/tests/codegen/asm/goto.rs
index 0e69c4d5840..3193d3aa145 100644
--- a/tests/codegen/asm/goto.rs
+++ b/tests/codegen/asm/goto.rs
@@ -43,11 +43,21 @@ pub unsafe fn asm_goto_with_outputs_use_in_label() -> u64 {
 // CHECK-LABEL: @asm_goto_noreturn
 #[no_mangle]
 pub unsafe fn asm_goto_noreturn() -> u64 {
-    let out: u64;
     // CHECK: callbr void asm sideeffect alignstack inteldialect "
     // CHECK-NEXT: to label %unreachable [label %[[JUMPBB:[a-b0-9]+]]]
     asm!("jmp {}", label { return 1; }, options(noreturn));
     // CHECK: [[JUMPBB]]:
     // CHECK-NEXT: ret i64 1
+}
+
+// CHECK-LABEL: @asm_goto_noreturn_with_outputs
+#[no_mangle]
+pub unsafe fn asm_goto_noreturn_with_outputs() -> u64 {
+    let out: u64;
+    // CHECK: [[RES:%[0-9]+]] = callbr i64 asm sideeffect alignstack inteldialect "
+    // CHECK-NEXT: to label %[[FALLTHROUGHBB:[a-b0-9]+]] [label %[[JUMPBB:[a-b0-9]+]]]
+    asm!("mov {}, 1", "jmp {}", out(reg) out, label { return out; });
+    // CHECK: [[JUMPBB]]:
+    // CHECK-NEXT: ret i64 [[RES]]
     out
 }
diff --git a/tests/ui/asm/x86_64/goto.rs b/tests/ui/asm/x86_64/goto.rs
index ab22317c458..d23018a05f5 100644
--- a/tests/ui/asm/x86_64/goto.rs
+++ b/tests/ui/asm/x86_64/goto.rs
@@ -65,6 +65,25 @@ fn goto_out_jump() {
     }
 }
 
+fn goto_out_jump_noreturn() {
+    unsafe {
+        let mut value = false;
+        let mut out: usize;
+        asm!(
+            "lea {}, [{} + 1]",
+            "jmp {}",
+            out(reg) out,
+            in(reg) 0x12345678usize,
+            label {
+                value = true;
+                assert_eq!(out, 0x12345679);
+            },
+            options(noreturn)
+        );
+        assert!(value);
+    }
+}
+
 // asm goto with outputs cause miscompilation in LLVM when multiple outputs are present.
 // The code sample below is adapted from https://github.com/llvm/llvm-project/issues/74483
 // and does not work with `-C opt-level=0`
@@ -131,6 +150,7 @@ fn main() {
     goto_jump();
     goto_out_fallthrough();
     goto_out_jump();
+    goto_out_jump_noreturn();
     // goto_multi_out();
     goto_noreturn();
     goto_noreturn_diverge();
diff --git a/tests/ui/asm/x86_64/goto.stderr b/tests/ui/asm/x86_64/goto.stderr
index 216b4d252ec..8f89e2b1290 100644
--- a/tests/ui/asm/x86_64/goto.stderr
+++ b/tests/ui/asm/x86_64/goto.stderr
@@ -1,5 +1,5 @@
 warning: unreachable statement
-  --> $DIR/goto.rs:124:9
+  --> $DIR/goto.rs:143:9
    |
 LL | /         asm!(
 LL | |             "jmp {}",
@@ -13,7 +13,7 @@ LL |           unreachable!();
    |           ^^^^^^^^^^^^^^ unreachable statement
    |
 note: the lint level is defined here
-  --> $DIR/goto.rs:114:8
+  --> $DIR/goto.rs:133:8
    |
 LL | #[warn(unreachable_code)]
    |        ^^^^^^^^^^^^^^^^