[LLD][MinGW] Add support for wrapped symbols on ARM64X (llvm#126296)

Apply `-wrap` arguments to both symbol tables.

Author: cjacek

lld/COFF/Driver.cpp

@@ -2573,11 +2573,13 @@ void LinkerDriver::linkerMain(ArrayRef<const char *> argsArr) {
     ctx.symtab.addUndefinedGlob(pat);
 
   // Create wrapped symbols for -wrap option.
-  std::vector<WrappedSymbol> wrapped = addWrappedSymbols(ctx, args);
-  // Load more object files that might be needed for wrapped symbols.
-  if (!wrapped.empty())
-    while (run())
-      ;
+  ctx.forEachSymtab([&](SymbolTable &symtab) {
+    addWrappedSymbols(symtab, args);
+    // Load more object files that might be needed for wrapped symbols.
+    if (!symtab.wrapped.empty())
+      while (run())
+        ;
+  });
 
   if (config->autoImport || config->stdcallFixup) {
     // MinGW specific.
@@ -2647,8 +2649,10 @@ void LinkerDriver::linkerMain(ArrayRef<const char *> argsArr) {
   run();
 
   // Apply symbol renames for -wrap.
-  if (!wrapped.empty())
-    wrapSymbols(ctx, wrapped);
+  ctx.forEachSymtab([](SymbolTable &symtab) {
+    if (!symtab.wrapped.empty())
+      wrapSymbols(symtab);
+  });
 
   if (isArm64EC(config->machine))
     createECExportThunks();

lld/COFF/MinGW.cpp

@@ -207,8 +207,8 @@ static StringRef mangle(Twine sym, MachineTypes machine) {
 // This function instantiates wrapper symbols. At this point, they seem
 // like they are not being used at all, so we explicitly set some flags so
 // that LTO won't eliminate them.
-std::vector<WrappedSymbol>
-lld::coff::addWrappedSymbols(COFFLinkerContext &ctx, opt::InputArgList &args) {
+void lld::coff::addWrappedSymbols(SymbolTable &symtab,
+                                  opt::InputArgList &args) {
   std::vector<WrappedSymbol> v;
   DenseSet<StringRef> seen;
 
@@ -217,14 +217,14 @@ lld::coff::addWrappedSymbols(COFFLinkerContext &ctx, opt::InputArgList &args) {
     if (!seen.insert(name).second)
       continue;
 
-    Symbol *sym = ctx.symtab.findUnderscore(name);
+    Symbol *sym = symtab.findUnderscore(name);
     if (!sym)
       continue;
 
     Symbol *real =
-        ctx.symtab.addUndefined(mangle("__real_" + name, ctx.config.machine));
+        symtab.addUndefined(mangle("__real_" + name, symtab.machine));
     Symbol *wrap =
-        ctx.symtab.addUndefined(mangle("__wrap_" + name, ctx.config.machine));
+        symtab.addUndefined(mangle("__wrap_" + name, symtab.machine));
     v.push_back({sym, real, wrap});
 
     // These symbols may seem undefined initially, but don't bail out
@@ -243,37 +243,36 @@ lld::coff::addWrappedSymbols(COFFLinkerContext &ctx, opt::InputArgList &args) {
     if (!isa<Undefined>(wrap))
       wrap->isUsedInRegularObj = true;
   }
-  return v;
+  symtab.wrapped = std::move(v);
 }
 
 // Do renaming for -wrap by updating pointers to symbols.
 //
 // When this function is executed, only InputFiles and symbol table
 // contain pointers to symbol objects. We visit them to replace pointers,
 // so that wrapped symbols are swapped as instructed by the command line.
-void lld::coff::wrapSymbols(COFFLinkerContext &ctx,
-                            ArrayRef<WrappedSymbol> wrapped) {
+void lld::coff::wrapSymbols(SymbolTable &symtab) {
   DenseMap<Symbol *, Symbol *> map;
-  for (const WrappedSymbol &w : wrapped) {
+  for (const WrappedSymbol &w : symtab.wrapped) {
     map[w.sym] = w.wrap;
     map[w.real] = w.sym;
     if (Defined *d = dyn_cast<Defined>(w.wrap)) {
-      Symbol *imp = ctx.symtab.find(("__imp_" + w.sym->getName()).str());
+      Symbol *imp = symtab.find(("__imp_" + w.sym->getName()).str());
       // Create a new defined local import for the wrap symbol. If
       // no imp prefixed symbol existed, there's no need for it.
       // (We can't easily distinguish whether any object file actually
       // referenced it or not, though.)
       if (imp) {
         DefinedLocalImport *wrapimp = make<DefinedLocalImport>(
-            ctx, saver().save("__imp_" + w.wrap->getName()), d);
-        ctx.symtab.localImportChunks.push_back(wrapimp->getChunk());
+            symtab.ctx, saver().save("__imp_" + w.wrap->getName()), d);
+        symtab.localImportChunks.push_back(wrapimp->getChunk());
         map[imp] = wrapimp;
       }
     }
   }
 
   // Update pointers in input files.
-  parallelForEach(ctx.objFileInstances, [&](ObjFile *file) {
+  parallelForEach(symtab.ctx.objFileInstances, [&](ObjFile *file) {
     MutableArrayRef<Symbol *> syms = file->getMutableSymbols();
     for (auto &sym : syms)
       if (Symbol *s = map.lookup(sym))

lld/COFF/MinGW.h

@@ -55,17 +55,9 @@ void writeDefFile(COFFLinkerContext &, StringRef name,
 // symbol becomes accessible as `__real_foo`, so you can call that from your
 // wrapper.
 //
-// This data structure is instantiated for each -wrap option.
-struct WrappedSymbol {
-  Symbol *sym;
-  Symbol *real;
-  Symbol *wrap;
-};
-
-std::vector<WrappedSymbol> addWrappedSymbols(COFFLinkerContext &ctx,
-                                             llvm::opt::InputArgList &args);
+void addWrappedSymbols(SymbolTable &symtab, llvm::opt::InputArgList &args);
 
-void wrapSymbols(COFFLinkerContext &ctx, ArrayRef<WrappedSymbol> wrapped);
+void wrapSymbols(SymbolTable &symtab);
 
 } // namespace lld::coff
 

lld/COFF/SymbolTable.h

@@ -34,6 +34,13 @@ class LazyArchive;
 class SectionChunk;
 class Symbol;
 
+// This data structure is instantiated for each -wrap option.
+struct WrappedSymbol {
+  Symbol *sym;
+  Symbol *real;
+  Symbol *wrap;
+};
+
 // SymbolTable is a bucket of all known symbols, including defined,
 // undefined, or lazy symbols (the last one is symbols in archive
 // files whose archive members are not yet loaded).
@@ -161,6 +168,9 @@ class SymbolTable {
   Symbol *delayLoadHelper = nullptr;
   Chunk *tailMergeUnwindInfoChunk = nullptr;
 
+  // A list of wrapped symbols.
+  std::vector<WrappedSymbol> wrapped;
+
   void fixupExports();
   void assignExportOrdinals();
   void parseModuleDefs(StringRef path);

lld/test/COFF/arm64x-wrap.s

@@ -0,0 +1,31 @@
+// REQUIRES: aarch64
+// RUN: split-file %s %t.dir && cd %t.dir
+
+// RUN: llvm-mc -filetype=obj -triple=arm64ec-windows test.s -o test-arm64ec.obj
+// RUN: llvm-mc -filetype=obj -triple=aarch64-windows test.s -o test-arm64.obj
+// RUN: llvm-mc -filetype=obj -triple=arm64ec-windows other.s -o other-arm64ec.obj
+// RUN: llvm-mc -filetype=obj -triple=aarch64-windows other.s -o other-arm64.obj
+// RUN: llvm-mc -filetype=obj -triple=arm64ec-windows %S/Inputs/loadconfig-arm64ec.s -o loadconfig-arm64ec.obj
+// RUN: llvm-mc -filetype=obj -triple=aarch64-windows %S/Inputs/loadconfig-arm64.s -o loadconfig-arm64.obj
+
+// RUN: lld-link -machine:arm64x -dll -noentry test-arm64.obj test-arm64ec.obj other-arm64.obj other-arm64ec.obj \
+// RUN:          loadconfig-arm64.obj loadconfig-arm64ec.obj -out:out.dll -wrap:sym -wrap:nosuchsym
+
+// RUN: llvm-readobj --hex-dump=.test out.dll | FileCheck %s
+// CHECK: 0x180004000 02000000 02000000 01000000 02000000
+// CHECK: 0x180004010 02000000 01000000
+
+#--- test.s
+        .section .test,"dr"
+        .word sym
+        .word __wrap_sym
+        .word __real_sym
+
+#--- other.s
+        .global sym
+        .global __wrap_sym
+        .global __real_sym
+
+sym = 1
+__wrap_sym = 2
+__real_sym = 3