Break critical edges in inline asm before code generation

compiler/rustc_mir_transform/src/add_call_guards.rs

@@ -40,32 +40,51 @@ 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 {
                 Some(Terminator {
                     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,
+    }
+}

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
+}