From 446965070149d53374ee4d3010789673083fb11c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Martin=20Storsj=C3=B6?= <martin@martin.st>
Date: Tue, 16 Jul 2024 23:17:01 +0300
Subject: [PATCH] [llvm-dlltool] Handle import renaming using other name types,
 when possible (#98228)

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 03af4921ddbca3..1ddc5d954de685 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 fc81f23d09d6c4..773e3762cc3d73 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