[LLD] [COFF] Allow wrapping dllimported functions

[mstorsjo]

Summary

GNU ld doesn't seem to do this though, but it looks like a reasonable use case, is easy to implement, and was requested in https://bugs.llvm.org/show_bug.cgi?id=47384.

[rnk]

lgtm

lld/COFF/MinGW.cpp

 // 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(ArrayRef<WrappedSymbol> wrapped) {
   DenseMap<Symbol *, Symbol *> map;
   for (const WrappedSymbol &w : wrapped) {
     map[w.sym] = w.wrap;
     map[w.real] = w.sym;
+    if (Defined *d = dyn_cast<Defined>(w.wrap)) {
+      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>(
+            saver.save("__imp_" + w.wrap->getName()), d);
+        symtab->localImportChunks.push_back(wrapimp->getChunk());
+        map[imp] = wrapimp;
+      }
+    }
   }
 
   // Update pointers in input files.
   parallelForEach(ObjFile::instances, [&](ObjFile *file) {
     MutableArrayRef<Symbol *> syms = file->getMutableSymbols();
     for (size_t i = 0, e = syms.size(); i != e; ++i)
       if (Symbol *s = map.lookup(syms[i]))
         syms[i] = s;
   });
 }

lld/test/COFF/wrap-dllimport.s

+// REQUIRES: x86
+
+// Check that we can wrap a dllimported symbol, so that references to
+// __imp_<symbol> gets redirected to a defined local import instead.
+
+// RUN: split-file %s %t.dir
+// RUN: llvm-mc -filetype=obj -triple=i686-win32-gnu %t.dir/main.s -o %t.main.obj
+// RUN: llvm-mc -filetype=obj -triple=i686-win32-gnu %t.dir/other.s -o %t.other.obj
+
+// RUN: lld-link -dll -out:%t.dll %t.other.obj -noentry -safeseh:no -export:foo -implib:%t.lib
+// RUN: lld-link -out:%t.exe %t.main.obj %t.lib -entry:entry -subsystem:console -debug:symtab -safeseh:no -wrap:foo -lldmap:%t.map
+// RUN: llvm-objdump -s -d --print-imm-hex %t.exe | FileCheck %s
+
+// CHECK:      Contents of section .rdata:
+// CHECK-NEXT:  402000 06104000
+
+// CHECK:      Disassembly of section .text:
+// CHECK-EMPTY:
+// CHECK:      00401000 <_entry>:
+// CHECK-NEXT:   401000: ff 25 00 20 40 00             jmpl    *0x402000
+// CHECK-EMPTY:
+// CHECK-NEXT: 00401006 <___wrap_foo>:
+// CHECK-NEXT:   401006: c3                            retl
+
+// The jmpl instruction in _entry points at an address in 0x402000,
+// which is the first 4 bytes of the .rdata section (above), which is a
+// pointer that points at ___wrap_foo.
+
+#--- main.s
+.global _entry
+_entry:
+  jmpl *__imp__foo
+
+.global ___wrap_foo
+___wrap_foo:
+  ret
+
+#--- other.s
+.global _foo
+
+_foo:
+  ret

lld/test/COFF/wrap-import.ll

 // REQUIRES: x86
 
 // Check that wrapping works when the wrapped symbol is imported from a
-// different DLL.
+// different DLL, redirecting references that used to point at the import
+// thunk, towards the local wrap function instead.
 
 // RUN: split-file %s %t.dir
 // RUN: llc %t.dir/main.ll -o %t.main.obj --filetype=obj
 // RUN: llvm-as %t.dir/main.ll -o %t.main.bc
 // RUN: llvm-mc -filetype=obj -triple=x86_64-win32-gnu %t.dir/lib.s -o %t.lib.obj
 
 // RUN: lld-link -dll -out:%t.lib.dll %t.lib.obj -noentry -export:func -implib:%t.lib.lib
 // RUN: lld-link -out:%t.exe %t.main.obj %t.lib.lib -entry:entry -subsystem:console -wrap:func