1 files changed, 16 insertions, 20 deletions
diff --git a/llvm/lib/IR/Mangler.cpp b/llvm/lib/IR/Mangler.cpp
index e6c3ea9d5688..884739b3212c 100644
--- a/llvm/lib/IR/Mangler.cpp
+++ b/llvm/lib/IR/Mangler.cpp
@@ -14,6 +14,7 @@
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/Twine.h"
+#include "llvm/Demangle/Demangle.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Function.h"
@@ -291,30 +292,25 @@ 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] == '#')
+  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());
+  }
+
+  // If the name contains $$h, then it is already mangled.
+  if (Name.contains("$$h"))
     return std::nullopt;
 
-  StringRef Prefix = "$$h";
-  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++;
-    }
-  } else {
-    Prefix = "#";
-  }
+  // Ask the demangler where we should insert "$$h".
+  auto InsertIdx = getArm64ECInsertionPointInMangledName(Name);
+  if (!InsertIdx)
+    return std::nullopt;
 
   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>