[WebAssebmly] Add support for --wrap

The code for implementing this features is taken almost verbatim
from the ELF backend.

Fixes: https://bugs.llvm.org/show_bug.cgi?id=41681

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

llvm-svn: 361639

lld/ELF/Driver.cpp

@@ -1522,7 +1522,7 @@ static void wrapSymbols(ArrayRef<WrappedSymbol> Wrapped) {
 
   // Update pointers in input files.
   parallelForEach(ObjectFiles, [&](InputFile *File) {
-    std::vector<Symbol *> &Syms = File->getMutableSymbols();
+    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/ELF/InputFiles.h

@@ -90,7 +90,7 @@ public:
   // function on files of other types.
   ArrayRef<Symbol *> getSymbols() { return getMutableSymbols(); }
 
-  std::vector<Symbol *> &getMutableSymbols() {
+  MutableArrayRef<Symbol *> getMutableSymbols() {
     assert(FileKind == BinaryKind || FileKind == ObjKind ||
            FileKind == BitcodeKind);
     return Symbols;

lld/include/lld/Common/LLVM.h

@@ -29,6 +29,7 @@ class Twine;
 class MemoryBuffer;
 class MemoryBufferRef;
 template <typename T> class ArrayRef;
+template <typename T> class MutableArrayRef;
 template <unsigned InternalLen> class SmallString;
 template <typename T, unsigned N> class SmallVector;
 template <typename T> class ErrorOr;
@@ -62,6 +63,7 @@ using llvm::isa;
 
 // ADT's.
 using llvm::ArrayRef;
+using llvm::MutableArrayRef;
 using llvm::Error;
 using llvm::ErrorOr;
 using llvm::Expected;

lld/test/wasm/wrap.ll

@@ -0,0 +1,40 @@
+; RUN: llc -filetype=obj %s -o %t.o
+; RUN: wasm-ld -wrap nosuchsym -wrap foo -o %t.wasm %t.o
+; RUN: wasm-ld -emit-relocs -wrap foo -o %t.wasm %t.o
+; RUN: obj2yaml %t.wasm | FileCheck %s
+
+target triple = "wasm32-unknown-unknown"
+
+define i32 @foo() {
+  ret i32 1
+}
+
+define void @_start() {
+entry:
+  call i32 @foo()
+  ret void
+}
+
+declare i32 @__real_foo()
+
+define i32 @__wrap_foo() {
+  %rtn = call i32 @__real_foo()
+  ret i32 %rtn
+}
+
+; CHECK:      - Type:            CODE
+; CHECK-NEXT:   Relocations:     
+; CHECK-NEXT:     - Type:            R_WASM_FUNCTION_INDEX_LEB
+; CHECK-NEXT:       Index:           2
+; CHECK-NEXT:       Offset:          0x00000009
+; CHECK-NEXT:     - Type:            R_WASM_FUNCTION_INDEX_LEB
+; CHECK-NEXT:       Index:           0
+; CHECK-NEXT:       Offset:          0x00000013
+
+; CHECK:        FunctionNames:
+; CHECK-NEXT:      - Index:           0
+; CHECK-NEXT:        Name:            foo
+; CHECK-NEXT:      - Index:           1
+; CHECK-NEXT:        Name:            _start
+; CHECK-NEXT:      - Index:           2
+; CHECK-NEXT:        Name:            __wrap_foo

lld/wasm/Driver.cpp

@@ -535,6 +535,84 @@ static std::string createResponseFile(const opt::InputArgList &Args) {
   return Data.str();
 }
 
+// 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;
+};
+
+static Symbol *addUndefined(StringRef Name) {
+  return Symtab->addUndefinedFunction(Name, "", "", 0, nullptr, nullptr);
+}
+
+// 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.
+static std::vector<WrappedSymbol> 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->find(Name);
+    if (!Sym)
+      continue;
+
+    Symbol *Real = addUndefined(Saver.save("__real_" + Name));
+    Symbol *Wrap = addUndefined(Saver.save("__wrap_" + Name));
+    V.push_back({Sym, Real, Wrap});
+
+    // 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;
+    Wrap->IsUsedInRegularObj = true;
+    Real->IsUsedInRegularObj = false;
+  }
+  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.
+static void 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(Symtab->ObjectFiles, [&](InputFile *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;
+  });
+
+  // Update pointers in the symbol table.
+  for (const WrappedSymbol &W : Wrapped)
+    Symtab->wrap(W.Sym, W.Real, W.Wrap);
+}
+
 void LinkerDriver::link(ArrayRef<const char *> ArgsArr) {
   WasmOptTable Parser;
   opt::InputArgList Args = Parser.parse(ArgsArr.slice(1));
@@ -628,6 +706,9 @@ void LinkerDriver::link(ArrayRef<const char *> ArgsArr) {
   for (auto *Arg : Args.filtered(OPT_export))
     handleUndefined(Arg->getValue());
 
+  // Create wrapped symbols for -wrap option.
+  std::vector<WrappedSymbol> Wrapped = addWrappedSymbols(Args);
+
   // Do link-time optimization if given files are LLVM bitcode files.
   // This compiles bitcode files into real object files.
   Symtab->addCombinedLTOObject();
@@ -640,6 +721,10 @@ void LinkerDriver::link(ArrayRef<const char *> ArgsArr) {
   if (errorCount())
     return;
 
+  // Apply symbol renames for -wrap.
+  if (!Wrapped.empty())
+    wrapSymbols(Wrapped);
+
   for (auto *Arg : Args.filtered(OPT_export)) {
     Symbol *Sym = Symtab->find(Arg->getValue());
     if (Sym && Sym->isDefined())

lld/wasm/InputFiles.h

@@ -61,6 +61,8 @@ public:
 
   ArrayRef<Symbol *> getSymbols() const { return Symbols; }
 
+  MutableArrayRef<Symbol *> getMutableSymbols() { return Symbols; }
+
 protected:
   InputFile(Kind K, MemoryBufferRef M) : MB(M), FileKind(K) {}
   MemoryBufferRef MB;

lld/wasm/LTO.cpp

@@ -108,6 +108,11 @@ void BitcodeCompiler::add(BitcodeFile &F) {
                             (R.Prevailing && Sym->isExported());
     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/wasm/Options.td

@@ -112,6 +112,9 @@ def version: F<"version">, HelpText<"Display the version number and exit">;
 def z: JoinedOrSeparate<["-"], "z">, MetaVarName<"<option>">,
   HelpText<"Linker option extensions">;
 
+defm wrap: Eq<"wrap", "Use wrapper functions for symbol">,
+  MetaVarName<"<symbol>=<symbol>">;
+
 // The follow flags are unique to wasm
 
 def allow_undefined: F<"allow-undefined">,

lld/wasm/SymbolTable.cpp

@@ -65,7 +65,8 @@ void SymbolTable::addCombinedLTOObject() {
 }
 
 void SymbolTable::reportRemainingUndefines() {
-  for (Symbol *Sym : SymVector) {
+  for (const auto& Pair : SymMap) {
+    const Symbol *Sym = SymVector[Pair.second];
     if (!Sym->isUndefined() || Sym->isWeak())
       continue;
     if (Config->AllowUndefinedSymbols.count(Sym->getName()) != 0)
@@ -104,6 +105,7 @@ std::pair<Symbol *, bool> SymbolTable::insertName(StringRef Name) {
 
   Symbol *Sym = reinterpret_cast<Symbol *>(make<SymbolUnion>());
   Sym->IsUsedInRegularObj = false;
+  Sym->CanInline = true;
   Sym->Traced = Trace;
   SymVector.emplace_back(Sym);
   return {Sym, true};
@@ -539,6 +541,19 @@ void SymbolTable::trace(StringRef Name) {
   SymMap.insert({CachedHashStringRef(Name), -1});
 }
 
+void SymbolTable::wrap(Symbol *Sym, Symbol *Real, Symbol *Wrap) {
+  // Swap symbols as instructed by -wrap.
+  int &OrigIdx = SymMap[CachedHashStringRef(Sym->getName())];
+  int &RealIdx= SymMap[CachedHashStringRef(Real->getName())];
+  int &WrapIdx = SymMap[CachedHashStringRef(Wrap->getName())];
+  LLVM_DEBUG(dbgs() << "wrap: " << Sym->getName() << "\n");
+
+  // Anyone looking up __real symbols should get the original
+  RealIdx = OrigIdx;
+  // Anyone looking up the original should get the __wrap symbol
+  OrigIdx = WrapIdx;
+}
+
 static const uint8_t UnreachableFn[] = {
     0x03 /* ULEB length */, 0x00 /* ULEB num locals */,
     0x00 /* opcode unreachable */, 0x0b /* opcode end */

lld/wasm/SymbolTable.h

@@ -35,14 +35,11 @@ class InputSegment;
 // There is one add* function per symbol type.
 class SymbolTable {
 public:
-  void addFile(InputFile *File);
-  void addCombinedLTOObject();
+  void wrap(Symbol *Sym, Symbol *Real, Symbol *Wrap);
 
-  std::vector<ObjFile *> ObjectFiles;
-  std::vector<SharedFile *> SharedFiles;
-  std::vector<BitcodeFile *> BitcodeFiles;
-  std::vector<InputFunction *> SyntheticFunctions;
-  std::vector<InputGlobal *> SyntheticGlobals;
+  void addFile(InputFile *File);
+
+  void addCombinedLTOObject();
 
   void reportRemainingUndefines();
 
@@ -87,6 +84,12 @@ public:
   void handleSymbolVariants();
   void handleWeakUndefines();
 
+  std::vector<ObjFile *> ObjectFiles;
+  std::vector<SharedFile *> SharedFiles;
+  std::vector<BitcodeFile *> BitcodeFiles;
+  std::vector<InputFunction *> SyntheticFunctions;
+  std::vector<InputGlobal *> SyntheticGlobals;
+
 private:
   std::pair<Symbol *, bool> insert(StringRef Name, const InputFile *File);
   std::pair<Symbol *, bool> insertName(StringRef Name);

lld/wasm/Symbols.h

@@ -114,6 +114,11 @@ public:
   // command line flag)
   unsigned ForceExport : 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;
+
   // True if this symbol is specified by --trace-symbol option.
   unsigned Traced : 1;
 
@@ -131,8 +136,8 @@ public:
 
 protected:
   Symbol(StringRef Name, Kind K, uint32_t Flags, InputFile *F)
-      : IsUsedInRegularObj(false), ForceExport(false), Traced(false),
-        Name(Name), SymbolKind(K), Flags(Flags), File(F),
+      : IsUsedInRegularObj(false), ForceExport(false), CanInline(false),
+        Traced(false), Name(Name), SymbolKind(K), Flags(Flags), File(F),
         Referenced(!Config->GcSections) {}
 
   StringRef Name;
@@ -474,6 +479,7 @@ T *replaceSymbol(Symbol *S, ArgT &&... Arg) {
   T *S2 = new (S) T(std::forward<ArgT>(Arg)...);
   S2->IsUsedInRegularObj = SymCopy.IsUsedInRegularObj;
   S2->ForceExport = SymCopy.ForceExport;
+  S2->CanInline = SymCopy.CanInline;
   S2->Traced = SymCopy.Traced;
 
   // Print out a log message if --trace-symbol was specified.