Reapply [LLD] [COFF] Implement a GNU/ELF like -wrap option

Add a simple forwarding option in the MinGW frontend, and implement
the private -wrap option in the COFF linker.

The feature in lld-link isn't gated by the -lldmingw option, but
the option is left as a private, undocumented option primarily
used by the MinGW driver.

The implementation is significantly based on the support for --wrap
in the ELF linker, but many small nuance details are different
between the ELF and COFF linkers, ending up with more than a few
implementation differences.

This fixes https://bugs.llvm.org/show_bug.cgi?id=47384.

Differential Revision: https://reviews.llvm.org/D89004

Reapplied with the bitfield member canInline fixed so it doesn't break
builds targeting windows.

Author: mstorsjo

lld/COFF/Driver.cpp

@@ -2008,6 +2008,12 @@ void LinkerDriver::link(ArrayRef<const char *> argsArr) {
     while (run());
   }
 
+  // Create wrapped symbols for -wrap option.
+  std::vector<WrappedSymbol> wrapped = addWrappedSymbols(args);
+  // Load more object files that might be needed for wrapped symbols.
+  if (!wrapped.empty())
+    while (run());
+
   if (config->autoImport) {
     // MinGW specific.
     // Load any further object files that might be needed for doing automatic
@@ -2051,6 +2057,10 @@ void LinkerDriver::link(ArrayRef<const char *> argsArr) {
   // references to the symbols we use from them.
   run();
 
+  // Apply symbol renames for -wrap.
+  if (!wrapped.empty())
+    wrapSymbols(wrapped);
+
   // Resolve remaining undefined symbols and warn about imported locals.
   symtab->resolveRemainingUndefines();
   if (errorCount())

lld/COFF/InputFiles.h

@@ -148,6 +148,8 @@ class ObjFile : public InputFile {
   ArrayRef<SectionChunk *> getGuardLJmpChunks() { return guardLJmpChunks; }
   ArrayRef<Symbol *> getSymbols() { return symbols; }
 
+  MutableArrayRef<Symbol *> getMutableSymbols() { return symbols; }
+
   ArrayRef<uint8_t> getDebugSection(StringRef secName);
 
   // Returns a Symbol object for the symbolIndex'th symbol in the

lld/COFF/LTO.cpp

@@ -139,6 +139,11 @@ void BitcodeCompiler::add(BitcodeFile &f) {
     r.VisibleToRegularObj = sym->isUsedInRegularObj;
     if (r.Prevailing)
       undefine(sym);
+
+    // We tell LTO to not apply interprocedural optimization for wrapped
+    // (with -wrap) symbols because otherwise LTO would inline them while
+    // their values are still not final.
+    r.LinkerRedefined = !sym->canInline;
   }
   checkError(ltoObj->add(std::move(f.obj), resols));
 }

lld/COFF/MinGW.cpp

@@ -7,9 +7,14 @@
 //===----------------------------------------------------------------------===//
 
 #include "MinGW.h"
+#include "Driver.h"
+#include "InputFiles.h"
 #include "SymbolTable.h"
 #include "lld/Common/ErrorHandler.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/Object/COFF.h"
+#include "llvm/Support/Parallel.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/raw_ostream.h"
 
@@ -173,3 +178,73 @@ void lld::coff::writeDefFile(StringRef name) {
     os << "\n";
   }
 }
+
+static StringRef mangle(Twine sym) {
+  assert(config->machine != IMAGE_FILE_MACHINE_UNKNOWN);
+  if (config->machine == I386)
+    return saver.save("_" + sym);
+  return saver.save(sym);
+}
+
+// Handles -wrap option.
+//
+// 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(opt::InputArgList &args) {
+  std::vector<WrappedSymbol> v;
+  DenseSet<StringRef> seen;
+
+  for (auto *arg : args.filtered(OPT_wrap)) {
+    StringRef name = arg->getValue();
+    if (!seen.insert(name).second)
+      continue;
+
+    Symbol *sym = symtab->findUnderscore(name);
+    if (!sym)
+      continue;
+
+    Symbol *real = symtab->addUndefined(mangle("__real_" + name));
+    Symbol *wrap = symtab->addUndefined(mangle("__wrap_" + name));
+    v.push_back({sym, real, wrap});
+
+    // These symbols may seem undefined initially, but don't bail out
+    // at symtab->reportUnresolvable() due to them, but let wrapSymbols
+    // below sort things out before checking finally with
+    // symtab->resolveRemainingUndefines().
+    sym->deferUndefined = true;
+    real->deferUndefined = true;
+    // We want to tell LTO not to inline symbols to be overwritten
+    // because LTO doesn't know the final symbol contents after renaming.
+    real->canInline = false;
+    sym->canInline = false;
+
+    // Tell LTO not to eliminate these symbols.
+    sym->isUsedInRegularObj = true;
+    if (!isa<Undefined>(wrap))
+      wrap->isUsedInRegularObj = true;
+  }
+  return 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(ArrayRef<WrappedSymbol> wrapped) {
+  DenseMap<Symbol *, Symbol *> map;
+  for (const WrappedSymbol &w : wrapped) {
+    map[w.sym] = w.wrap;
+    map[w.real] = w.sym;
+  }
+
+  // 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/COFF/MinGW.h

@@ -12,7 +12,10 @@
 #include "Config.h"
 #include "Symbols.h"
 #include "lld/Common/LLVM.h"
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/StringSet.h"
+#include "llvm/Option/ArgList.h"
+#include <vector>
 
 namespace lld {
 namespace coff {
@@ -36,6 +39,24 @@ class AutoExporter {
 
 void writeDefFile(StringRef name);
 
+// The -wrap option is a feature to rename symbols so that you can write
+// wrappers for existing functions. If you pass `-wrap:foo`, all
+// occurrences of symbol `foo` are resolved to `__wrap_foo` (so, you are
+// expected to write `__wrap_foo` function as a wrapper). The original
+// 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(llvm::opt::InputArgList &args);
+
+void wrapSymbols(ArrayRef<WrappedSymbol> wrapped);
+
 } // namespace coff
 } // namespace lld
 

lld/COFF/Options.td

@@ -252,6 +252,7 @@ def print_symbol_order: P<
     "print-symbol-order",
     "Print a symbol order specified by /call-graph-ordering-file and "
     "/call-graph-profile-sort into the specified file">;
+def wrap : P_priv<"wrap">;
 
 // Flags for debugging
 def lldmap : F<"lldmap">;

lld/COFF/SymbolTable.cpp

@@ -390,7 +390,7 @@ void SymbolTable::reportUnresolvable() {
   for (auto &i : symMap) {
     Symbol *sym = i.second;
     auto *undef = dyn_cast<Undefined>(sym);
-    if (!undef)
+    if (!undef || sym->deferUndefined)
       continue;
     if (undef->getWeakAlias())
       continue;
@@ -482,6 +482,7 @@ std::pair<Symbol *, bool> SymbolTable::insert(StringRef name) {
     sym = reinterpret_cast<Symbol *>(make<SymbolUnion>());
     sym->isUsedInRegularObj = false;
     sym->pendingArchiveLoad = false;
+    sym->canInline = true;
     inserted = true;
   }
   return {sym, inserted};

lld/COFF/Symbols.h

@@ -103,8 +103,8 @@ class Symbol {
   explicit Symbol(Kind k, StringRef n = "")
       : symbolKind(k), isExternal(true), isCOMDAT(false),
         writtenToSymtab(false), pendingArchiveLoad(false), isGCRoot(false),
-        isRuntimePseudoReloc(false), nameSize(n.size()),
-        nameData(n.empty() ? nullptr : n.data()) {}
+        isRuntimePseudoReloc(false), deferUndefined(false), canInline(true),
+        nameSize(n.size()), nameData(n.empty() ? nullptr : n.data()) {}
 
   const unsigned symbolKind : 8;
   unsigned isExternal : 1;
@@ -130,6 +130,16 @@ class Symbol {
 
   unsigned isRuntimePseudoReloc : 1;
 
+  // True if we want to allow this symbol to be undefined in the early
+  // undefined check pass in SymbolTable::reportUnresolvable(), as it
+  // might be fixed up later.
+  unsigned deferUndefined : 1;
+
+  // False if LTO shouldn't inline whatever this symbol points to. If a symbol
+  // is overwritten after LTO, LTO shouldn't inline the symbol because it
+  // doesn't know the final contents of the symbol.
+  unsigned canInline : 1;
+
 protected:
   // Symbol name length. Assume symbol lengths fit in a 32-bit integer.
   uint32_t nameSize;
@@ -468,7 +478,9 @@ void replaceSymbol(Symbol *s, ArgT &&... arg) {
                 "SymbolUnion not aligned enough");
   assert(static_cast<Symbol *>(static_cast<T *>(nullptr)) == nullptr &&
          "Not a Symbol");
+  bool canInline = s->canInline;
   new (s) T(std::forward<ArgT>(arg)...);
+  s->canInline = canInline;
 }
 } // namespace coff
 

lld/MinGW/Driver.cpp

@@ -377,6 +377,8 @@ bool mingw::link(ArrayRef<const char *> argsArr, bool canExitEarly,
     add("-includeoptional:" + StringRef(a->getValue()));
   for (auto *a : args.filtered(OPT_delayload))
     add("-delayload:" + StringRef(a->getValue()));
+  for (auto *a : args.filtered(OPT_wrap))
+    add("-wrap:" + StringRef(a->getValue()));
 
   std::vector<StringRef> searchPaths;
   for (auto *a : args.filtered(OPT_L)) {

lld/MinGW/Options.td

@@ -91,6 +91,8 @@ defm whole_archive: B<"whole-archive",
 def v: Flag<["-"], "v">, HelpText<"Display the version number">;
 def verbose: F<"verbose">, HelpText<"Verbose mode">;
 def version: F<"version">, HelpText<"Display the version number and exit">;
+defm wrap: Eq<"wrap", "Use wrapper functions for symbol">,
+     MetaVarName<"<symbol>">;
 
 // LLD specific options
 def _HASH_HASH_HASH : Flag<["-"], "###">,

lld/test/COFF/wrap-i386.s

@@ -0,0 +1,49 @@
+// REQUIRES: x86
+// 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 -out:%t.exe %t.main.obj %t.other.obj -entry:entry -subsystem:console -debug:symtab -safeseh:no -wrap:foo -wrap:nosuchsym
+// RUN: llvm-objdump -d --print-imm-hex %t.exe | FileCheck %s
+
+// CHECK: <_entry>:
+// CHECK-NEXT: movl $0x11010, %edx
+// CHECK-NEXT: movl $0x11010, %edx
+// CHECK-NEXT: movl $0x11000, %edx
+
+// RUN: llvm-readobj --symbols %t.exe > %t.dump
+// RUN: FileCheck --check-prefix=SYM1 %s < %t.dump
+// RUN: FileCheck --check-prefix=SYM2 %s < %t.dump
+// RUN: FileCheck --check-prefix=SYM3 %s < %t.dump
+
+// _foo = 0xffc11000 = 4290842624
+// ___wrap_foo = ffc11010 = 4290842640
+// SYM1:      Name: _foo
+// SYM1-NEXT: Value: 4290842624
+// SYM1-NEXT: Section: IMAGE_SYM_ABSOLUTE
+// SYM1-NEXT: BaseType: Null
+// SYM1-NEXT: ComplexType: Null
+// SYM1-NEXT: StorageClass: External
+// SYM2:      Name: ___wrap_foo
+// SYM2-NEXT: Value: 4290842640
+// SYM2-NEXT: Section: IMAGE_SYM_ABSOLUTE
+// SYM2-NEXT: BaseType: Null
+// SYM2-NEXT: ComplexType: Null
+// SYM2-NEXT: StorageClass: External
+// SYM3-NOT:  Name: ___real_foo
+
+#--- main.s
+.global _entry
+_entry:
+  movl $_foo, %edx
+  movl $___wrap_foo, %edx
+  movl $___real_foo, %edx
+
+#--- other.s
+.global _foo
+.global ___wrap_foo
+.global ___real_foo
+
+_foo = 0x11000
+___wrap_foo = 0x11010
+___real_foo = 0x11020

lld/test/COFF/wrap-import.ll

@@ -0,0 +1,36 @@
+// REQUIRES: x86
+
+// Check that wrapping works when the wrapped symbol is imported from a
+// different DLL.
+
+// 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
+// RUN: lld-link -out:%t.exe %t.main.bc %t.lib.lib -entry:entry -subsystem:console -wrap:func
+
+#--- main.ll
+target datalayout = "e-m:w-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-w64-windows-gnu"
+
+declare void @func()
+
+define void @entry() {
+  call void @func()
+  ret void
+}
+
+declare void @__real_func()
+
+define void @__wrap_func() {
+  call void @__real_func()
+  ret void
+}
+
+#--- lib.s
+.global func
+func:
+  ret

lld/test/COFF/wrap-lto-1.ll

@@ -0,0 +1,36 @@
+; REQUIRES: x86
+; LTO
+; RUN: llvm-as %s -o %t.obj
+; RUN: lld-link -out:%t.exe %t.obj -entry:entry -subsystem:console -wrap:bar -debug:symtab -lldsavetemps
+; RUN: cat %t.exe.resolution.txt | FileCheck -check-prefix=RESOLS %s
+
+; ThinLTO
+; RUN: opt -module-summary %s -o %t.obj
+; RUN: lld-link -out:%t.exe %t.obj -entry:entry -subsystem:console -wrap:bar -debug:symtab -lldsavetemps
+; RUN: cat %t.exe.resolution.txt | FileCheck -check-prefix=RESOLS %s
+
+; Make sure that the 'r' (linker redefined) bit is set for bar and __real_bar
+; in the resolutions file. The calls to bar and __real_bar will be routed to
+; __wrap_bar and bar, respectively. So they cannot be inlined.
+; RESOLS: ,bar,pxr{{$}}
+; RESOLS: ,__real_bar,xr{{$}}
+; RESOLS: ,__wrap_bar,px{{$}}
+
+target datalayout = "e-m:w-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-w64-windows-gnu"
+
+define void @bar() {
+  ret void
+}
+
+define void @entry() {
+  call void @bar()
+  ret void
+}
+
+declare void @__real_bar()
+
+define void @__wrap_bar() {
+  call void @__real_bar()
+  ret void
+}