Rollup merge of #138073 - tmiasko:inline-asm-critical-edges, r=bjorn3

Break critical edges in inline asm before code generation

An inline asm terminator defines outputs along its target edges -- a
fallthrough target and labeled targets. Code generation implements this
by inserting code directly into the target blocks. This approach works
only if the target blocks don't have other predecessors.

Establish required invariant by extending existing code that breaks
critical edges before code generation.

Fixes #137867.

r? ``@bjorn3``

2 files changed, 81 insertions, 17 deletions

diff --git a/compiler/rustc_mir_transform/src/add_call_guards.rs b/compiler/rustc_mir_transform/src/add_call_guards.rs
index 55694cacd92..bc335cee147 100644
--- a/compiler/rustc_mir_transform/src/add_call_guards.rs
+++ b/compiler/rustc_mir_transform/src/add_call_guards.rs
@@ -40,6 +40,16 @@ impl<'tcx> crate::MirPass<'tcx> for AddCallGuards {
         let mut new_blocks = Vec::new();
 
         let cur_len = body.basic_blocks.len();
+        let mut new_block = |source_info: SourceInfo, is_cleanup: bool, target: BasicBlock| {
+            let block = BasicBlockData {
+                statements: vec![],
+                is_cleanup,
+                terminator: Some(Terminator { source_info, kind: TerminatorKind::Goto { target } }),
+            };
+            let idx = cur_len + new_blocks.len();
+            new_blocks.push(block);
+            BasicBlock::new(idx)
+        };
 
         for block in body.basic_blocks_mut() {
             match block.terminator {
@@ -47,25 +57,34 @@ impl<'tcx> crate::MirPass<'tcx> for AddCallGuards {
                     kind: TerminatorKind::Call { target: Some(ref mut destination), unwind, .. },
                     source_info,
                 }) if pred_count[*destination] > 1
-                    && (matches!(
-                        unwind,
-                        UnwindAction::Cleanup(_) | UnwindAction::Terminate(_)
-                    ) || self == &AllCallEdges) =>
+                    && (generates_invoke(unwind) || self == &AllCallEdges) =>
                 {
                     // It's a critical edge, break it
-                    let call_guard = BasicBlockData {
-                        statements: vec![],
-                        is_cleanup: block.is_cleanup,
-                        terminator: Some(Terminator {
-                            source_info,
-                            kind: TerminatorKind::Goto { target: *destination },
-                        }),
-                    };
-
-                    // Get the index it will be when inserted into the MIR
-                    let idx = cur_len + new_blocks.len();
-                    new_blocks.push(call_guard);
-                    *destination = BasicBlock::new(idx);
+                    *destination = new_block(source_info, block.is_cleanup, *destination);
+                }
+                Some(Terminator {
+                    kind:
+                        TerminatorKind::InlineAsm {
+                            asm_macro: InlineAsmMacro::Asm,
+                            ref mut targets,
+                            ref operands,
+                            unwind,
+                            ..
+                        },
+                    source_info,
+                }) if self == &CriticalCallEdges => {
+                    let has_outputs = operands.iter().any(|op| {
+                        matches!(op, InlineAsmOperand::InOut { .. } | InlineAsmOperand::Out { .. })
+                    });
+                    let has_labels =
+                        operands.iter().any(|op| matches!(op, InlineAsmOperand::Label { .. }));
+                    if has_outputs && (has_labels || generates_invoke(unwind)) {
+                        for target in targets.iter_mut() {
+                            if pred_count[*target] > 1 {
+                                *target = new_block(source_info, block.is_cleanup, *target);
+                            }
+                        }
+                    }
                 }
                 _ => {}
             }
@@ -80,3 +99,11 @@ impl<'tcx> crate::MirPass<'tcx> for AddCallGuards {
         true
     }
 }
+
+/// Returns true if this unwind action is code generated as an invoke as opposed to a call.
+fn generates_invoke(unwind: UnwindAction) -> bool {
+    match unwind {
+        UnwindAction::Continue | UnwindAction::Unreachable => false,
+        UnwindAction::Cleanup(_) | UnwindAction::Terminate(_) => true,
+    }
+}
diff --git a/tests/codegen/asm/critical.rs b/tests/codegen/asm/critical.rs
new file mode 100644
index 00000000000..8c039900cab
--- /dev/null
+++ b/tests/codegen/asm/critical.rs
@@ -0,0 +1,37 @@
+//@ only-x86_64
+//@ compile-flags: -C no-prepopulate-passes
+#![feature(asm_goto)]
+#![feature(asm_goto_with_outputs)]
+#![crate_type = "lib"]
+use std::arch::asm;
+
+// Regression test for #137867. Check that critical edges have been split before code generation,
+// and so all stores to the asm output occur on disjoint paths without any of them jumping to
+// another callbr label.
+//
+// CHECK-LABEL: @f(
+// CHECK:        [[OUT:%.*]] = callbr i32 asm
+// CHECK-NEXT:   to label %[[BB0:.*]] [label %[[BB1:.*]], label %[[BB2:.*]]],
+// CHECK:       [[BB1]]:
+// CHECK-NEXT:    store i32 [[OUT]], ptr %a
+// CHECK-NEXT:    br label %[[BBR:.*]]
+// CHECK:       [[BB2]]:
+// CHECK-NEXT:    store i32 [[OUT]], ptr %a
+// CHECK-NEXT:    br label %[[BBR]]
+// CHECK:       [[BB0]]:
+// CHECK-NEXT:    store i32 [[OUT]], ptr %a
+// CHECK-NEXT:    br label %[[BBR]]
+// CHECK:       [[BBR]]:
+// CHECK-NEXT:    [[RET:%.*]] = load i32, ptr %a
+// CHECK-NEXT:    ret i32 [[RET]]
+#[unsafe(no_mangle)]
+pub unsafe fn f(mut a: u32) -> u32 {
+    asm!(
+        "jmp {}
+         jmp {}",
+        label {},
+        label {},
+        inout("eax") a,
+    );
+    a
+}