From 0fad0b99b214cbd0a4e9fa1831ad9f4b96bdf059 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Martin=20Storsj=C3=B6?= <martin@martin.st>
Date: Tue, 16 Apr 2024 15:22:37 +0300
Subject: [PATCH] [llvm-dlltool] Handle import renaming using other name types,
 when possible

This avoids needing to use weak aliases for these cases. (Weak
aliases only work if there's another, regular import entry that
provide the desired symbol from the DLL.)
---
 llvm/lib/Object/COFFImportFile.cpp            | 66 +++++++++++++------
 .../tools/llvm-dlltool/coff-decorated.def     | 16 +++++
 2 files changed, 61 insertions(+), 21 deletions(-)

diff --git a/llvm/lib/Object/COFFImportFile.cpp b/llvm/lib/Object/COFFImportFile.cpp
index 03af4921ddbca..1ddc5d954de68 100644
--- a/llvm/lib/Object/COFFImportFile.cpp
+++ b/llvm/lib/Object/COFFImportFile.cpp
@@ -52,18 +52,12 @@ StringRef COFFImportFile::getFileFormatName() const {
   }
 }
 
-StringRef COFFImportFile::getExportName() const {
-  const coff_import_header *hdr = getCOFFImportHeader();
-  StringRef name = Data.getBuffer().substr(sizeof(*hdr)).split('\0').first;
-
+static StringRef applyNameType(ImportNameType Type, StringRef name) {
   auto ltrim1 = [](StringRef s, StringRef chars) {
     return !s.empty() && chars.contains(s[0]) ? s.substr(1) : s;
   };
 
-  switch (hdr->getNameType()) {
-  case IMPORT_ORDINAL:
-    name = "";
-    break;
+  switch (Type) {
   case IMPORT_NAME_NOPREFIX:
     name = ltrim1(name, "?@_");
     break;
@@ -71,6 +65,24 @@ StringRef COFFImportFile::getExportName() const {
     name = ltrim1(name, "?@_");
     name = name.substr(0, name.find('@'));
     break;
+  default:
+    break;
+  }
+  return name;
+}
+
+StringRef COFFImportFile::getExportName() const {
+  const coff_import_header *hdr = getCOFFImportHeader();
+  StringRef name = Data.getBuffer().substr(sizeof(*hdr)).split('\0').first;
+
+  switch (hdr->getNameType()) {
+  case IMPORT_ORDINAL:
+    name = "";
+    break;
+  case IMPORT_NAME_NOPREFIX:
+  case IMPORT_NAME_UNDECORATE:
+    name = applyNameType(static_cast<ImportNameType>(hdr->getNameType()), name);
+    break;
   case IMPORT_NAME_EXPORTAS: {
     // Skip DLL name
     name = Data.getBuffer().substr(sizeof(*hdr) + name.size() + 1);
@@ -691,19 +703,6 @@ Error writeImportLibrary(StringRef ImportName, StringRef Path,
         Name.swap(*ReplacedName);
       }
 
-      if (!E.ImportName.empty() && Name != E.ImportName) {
-        StringRef Prefix = "";
-        if (Machine == IMAGE_FILE_MACHINE_I386 && AddUnderscores)
-          Prefix = "_";
-
-        if (ImportType == IMPORT_CODE)
-          Members.push_back(OF.createWeakExternal((Prefix + E.ImportName).str(),
-                                                  Name, false, M));
-        Members.push_back(OF.createWeakExternal((Prefix + E.ImportName).str(),
-                                                Name, true, M));
-        continue;
-      }
-
       ImportNameType NameType;
       std::string ExportName;
       if (E.Noname) {
@@ -711,6 +710,31 @@ Error writeImportLibrary(StringRef ImportName, StringRef Path,
       } else if (!E.ExportAs.empty()) {
         NameType = IMPORT_NAME_EXPORTAS;
         ExportName = E.ExportAs;
+      } else if (!E.ImportName.empty()) {
+        // If we need to import from a specific ImportName, we may need to use
+        // a weak alias (which needs another import to point at). But if we can
+        // express ImportName based on the symbol name and a specific NameType,
+        // prefer that over an alias.
+        if (Machine == IMAGE_FILE_MACHINE_I386 &&
+            applyNameType(IMPORT_NAME_UNDECORATE, Name) == E.ImportName)
+          NameType = IMPORT_NAME_UNDECORATE;
+        else if (Machine == IMAGE_FILE_MACHINE_I386 &&
+                 applyNameType(IMPORT_NAME_NOPREFIX, Name) == E.ImportName)
+          NameType = IMPORT_NAME_NOPREFIX;
+        else if (Name == E.ImportName)
+          NameType = IMPORT_NAME;
+        else {
+          StringRef Prefix = "";
+          if (Machine == IMAGE_FILE_MACHINE_I386 && AddUnderscores)
+            Prefix = "_";
+
+          if (ImportType == IMPORT_CODE)
+            Members.push_back(OF.createWeakExternal(
+                (Prefix + E.ImportName).str(), Name, false, M));
+          Members.push_back(OF.createWeakExternal((Prefix + E.ImportName).str(),
+                                                  Name, true, M));
+          continue;
+        }
       } else {
         NameType = getNameType(SymbolName, E.Name, M, MinGW);
       }
diff --git a/llvm/test/tools/llvm-dlltool/coff-decorated.def b/llvm/test/tools/llvm-dlltool/coff-decorated.def
index fc81f23d09d6c..773e3762cc3d7 100644
--- a/llvm/test/tools/llvm-dlltool/coff-decorated.def
+++ b/llvm/test/tools/llvm-dlltool/coff-decorated.def
@@ -13,6 +13,10 @@ StdcallExportName@4=StdcallInternalFunction@4
 OtherStdcallExportName@4=CdeclInternalFunction
 CdeclExportName=StdcallInternalFunction@4
 
+NoprefixStdcall@4 == NoprefixStdcall@4
+DecoratedStdcall@4 == _DecoratedStdcall@4
+UndecoratedStdcall@4 == UndecoratedStdcall
+
 ; CHECK:      Name type: noprefix
 ; CHECK-NEXT: Export name: CdeclFunction
 ; CHECK-NEXT: Symbol: __imp__CdeclFunction
@@ -43,3 +47,15 @@ CdeclExportName=StdcallInternalFunction@4
 ; CHECK-NEXT: Export name: CdeclExportName
 ; CHECK-NEXT: Symbol: __imp__CdeclExportName
 ; CHECK-NEXT: Symbol: _CdeclExportName
+; CHECK:      Name type: noprefix
+; CHECK-NEXT: Export name: NoprefixStdcall@4
+; CHECK-NEXT: Symbol: __imp__NoprefixStdcall@4
+; CHECK-NEXT: Symbol: _NoprefixStdcall@4
+; CHECK:      Name type: name
+; CHECK-NEXT: Export name: _DecoratedStdcall@4
+; CHECK-NEXT: Symbol: __imp__DecoratedStdcall@4
+; CHECK-NEXT: Symbol: _DecoratedStdcall@4
+; CHECK:      Name type: undecorate
+; CHECK-NEXT: Export name: UndecoratedStdcall
+; CHECK-NEXT: Symbol: __imp__UndecoratedStdcall@4
+; CHECK-NEXT: Symbol: _UndecoratedStdcall@4