[lld] [LLD][COFF] Fix handling of weak aliases referencing lazy symbols (PR #112243)

[Jacek Caban via llvm-commits]

https://github.com/cjacek created https://github.com/llvm/llvm-project/pull/112243

The assumption that a symbol is either `Defined` or `Undefined` is not always true for some cases. For example, `mangleMaybe` may create a weak alias to a lazy archive symbol.

>From 5544b21bd498d39e9aa8699c75e25c87069e753c Mon Sep 17 00:00:00 2001
From: Jacek Caban <jacek at codeweavers.com>
Date: Mon, 14 Oct 2024 17:11:00 +0200
Subject: [PATCH] [LLD][COFF] Fix handling of weak aliases referencing lazy
 symbols

The assumption that a symbol is either `Defined` or `Undefined` is not always true
for some cases. For example, `mangleMaybe` may create a weak alias to a lazy archive symbol.
---
 lld/COFF/Driver.cpp       |  4 ++--
 lld/COFF/Symbols.cpp      |  8 ++++----
 lld/COFF/Symbols.h        |  5 ++++-
 lld/test/COFF/weak-lazy.s | 18 ++++++++++++++++++
 4 files changed, 28 insertions(+), 7 deletions(-)
 create mode 100644 lld/test/COFF/weak-lazy.s

diff --git a/lld/COFF/Driver.cpp b/lld/COFF/Driver.cpp
index 85a58a3677181a..12e1ae62811270 100644
--- a/lld/COFF/Driver.cpp
+++ b/lld/COFF/Driver.cpp
@@ -1340,7 +1340,7 @@ void LinkerDriver::maybeCreateECExportThunk(StringRef name, Symbol *&sym) {
   if (!sym)
     return;
   if (auto undef = dyn_cast<Undefined>(sym))
-    def = undef->getWeakAlias();
+    def = undef->getDefinedWeakAlias();
   else
     def = dyn_cast<Defined>(sym);
   if (!def)
@@ -1376,7 +1376,7 @@ void LinkerDriver::createECExportThunks() {
       continue;
     Defined *targetSym;
     if (auto undef = dyn_cast<Undefined>(sym))
-      targetSym = undef->getWeakAlias();
+      targetSym = undef->getDefinedWeakAlias();
     else
       targetSym = dyn_cast<Defined>(sym);
     if (!targetSym)
diff --git a/lld/COFF/Symbols.cpp b/lld/COFF/Symbols.cpp
index 567c2b93776c94..89f2da02bdcf42 100644
--- a/lld/COFF/Symbols.cpp
+++ b/lld/COFF/Symbols.cpp
@@ -112,12 +112,12 @@ DefinedImportThunk::DefinedImportThunk(COFFLinkerContext &ctx, StringRef name,
                                        ImportThunkChunk *chunk)
     : Defined(DefinedImportThunkKind, name), wrappedSym(s), data(chunk) {}
 
-Defined *Undefined::getWeakAlias() {
+Symbol *Undefined::getWeakAlias() {
   // A weak alias may be a weak alias to another symbol, so check recursively.
   DenseSet<Symbol *> weakChain;
   for (Symbol *a = weakAlias; a; a = cast<Undefined>(a)->weakAlias) {
-    if (auto *d = dyn_cast<Defined>(a))
-      return d;
+    if (!isa<Undefined>(a))
+      return a;
     if (!weakChain.insert(a).second)
       break; // We have a cycle.
   }
@@ -125,7 +125,7 @@ Defined *Undefined::getWeakAlias() {
 }
 
 bool Undefined::resolveWeakAlias() {
-  Defined *d = getWeakAlias();
+  Defined *d = getDefinedWeakAlias();
   if (!d)
     return false;
 
diff --git a/lld/COFF/Symbols.h b/lld/COFF/Symbols.h
index 9b21e09bf83a42..a898ebf05fd809 100644
--- a/lld/COFF/Symbols.h
+++ b/lld/COFF/Symbols.h
@@ -340,7 +340,10 @@ class Undefined : public Symbol {
   // If this symbol is external weak, try to resolve it to a defined
   // symbol by searching the chain of fallback symbols. Returns the symbol if
   // successful, otherwise returns null.
-  Defined *getWeakAlias();
+  Symbol *getWeakAlias();
+  Defined *getDefinedWeakAlias() {
+    return dyn_cast_or_null<Defined>(getWeakAlias());
+  }
 
   // If this symbol is external weak, replace this object with aliased symbol.
   bool resolveWeakAlias();
diff --git a/lld/test/COFF/weak-lazy.s b/lld/test/COFF/weak-lazy.s
new file mode 100644
index 00000000000000..2812ba7af8b5f8
--- /dev/null
+++ b/lld/test/COFF/weak-lazy.s
@@ -0,0 +1,18 @@
+# REQUIRES: x86
+
+# RUN: llvm-mc -filetype=obj -triple=i686-windows %s -o %t.obj
+# RUN: llvm-lib -machine:x86 -out:%t-func.lib %t.obj
+
+# -export:func creates a weak alias to a lazy symbol. Make sure we can handle that when processing -export:func2=func.
+# RUN: lld-link -dll -noentry -machine:x86 -out:%t.dll %t-func.lib -export:func -export:func2=func
+
+        .text
+        .def    @feat.00;
+        .scl    3;
+        .type   0;
+        .endef
+        .globl  @feat.00
+.set @feat.00, 1
+        .globl _func at 0
+_func at 0:
+        retl