[llvm] [win] Fix EH Cont Guard targets when SEH personality is used (PR #129612)

[Daniel Paoliello via llvm-commits]

https://github.com/dpaoliello updated https://github.com/llvm/llvm-project/pull/129612

>From 5cf266d173825c9c3fa02daa003b3671f417545d Mon Sep 17 00:00:00 2001
From: "Daniel Paoliello (HE/HIM)" <danpao at microsoft.com>
Date: Mon, 3 Mar 2025 09:57:45 -0800
Subject: [PATCH] [win] Fix EH Cont Guard targets when SEH personality is used

---
 llvm/lib/CodeGen/AsmPrinter/WinException.cpp         |  6 ++++++
 .../lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp | 12 +++++++++---
 llvm/test/CodeGen/AArch64/arm64ec-eh.ll              |  2 +-
 llvm/test/CodeGen/AArch64/landingpad-ifcvt.ll        |  1 +
 llvm/test/CodeGen/X86/seh-except-restore.ll          |  1 +
 llvm/test/CodeGen/X86/win32-seh-catchpad.ll          |  1 +
 6 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/CodeGen/AsmPrinter/WinException.cpp b/llvm/lib/CodeGen/AsmPrinter/WinException.cpp
index d51ac2ac89f35..062283975851a 100644
--- a/llvm/lib/CodeGen/AsmPrinter/WinException.cpp
+++ b/llvm/lib/CodeGen/AsmPrinter/WinException.cpp
@@ -290,6 +290,12 @@ void WinException::endFuncletImpl() {
       // functions that need it in the end anyway.
     }
 
+    if (!MF->getEHContTargets().empty()) {
+      // Copy the function's EH Continuation targets to a module-level list.
+      EHContTargets.insert(EHContTargets.end(), MF->getEHContTargets().begin(),
+                           MF->getEHContTargets().end());
+    }
+
     // Switch back to the funclet start .text section now that we are done
     // writing to .xdata, and emit an .seh_endproc directive to mark the end of
     // the function.
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
index d2df323fce638..1618f25f3908e 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -1970,7 +1970,11 @@ void SelectionDAGBuilder::visitCatchPad(const CatchPadInst &I) {
   bool IsCoreCLR = Pers == EHPersonality::CoreCLR;
   bool IsSEH = isAsynchronousEHPersonality(Pers);
   MachineBasicBlock *CatchPadMBB = FuncInfo.MBB;
-  if (!IsSEH)
+  if (IsSEH) {
+    // For SEH, EHCont Guard needs to know that this catchpad is a target.
+    CatchPadMBB->setIsEHContTarget(true);
+    DAG.getMachineFunction().setHasEHContTarget(true);
+  } else
     CatchPadMBB->setIsEHScopeEntry();
   // In MSVC C++ and CoreCLR, catchblocks are funclets and need prologues.
   if (IsMSVCCXX || IsCoreCLR)
@@ -1981,8 +1985,6 @@ void SelectionDAGBuilder::visitCatchRet(const CatchReturnInst &I) {
   // Update machine-CFG edge.
   MachineBasicBlock *TargetMBB = FuncInfo.getMBB(I.getSuccessor());
   FuncInfo.MBB->addSuccessor(TargetMBB);
-  TargetMBB->setIsEHContTarget(true);
-  DAG.getMachineFunction().setHasEHContTarget(true);
 
   auto Pers = classifyEHPersonality(FuncInfo.Fn->getPersonalityFn());
   bool IsSEH = isAsynchronousEHPersonality(Pers);
@@ -1996,6 +1998,10 @@ void SelectionDAGBuilder::visitCatchRet(const CatchReturnInst &I) {
     return;
   }
 
+  // For non-SEH, EHCont Guard needs to know that this catchret is a target.
+  TargetMBB->setIsEHContTarget(true);
+  DAG.getMachineFunction().setHasEHContTarget(true);
+
   // Figure out the funclet membership for the catchret's successor.
   // This will be used by the FuncletLayout pass to determine how to order the
   // BB's.
diff --git a/llvm/test/CodeGen/AArch64/arm64ec-eh.ll b/llvm/test/CodeGen/AArch64/arm64ec-eh.ll
index 4b6bbb01c6c2d..64c4a871647f1 100644
--- a/llvm/test/CodeGen/AArch64/arm64ec-eh.ll
+++ b/llvm/test/CodeGen/AArch64/arm64ec-eh.ll
@@ -22,7 +22,6 @@ define dso_local i32 @test() #0 personality ptr @__C_specific_handler {
 ; CHECK-NEXT:    bl      "#ext"
 ; CHECK-NEXT:  .Ltmp1:
 ; CHECK-NEXT:  .LBB0_1:
-; CHECK-NEXT:  $ehgcr_0_1:
 ; CHECK-NEXT:    stur    w0, [x29, #-4]
 ; CHECK-NEXT:    .seh_startepilogue
 ; CHECK-NEXT:    ldp     x29, x30, [sp, #16]             // 16-byte Folded Reload
@@ -32,6 +31,7 @@ define dso_local i32 @test() #0 personality ptr @__C_specific_handler {
 ; CHECK-NEXT:    .seh_endepilogue
 ; CHECK-NEXT:    ret
 ; CHECK-NEXT:  .LBB0_2:
+; CHECK-NEXT:  $ehgcr_0_2:
 ; CHECK-NEXT:    mov     w0, wzr
 ; CHECK-NEXT:    b       .LBB0_1
   %1 = alloca i32, align 4
diff --git a/llvm/test/CodeGen/AArch64/landingpad-ifcvt.ll b/llvm/test/CodeGen/AArch64/landingpad-ifcvt.ll
index 367b4975e7558..c6dd152d7ff54 100644
--- a/llvm/test/CodeGen/AArch64/landingpad-ifcvt.ll
+++ b/llvm/test/CodeGen/AArch64/landingpad-ifcvt.ll
@@ -2,6 +2,7 @@
 
 ; Make sure this doesn't crash (and the output is sane).
 ; CHECK: // %__except.ret
+; CHECK-NEXT: $ehgcr_0_2:
 ; CHECK-NEXT: mov     x0, xzr
 
 target datalayout = "e-m:w-p:64:64-i32:32-i64:64-i128:128-n32:64-S128"
diff --git a/llvm/test/CodeGen/X86/seh-except-restore.ll b/llvm/test/CodeGen/X86/seh-except-restore.ll
index e362f01791d06..f428a83cc2aef 100644
--- a/llvm/test/CodeGen/X86/seh-except-restore.ll
+++ b/llvm/test/CodeGen/X86/seh-except-restore.ll
@@ -48,6 +48,7 @@ return:                                           ; preds = %entry, %__except
 ; CHECK: LBB0_2:                                 # %return
 
 ; CHECK: LBB0_1:                                 # %__except.ret
+; CHECK-NEXT:         $ehgcr_0_1:
 ; CHECK-NEXT:         movl    -24(%ebp), %esp
 ; CHECK-NEXT:         addl    $12, %ebp
 
diff --git a/llvm/test/CodeGen/X86/win32-seh-catchpad.ll b/llvm/test/CodeGen/X86/win32-seh-catchpad.ll
index fa7335cb7bf64..832ddcab5227b 100644
--- a/llvm/test/CodeGen/X86/win32-seh-catchpad.ll
+++ b/llvm/test/CodeGen/X86/win32-seh-catchpad.ll
@@ -196,6 +196,7 @@ __except:
 
 ; CHECK-LABEL: _code_in_catchpad:
 ; CHECK: # %__except.ret
+; CHECK-NEXT:         $ehgcr_4_1:
 ; CHECK-NEXT:         movl    -24(%ebp), %esp
 ; CHECK-NEXT:         addl    $12, %ebp
 ; CHECK-NEXT:         movl    $-1, -16(%ebp)