[llvm-lib][Object] Add support for EC importlib symbols. (#81059)

ARM64EC import libraries expose two additional symbols: mangled thunk
symbol (like `#func`) and auxiliary import symbol (like`__imp_aux_func`).
The main functional change with this patch is that those symbols are
properly added to static library ECSYMBOLS.

7 files changed, 225 insertions, 32 deletions

diff --git a/llvm/include/llvm/Object/COFF.h b/llvm/include/llvm/Object/COFF.h
index a548b2c..2a5c3d8 100644
--- a/llvm/include/llvm/Object/COFF.h
+++ b/llvm/include/llvm/Object/COFF.h
@@ -1362,6 +1362,47 @@ public:
   SectionStrippedError() { setErrorCode(object_error::section_stripped); }
 };
 
+inline std::optional<std::string>
+getArm64ECMangledFunctionName(StringRef Name) {
+  bool IsCppFn = Name[0] == '?';
+  if (IsCppFn && Name.find("$$h") != std::string::npos)
+    return std::nullopt;
+  if (!IsCppFn && Name[0] == '#')
+    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 = "#";
+  }
+
+  return std::optional<std::string>(
+      (Name.substr(0, InsertIdx) + Prefix + Name.substr(InsertIdx)).str());
+}
+
+inline std::optional<std::string>
+getArm64ECDemangledFunctionName(StringRef Name) {
+  if (Name[0] == '#')
+    return std::string(Name.substr(1));
+  if (Name[0] != '?')
+    return std::nullopt;
+
+  std::pair<StringRef, StringRef> Pair = Name.split("$$h");
+  if (Pair.second.empty())
+    return std::nullopt;
+  return (Pair.first + Pair.second).str();
+}
+
 } // end namespace object
 
 } // end namespace llvm
