[IR][ARM64EC][NFC] Clean up and document ARM64EC mangling helpers.

[cjacek]

No description provided.

[cjacek]

This PR addresses @mstorsjo's comment from #107164. Since the C++ and non-C++ code paths don't share much beyond the return statement, I've separated them further. This change eliminates the need for the Prefix variable (which triggered the original comment) and makes the non-C++ case easier to follow. I also added some comments for clarity.

[llvmbot]

@llvm/pr-subscribers-llvm-ir

Author: Jacek Caban (cjacek)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/107230.diff

2 Files Affected:

• (modified) llvm/include/llvm/IR/Mangler.h (+5)
• (modified) llvm/lib/IR/Mangler.cpp (+20-18)

diff --git a/llvm/include/llvm/IR/Mangler.h b/llvm/include/llvm/IR/Mangler.h
index f28ffc961b6db0..349f9e6e752339 100644
--- a/llvm/include/llvm/IR/Mangler.h
+++ b/llvm/include/llvm/IR/Mangler.h
@@ -53,7 +53,12 @@ void emitLinkerFlagsForGlobalCOFF(raw_ostream &OS, const GlobalValue *GV,
 void emitLinkerFlagsForUsedCOFF(raw_ostream &OS, const GlobalValue *GV,
                                 const Triple &T, Mangler &M);
 
+/// Returns the ARM64EC mangled function name unless the input is already
+/// mangled.
 std::optional<std::string> getArm64ECMangledFunctionName(StringRef Name);
+
+/// Returns the ARM64EC demangled function name, unless the input is not
+/// mangled.
 std::optional<std::string> getArm64ECDemangledFunctionName(StringRef Name);
 
 } // End llvm namespace
diff --git a/llvm/lib/IR/Mangler.cpp b/llvm/lib/IR/Mangler.cpp
index e6c3ea9d568831..0dacf76743abdd 100644
--- a/llvm/lib/IR/Mangler.cpp
+++ b/llvm/lib/IR/Mangler.cpp
@@ -291,39 +291,41 @@ void llvm::emitLinkerFlagsForUsedCOFF(raw_ostream &OS, const GlobalValue *GV,
 }
 
 std::optional<std::string> llvm::getArm64ECMangledFunctionName(StringRef Name) {
-  bool IsCppFn = Name[0] == '?';
-  if (IsCppFn && Name.contains("$$h"))
-    return std::nullopt;
-  if (!IsCppFn && Name[0] == '#')
-    return std::nullopt;
+  if (Name[0] != '?') {
+    // For non-C++ symbols, prefix the name with "#" unless it's already
+    // mangled.
+    if (Name[0] == '#')
+      return std::nullopt;
+    return std::optional<std::string>(("#" + Name).str());
+  }
 
-  StringRef Prefix = "$$h";
+  // Insert the ARM64EC "$$h" tag after the mangled function name.
+  if (Name.contains("$$h"))
+    return std::nullopt;
   size_t InsertIdx = 0;
-  if (IsCppFn) {
-    InsertIdx = Name.find("@@");
-    size_t ThreeAtSignsIdx = Name.find("@@@");
-    if (InsertIdx != std::string::npos && InsertIdx != ThreeAtSignsIdx) {
-      InsertIdx += 2;
-    } else {
-      InsertIdx = Name.find("@");
-      if (InsertIdx != std::string::npos)
-        InsertIdx++;
-    }
+  InsertIdx = Name.find("@@");
+  size_t ThreeAtSignsIdx = Name.find("@@@");
+  if (InsertIdx != std::string::npos && InsertIdx != ThreeAtSignsIdx) {
+    InsertIdx += 2;
   } else {
-    Prefix = "#";
+    InsertIdx = Name.find("@");
+    if (InsertIdx != std::string::npos)
+      InsertIdx++;
   }
 
   return std::optional<std::string>(
-      (Name.substr(0, InsertIdx) + Prefix + Name.substr(InsertIdx)).str());
+      (Name.substr(0, InsertIdx) + "$$h" + Name.substr(InsertIdx)).str());
 }
 
 std::optional<std::string>
 llvm::getArm64ECDemangledFunctionName(StringRef Name) {
+  // For non-C++ names, drop the "#" prefix.
   if (Name[0] == '#')
     return std::optional<std::string>(Name.substr(1));
   if (Name[0] != '?')
     return std::nullopt;
 
+  // Drop the ARM64EC "$$h" tag.
   std::pair<StringRef, StringRef> Pair = Name.split("$$h");
   if (Pair.second.empty())
     return std::nullopt;

[mstorsjo]

LGTM, just one minor nit.

Thanks!

[mstorsjo]

You can probably merge this with the declaration of InsertIdx on the previous line now?