[llvm] [CodeGen][ARM64EC] Mangle EH personality handler names (PR #121652)

[Billy Laws via llvm-commits]

https://github.com/bylaws created https://github.com/llvm/llvm-project/pull/121652

Since the target isn't passed to `llvm::classifyEHPersonality` this currently demangles the handler name unconditionally when classifying, this shouldn't cause any issues but if preferred I can change that.

>From 05019aca4704dd8d03eb08e087483c8529be9cc3 Mon Sep 17 00:00:00 2001
From: Billy Laws <blaws05 at gmail.com>
Date: Thu, 12 Dec 2024 15:23:00 +0000
Subject: [PATCH] [CodeGen][ARM64EC] Mangle EH personality handler names

---
 llvm/lib/IR/EHPersonalities.cpp               |  8 ++-
 .../AArch64/AArch64Arm64ECCallLowering.cpp    | 10 +++
 llvm/test/CodeGen/AArch64/arm64ec-eh.ll       | 65 +++++++++++++++++++
 3 files changed, 82 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/AArch64/arm64ec-eh.ll

diff --git a/llvm/lib/IR/EHPersonalities.cpp b/llvm/lib/IR/EHPersonalities.cpp
index 7c32601b8a83ea..8320e738531937 100644
--- a/llvm/lib/IR/EHPersonalities.cpp
+++ b/llvm/lib/IR/EHPersonalities.cpp
@@ -12,6 +12,7 @@
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/IR/Mangler.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
@@ -25,7 +26,12 @@ EHPersonality llvm::classifyEHPersonality(const Value *Pers) {
       Pers ? dyn_cast<GlobalValue>(Pers->stripPointerCasts()) : nullptr;
   if (!F || !F->getValueType() || !F->getValueType()->isFunctionTy())
     return EHPersonality::Unknown;
-  return StringSwitch<EHPersonality>(F->getName())
+
+  std::string Name = F->getName().str();
+  if (auto DemangledName = getArm64ECDemangledFunctionName(Name))
+    Name = *DemangledName;
+
+  return StringSwitch<EHPersonality>(Name)
       .Case("__gnat_eh_personality", EHPersonality::GNU_Ada)
       .Case("__gxx_personality_v0", EHPersonality::GNU_CXX)
       .Case("__gxx_personality_seh0", EHPersonality::GNU_CXX)
diff --git a/llvm/lib/Target/AArch64/AArch64Arm64ECCallLowering.cpp b/llvm/lib/Target/AArch64/AArch64Arm64ECCallLowering.cpp
index 415edb189e60c4..f33cc5b6831ce3 100644
--- a/llvm/lib/Target/AArch64/AArch64Arm64ECCallLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64Arm64ECCallLowering.cpp
@@ -807,6 +807,16 @@ bool AArch64Arm64ECCallLowering::runOnModule(Module &Mod) {
   SetVector<GlobalAlias *> PatchableFns;
 
   for (Function &F : Mod) {
+    if (F.hasPersonalityFn()) {
+      GlobalValue *PersFn = dyn_cast<GlobalValue>(F.getPersonalityFn()->stripPointerCasts());
+      if (PersFn->getValueType() && PersFn->getValueType()->isFunctionTy()) {
+        if (std::optional<std::string> MangledName =
+            getArm64ECMangledFunctionName(PersFn->getName().str())) {
+          PersFn->setName(MangledName.value());
+        }
+      }
+    }
+
     if (!F.hasFnAttribute(Attribute::HybridPatchable) || F.isDeclaration() ||
         F.hasLocalLinkage() || F.getName().ends_with("$hp_target"))
       continue;
diff --git a/llvm/test/CodeGen/AArch64/arm64ec-eh.ll b/llvm/test/CodeGen/AArch64/arm64ec-eh.ll
new file mode 100644
index 00000000000000..4b6bbb01c6c2dc
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/arm64ec-eh.ll
@@ -0,0 +1,65 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=arm64ec-pc-windows-msvc %s -o - | FileCheck %s
+
+define dso_local i32 @test() #0 personality ptr @__C_specific_handler {
+; CHECK-LABEL:  "#test"
+; CHECK:  .Lfunc_begin0:
+; CHECK-NEXT:  .seh_proc "#test"
+; CHECK-NEXT:    .seh_handler "#__C_specific_handler", @unwind, @except
+; CHECK-NEXT:  // %bb.0:
+; CHECK-NEXT:    sub sp, sp, #48
+; CHECK-NEXT:    .seh_stackalloc 48
+; CHECK-NEXT:    stp x29, x30, [sp, #16]             // 16-byte Folded Spill
+; CHECK-NEXT:    .seh_save_fplr 16
+; CHECK-NEXT:    add x29, sp, #16
+; CHECK-NEXT:    .seh_add_fp 16
+; CHECK-NEXT:    .seh_endprologue
+; CHECK-NEXT:    mov x0, #-2                         // =0xfffffffffffffffe
+; CHECK-NEXT:    mov w8, #-1                         // =0xffffffff
+; CHECK-NEXT:    stur x0, [x29, #16]
+; CHECK-NEXT:    stur w8, [x29, #-4]
+; CHECK-NEXT:  .Ltmp0:
+; 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
+; CHECK-NEXT:    .seh_save_fplr  16
+; CHECK-NEXT:    add     sp, sp, #48
+; CHECK-NEXT:    .seh_stackalloc 48
+; CHECK-NEXT:    .seh_endepilogue
+; CHECK-NEXT:    ret
+; CHECK-NEXT:  .LBB0_2:
+; CHECK-NEXT:    mov     w0, wzr
+; CHECK-NEXT:    b       .LBB0_1
+  %1 = alloca i32, align 4
+  %2 = alloca i32, align 4
+  store i32 -1, ptr %1, align 4
+  %3 = invoke i32 @ext() #3
+          to label %12 unwind label %4
+
+4:                                                ; preds = %0
+  %5 = catchswitch within none [label %6] unwind to caller
+
+6:                                                ; preds = %4
+  %7 = catchpad within %5 [ptr null]
+  catchret from %7 to label %8
+
+8:                                                ; preds = %6
+  store i32 0, ptr %1, align 4
+  br label %10
+
+10:                                               ; preds = %8, %12
+  %11 = load i32, ptr %1, align 4
+  ret i32 %11
+
+12:                                               ; preds = %0
+  store i32 %3, ptr %1, align 4
+  br label %10
+}
+
+declare dso_local i32 @ext() #1
+
+declare dso_local i32 @__C_specific_handler(...)