diff --git a/llvm/include/llvm/Object/COFFImportFile.h b/llvm/include/llvm/Object/COFFImportFile.h
index 7c5846e9..46a982d 100644
--- a/llvm/include/llvm/Object/COFFImportFile.h
+++ b/llvm/include/llvm/Object/COFFImportFile.h
@@ -27,6 +27,9 @@ namespace llvm {
 namespace object {
 
 class COFFImportFile : public SymbolicFile {
+private:
+  enum SymbolIndex { ImpSymbol, ThunkSymbol, ECAuxSymbol, ECThunkSymbol };
+
 public:
   COFFImportFile(MemoryBufferRef Source)
       : SymbolicFile(ID_COFFImportFile, Source) {}
@@ -36,9 +39,23 @@ public:
   void moveSymbolNext(DataRefImpl &Symb) const override { ++Symb.p; }
 
   Error printSymbolName(raw_ostream &OS, DataRefImpl Symb) const override {
-    if (Symb.p == 0)
+    switch (Symb.p) {
+    case ImpSymbol:
       OS << "__imp_";
-    OS << StringRef(Data.getBufferStart() + sizeof(coff_import_header));
+      break;
+    case ECAuxSymbol:
+      OS << "__imp_aux_";
+      break;
+    }
+    const char *Name = Data.getBufferStart() + sizeof(coff_import_header);
+    if (Symb.p != ECThunkSymbol && COFF::isArm64EC(getMachine())) {
+      if (std::optional<std::string> DemangledName =
+              getArm64ECDemangledFunctionName(Name)) {
+        OS << StringRef(*DemangledName);
+        return Error::success();
+      }
+    }
+    OS << StringRef(Name);
     return Error::success();
   }
 
@@ -52,7 +69,12 @@ public:
 
   basic_symbol_iterator symbol_end() const override {
     DataRefImpl Symb;
-    Symb.p = isData() ? 1 : 2;
+    if (isData())
+      Symb.p = ImpSymbol + 1;
+    else if (COFF::isArm64EC(getMachine()))
+      Symb.p = ECThunkSymbol + 1;
+    else
+      Symb.p = ThunkSymbol + 1;
     return BasicSymbolRef(Symb, this);
   }
 
diff --git a/llvm/lib/Object/COFFImportFile.cpp b/llvm/lib/Object/COFFImportFile.cpp
index 51e6274..a3e5e78 100644
--- a/llvm/lib/Object/COFFImportFile.cpp
+++ b/llvm/lib/Object/COFFImportFile.cpp
@@ -684,6 +684,21 @@ Error writeImportLibrary(StringRef ImportName, StringRef Path,
       NameType = getNameType(SymbolName, E.Name, Machine, MinGW);
     }
 
+    // On ARM64EC, use EXPORTAS to import demangled name for mangled symbols.
+    if (ImportType == IMPORT_CODE && isArm64EC(Machine)) {
+      if (std::optional<std::string> MangledName =
+              getArm64ECMangledFunctionName(Name)) {
+        if (ExportName.empty()) {
+          NameType = IMPORT_NAME_EXPORTAS;
+          ExportName.swap(Name);
+        }
+        Name = std::move(*MangledName);
+      } else if (ExportName.empty()) {
+        NameType = IMPORT_NAME_EXPORTAS;
+        ExportName = std::move(*getArm64ECDemangledFunctionName(Name));
+      }
+    }
+
     Members.push_back(OF.createShortImport(Name, E.Ordinal, ImportType,
                                            NameType, ExportName, Machine));
   }
diff --git a/llvm/lib/Target/AArch64/AArch64Arm64ECCallLowering.cpp b/llvm/lib/Target/AArch64/AArch64Arm64ECCallLowering.cpp
index 03d641d..55c5bbc 100644
--- a/llvm/lib/Target/AArch64/AArch64Arm64ECCallLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64Arm64ECCallLowering.cpp
@@ -24,11 +24,13 @@
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/InitializePasses.h"
+#include "llvm/Object/COFF.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/TargetParser/Triple.h"
 
 using namespace llvm;
+using namespace llvm::object;
 
 using OperandBundleDef = OperandBundleDefT<Value *>;
 
diff --git a/llvm/lib/Target/AArch64/AArch64MCInstLower.cpp b/llvm/lib/Target/AArch64/AArch64MCInstLower.cpp
index 1e12cf5..37d621c 100644
--- a/llvm/lib/Target/AArch64/AArch64MCInstLower.cpp
+++ b/llvm/lib/Target/AArch64/AArch64MCInstLower.cpp
@@ -23,11 +23,13 @@
 #include "llvm/MC/MCExpr.h"
 #include "llvm/MC/MCInst.h"
 #include "llvm/MC/MCStreamer.h"
+#include "llvm/Object/COFF.h"
 #include "llvm/Support/CodeGen.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Target/TargetLoweringObjectFile.h"
 #include "llvm/Target/TargetMachine.h"
 using namespace llvm;
+using namespace llvm::object;
 
 extern cl::opt<bool> EnableAArch64ELFLocalDynamicTLSGeneration;
 
diff --git a/llvm/lib/Target/AArch64/Utils/AArch64BaseInfo.h b/llvm/lib/Target/AArch64/Utils/AArch64BaseInfo.h
index 10e6965..8b32d59 100644
--- a/llvm/lib/Target/AArch64/Utils/AArch64BaseInfo.h
+++ b/llvm/lib/Target/AArch64/Utils/AArch64BaseInfo.h
@@ -248,34 +248,6 @@ static inline bool atomicBarrierDroppedOnZero(unsigned Opcode) {
   return false;
 }
 
-static inline std::optional<std::string>
-getArm64ECMangledFunctionName(std::string Name) {
-  bool IsCppFn = Name[0] == '?';
-  if (IsCppFn && Name.find("$$h") != std::string::npos)
-    return std::nullopt;
-  if (!IsCppFn && Name[0] == '#')
-    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 = "#";
-  }
-
-  Name.insert(Name.begin() + InsertIdx, Prefix.begin(), Prefix.end());
-  return std::optional<std::string>(Name);
-}
-
 namespace AArch64CC {
 
 // The CondCodes constants map directly to the 4-bit encoding of the condition
diff --git a/llvm/test/tools/llvm-lib/arm64ec-implib.test b/llvm/test/tools/llvm-lib/arm64ec-implib.test
index 4250c77..c583ef7 100644
--- a/llvm/test/tools/llvm-lib/arm64ec-implib.test
+++ b/llvm/test/tools/llvm-lib/arm64ec-implib.test
@@ -11,9 +11,23 @@ ARMAP-NEXT: __NULL_IMPORT_DESCRIPTOR in test.dll
 ARMAP-NEXT: test_NULL_THUNK_DATA in test.dll
 ARMAP-EMPTY:
 ARMAP-NEXT: Archive EC map
+ARMAP-NEXT: #expname in test.dll
+ARMAP-NEXT: #funcexp in test.dll
+ARMAP-NEXT: #mangledfunc in test.dll
+ARMAP-NEXT: ?test_cpp_func@@$$hYAHPEAX@Z in test.dll
+ARMAP-NEXT: ?test_cpp_func@@YAHPEAX@Z in test.dll
+ARMAP-NEXT: __imp_?test_cpp_func@@YAHPEAX@Z in test.dll
+ARMAP-NEXT: __imp_aux_?test_cpp_func@@YAHPEAX@Z in test.dll
+ARMAP-NEXT: __imp_aux_expname in test.dll
+ARMAP-NEXT: __imp_aux_funcexp in test.dll
+ARMAP-NEXT: __imp_aux_mangledfunc in test.dll
 ARMAP-NEXT: __imp_dataexp in test.dll
+ARMAP-NEXT: __imp_expname in test.dll
 ARMAP-NEXT: __imp_funcexp in test.dll
+ARMAP-NEXT: __imp_mangledfunc in test.dll
+ARMAP-NEXT: expname in test.dll
 ARMAP-NEXT: funcexp in test.dll
+ARMAP-NEXT: mangledfunc in test.dll
 
 RUN: llvm-readobj test.lib | FileCheck -check-prefix=READOBJ %s
 
@@ -35,10 +49,42 @@ READOBJ-EMPTY:
 READOBJ-NEXT: File: test.dll
 READOBJ-NEXT: Format: COFF-import-file-ARM64EC
 READOBJ-NEXT: Type: code
-READOBJ-NEXT: Name type: name
+READOBJ-NEXT: Name type: export as
 READOBJ-NEXT: Export name: funcexp
 READOBJ-NEXT: Symbol: __imp_funcexp
 READOBJ-NEXT: Symbol: funcexp
+READOBJ-NEXT: Symbol: __imp_aux_funcexp
+READOBJ-NEXT: Symbol: #funcexp
+READOBJ-EMPTY:
+READOBJ-NEXT: File: test.dll
+READOBJ-NEXT: Format: COFF-import-file-ARM64EC
+READOBJ-NEXT: Type: code
+READOBJ-NEXT: Name type: export as
+READOBJ-NEXT: Export name: mangledfunc
+READOBJ-NEXT: Symbol: __imp_mangledfunc
+READOBJ-NEXT: Symbol: mangledfunc
+READOBJ-NEXT: Symbol: __imp_aux_mangledfunc
+READOBJ-NEXT: Symbol: #mangledfunc
+READOBJ-EMPTY:
+READOBJ-NEXT: File: test.dll
+READOBJ-NEXT: Format: COFF-import-file-ARM64EC
+READOBJ-NEXT: Type: code
+READOBJ-NEXT: Name type: export as
+READOBJ-NEXT: Export name: ?test_cpp_func@@YAHPEAX@Z
+READOBJ-NEXT: Symbol: __imp_?test_cpp_func@@YAHPEAX@Z
+READOBJ-NEXT: Symbol: ?test_cpp_func@@YAHPEAX@Z
+READOBJ-NEXT: Symbol: __imp_aux_?test_cpp_func@@YAHPEAX@Z
+READOBJ-NEXT: Symbol: ?test_cpp_func@@$$hYAHPEAX@Z
+READOBJ-EMPTY:
+READOBJ-NEXT: File: test.dll
+READOBJ-NEXT: Format: COFF-import-file-ARM64EC
+READOBJ-NEXT: Type: code
+READOBJ-NEXT: Name type: export as
+READOBJ-NEXT: Export name: expname
+READOBJ-NEXT: Symbol: __imp_expname
+READOBJ-NEXT: Symbol: expname
+READOBJ-NEXT: Symbol: __imp_aux_expname
+READOBJ-NEXT: Symbol: #expname
 READOBJ-EMPTY:
 READOBJ-NEXT: File: test.dll
 READOBJ-NEXT: Format: COFF-import-file-ARM64EC
@@ -51,8 +97,101 @@ Creating a new lib containing the existing lib:
 RUN: llvm-lib -machine:arm64ec test.lib -out:test2.lib
 RUN: llvm-nm --print-armap test2.lib | FileCheck -check-prefix=ARMAP %s
 
+
+RUN: llvm-lib -machine:arm64ec -def:exportas.def -out:exportas.lib
+RUN: llvm-nm --print-armap exportas.lib | FileCheck -check-prefix=EXPAS-ARMAP %s
+RUN: llvm-readobj exportas.lib | FileCheck -check-prefix=EXPAS-READOBJ %s
+
+EXPAS-ARMAP:      Archive EC map
+EXPAS-ARMAP-NEXT: #func1 in test.dll
+EXPAS-ARMAP-NEXT: #func2 in test.dll
+EXPAS-ARMAP-NEXT: #func3 in test.dll
+EXPAS-ARMAP-NEXT: #func4 in test.dll
+EXPAS-ARMAP-NEXT: __imp_aux_func1 in test.dll
+EXPAS-ARMAP-NEXT: __imp_aux_func2 in test.dll
+EXPAS-ARMAP-NEXT: __imp_aux_func3 in test.dll
+EXPAS-ARMAP-NEXT: __imp_aux_func4 in test.dll
+EXPAS-ARMAP-NEXT: __imp_data1 in test.dll
+EXPAS-ARMAP-NEXT: __imp_data2 in test.dll
+EXPAS-ARMAP-NEXT: __imp_func1 in test.dll
+EXPAS-ARMAP-NEXT: __imp_func2 in test.dll
+EXPAS-ARMAP-NEXT: __imp_func3 in test.dll
+EXPAS-ARMAP-NEXT: __imp_func4 in test.dll
+EXPAS-ARMAP-NEXT: func1 in test.dll
+EXPAS-ARMAP-NEXT: func2 in test.dll
+EXPAS-ARMAP-NEXT: func3 in test.dll
+EXPAS-ARMAP-NEXT: func4 in test.dll
+
+EXPAS-READOBJ:      File: test.dll
+EXPAS-READOBJ-NEXT: Format: COFF-import-file-ARM64EC
+EXPAS-READOBJ-NEXT: Type: code
+EXPAS-READOBJ-NEXT: Name type: export as
+EXPAS-READOBJ-NEXT: Export name: func1
+EXPAS-READOBJ-NEXT: Symbol: __imp_func1
+EXPAS-READOBJ-NEXT: Symbol: func1
+EXPAS-READOBJ-NEXT: Symbol: __imp_aux_func1
+EXPAS-READOBJ-NEXT: Symbol: #func1
+EXPAS-READOBJ-EMPTY:
+EXPAS-READOBJ-NEXT: File: test.dll
+EXPAS-READOBJ-NEXT: Format: COFF-import-file-ARM64EC
+EXPAS-READOBJ-NEXT: Type: code
+EXPAS-READOBJ-NEXT: Name type: export as
+EXPAS-READOBJ-NEXT: Export name: func2
+EXPAS-READOBJ-NEXT: Symbol: __imp_func2
+EXPAS-READOBJ-NEXT: Symbol: func2
+EXPAS-READOBJ-NEXT: Symbol: __imp_aux_func2
+EXPAS-READOBJ-NEXT: Symbol: #func2
+EXPAS-READOBJ-EMPTY:
+EXPAS-READOBJ-NEXT: File: test.dll
+EXPAS-READOBJ-NEXT: Format: COFF-import-file-ARM64EC
+EXPAS-READOBJ-NEXT: Type: code
+EXPAS-READOBJ-NEXT: Name type: export as
+EXPAS-READOBJ-NEXT: Export name: #func3
+EXPAS-READOBJ-NEXT: Symbol: __imp_func3
+EXPAS-READOBJ-NEXT: Symbol: func3
+EXPAS-READOBJ-NEXT: Symbol: __imp_aux_func3
+EXPAS-READOBJ-NEXT: Symbol: #func3
+EXPAS-READOBJ-EMPTY:
+EXPAS-READOBJ-NEXT: File: test.dll
+EXPAS-READOBJ-NEXT: Format: COFF-import-file-ARM64EC
+EXPAS-READOBJ-NEXT: Type: code
+EXPAS-READOBJ-NEXT: Name type: export as
+EXPAS-READOBJ-NEXT: Export name: #func4
+EXPAS-READOBJ-NEXT: Symbol: __imp_func4
+EXPAS-READOBJ-NEXT: Symbol: func4
+EXPAS-READOBJ-NEXT: Symbol: __imp_aux_func4
+EXPAS-READOBJ-NEXT: Symbol: #func4
+EXPAS-READOBJ-EMPTY:
+EXPAS-READOBJ-NEXT: File: test.dll
+EXPAS-READOBJ-NEXT: Format: COFF-import-file-ARM64EC
+EXPAS-READOBJ-NEXT: Type: data
+EXPAS-READOBJ-NEXT: Name type: export as
+EXPAS-READOBJ-NEXT: Export name: #data1
+EXPAS-READOBJ-NEXT: Symbol: __imp_data1
+EXPAS-READOBJ-EMPTY:
+EXPAS-READOBJ-NEXT: File: test.dll
+EXPAS-READOBJ-NEXT: Format: COFF-import-file-ARM64EC
+EXPAS-READOBJ-NEXT: Type: data
+EXPAS-READOBJ-NEXT: Name type: export as
+EXPAS-READOBJ-NEXT: Export name: data2
+EXPAS-READOBJ-NEXT: Symbol: __imp_data2
+
+
 #--- test.def
 LIBRARY test.dll
 EXPORTS
     funcexp
+    #mangledfunc
+    ?test_cpp_func@@YAHPEAX@Z
+    expname=impname
     dataexp DATA
+
+#--- exportas.def
+LIBRARY test.dll
+EXPORTS
+    #func1 EXPORTAS func1
+    func2 EXPORTAS func2
+    func3 EXPORTAS #func3
+    #func4 EXPORTAS #func4
+    data1 DATA EXPORTAS #data1
+    #data2 DATA EXPORTAS data2