From 20a099c481186fcec7060fef126cfe09bda61316 Mon Sep 17 00:00:00 2001
From: Maurice Heumann <maurice.heumann@wibu.com>
Date: Thu, 19 Sep 2024 10:05:43 +0200
Subject: [PATCH] [ARM64] [Windows] Mark block address as taken when expanding
 catchrets

This fixes issue #109250

The issue happens during the `MachineBlockPlacement` pass.
The block, whose address was previously not taken, is deemed redundant
by the pass and subsequently replaced using
`MachineBasicBlock::ReplaceUsesOfBlockWith` in `BranchFolding`.

ReplaceUsesOfBlockWith only replaces uses in the terminator.
However, `expandPostRAPseudo` introduces new block uses
when expanding catchrets. These uses do not get replaced, which results
in undefined label errors later on.

Marking the block addresss as taken prevents the replacement of the block,
without also replacing non-terminator uses.
---
 llvm/lib/Target/AArch64/AArch64InstrInfo.cpp  |  1 +
 llvm/test/CodeGen/AArch64/pr58516.ll          |  3 +-
 .../AArch64/wineh-catchret-obj-generation.ll  | 45 +++++++++++++++++++
 llvm/test/CodeGen/AArch64/wineh-try-catch.ll  |  2 +-
 4 files changed, 49 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/CodeGen/AArch64/wineh-catchret-obj-generation.ll

diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
index 3b38a5f78dee51..32bc0e7d0d6475 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
@@ -1994,6 +1994,7 @@ bool AArch64InstrInfo::expandPostRAPseudo(MachineInstr &MI) const {
         .addReg(AArch64::X0)
         .addMBB(TargetMBB)
         .addImm(0);
+    TargetMBB->setMachineBlockAddressTaken();
     return true;
   }
 
diff --git a/llvm/test/CodeGen/AArch64/pr58516.ll b/llvm/test/CodeGen/AArch64/pr58516.ll
index b4840f01ce116b..3361ded48d4e2a 100644
--- a/llvm/test/CodeGen/AArch64/pr58516.ll
+++ b/llvm/test/CodeGen/AArch64/pr58516.ll
@@ -40,7 +40,8 @@ define void @osfx(ptr %this) comdat personality ptr @__CxxFrameHandler3 {
 ; CHECK-NEXT:  // %bb.1: // %invoke.cont12
 ; CHECK-NEXT:    str wzr, [x20]
 ; CHECK-NEXT:    str wzr, [x21]
-; CHECK-NEXT:  .LBB0_2: // %try.cont
+; CHECK-NEXT:  .LBB0_2: // Block address taken
+; CHECK-NEXT:    // %try.cont
 ; CHECK-NEXT:  $ehgcr_0_2:
 ; CHECK-NEXT:    .seh_startepilogue
 ; CHECK-NEXT:    sub sp, x29, #24
diff --git a/llvm/test/CodeGen/AArch64/wineh-catchret-obj-generation.ll b/llvm/test/CodeGen/AArch64/wineh-catchret-obj-generation.ll
new file mode 100644
index 00000000000000..68f5cf843ddee3
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/wineh-catchret-obj-generation.ll
@@ -0,0 +1,45 @@
+; RUN: llc -mtriple aarch64-unknown-windows-msvc -filetype=obj %s -o %t.o
+
+declare i32 @__CxxFrameHandler3(...)
+
+declare void @llvm.seh.try.begin() #0
+
+define fastcc ptr @test_function(i1 %0, ptr %_Fmtfl.i.i, i1 %1) personality ptr @__CxxFrameHandler3 {
+entry:
+  br i1 %0, label %right-block527, label %left-block526
+
+common.ret1:
+  %common.ret1.op = phi ptr [ null, %left-block530 ], [ null, %some-block ], [ %_Fmtfl.i.i, %invoke.cont.i124 ], [ null, %left-block526 ]
+  ret ptr %common.ret1.op
+
+invoke.cont.i124:
+  %.not657 = icmp eq i32 1, 0
+  br i1 %.not657, label %some-block, label %common.ret1
+
+catch.dispatch.i:
+  %2 = catchswitch within none [label %catch.i] unwind to caller
+
+catch.i:
+  %3 = catchpad within %2 [ptr null, i32 0, ptr null]
+  catchret from %3 to label %some-block
+
+some-block:
+  br label %common.ret1
+
+left-block526:
+  br i1 %1, label %common.ret1, label %left-block530
+
+right-block527:
+  invoke void @llvm.seh.try.begin()
+          to label %invoke.cont.i124 unwind label %catch.dispatch.i
+
+left-block530:
+  %.not = icmp eq i32 0, 0
+  br label %common.ret1
+}
+
+attributes #0 = { nounwind willreturn memory(write) }
+
+!llvm.module.flags = !{!0}
+
+!0 = !{i32 2, !"eh-asynch", i32 1}
diff --git a/llvm/test/CodeGen/AArch64/wineh-try-catch.ll b/llvm/test/CodeGen/AArch64/wineh-try-catch.ll
index b27e5374b25762..c3b5a8968d7bba 100644
--- a/llvm/test/CodeGen/AArch64/wineh-try-catch.ll
+++ b/llvm/test/CodeGen/AArch64/wineh-try-catch.ll
@@ -41,7 +41,7 @@
 ; CHECK-LABEL: .Ltmp0:
 ; CHECK:       bl      "?func2@@YAHXZ
 
-; CHECK:        [[CATCHRETDEST:.LBB0_[0-9]+]]:      // %catchret.dest
+; CHECK:        [[CATCHRETDEST:.LBB0_[0-9]+]]:      // Block address taken
 
 ; Check the catch funclet.
 ; CHECK-LABEL: "?catch$4@?0??func@@YAHXZ@4HA":