[WebAssebmly] Add support for -wrap

[sbc100]

Summary

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

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

[ruiu]

LGTM

[sbc100]

• improve test

lld/ELF/Driver.cpp

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(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;
   });
 
   // Update pointers in the symbol table.
   for (const WrappedSymbol &W : Wrapped)
     Symtab->wrap(W.Sym, W.Real, W.Wrap);

lld/ELF/InputFiles.h

ArrayRef<InputSectionBase *> getSections() const {
     assert(FileKind == ObjKind || FileKind == BinaryKind);
     return Sections;
   }
 
   // Returns object file symbols. It is a runtime error to call this
   // 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;
   }
 
   // Filename of .a which contained this file. If this file was
   // not in an archive file, it is the empty string. We use this
   // string for creating error messages.

lld/include/lld/Common/LLVM.h

 // ADT's.
 class raw_ostream;
 class Error;
 class StringRef;
 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;
 template <typename T> class Expected;
 
 namespace object {
 class WasmObjectFile;
 struct WasmSection;

 using llvm::cast;
 using llvm::cast_or_null;
 using llvm::dyn_cast;
 using llvm::dyn_cast_or_null;
 using llvm::isa;
 
 // ADT's.
 using llvm::ArrayRef;
+using llvm::MutableArrayRef;
 using llvm::Error;
 using llvm::ErrorOr;
 using llvm::Expected;
 using llvm::MemoryBuffer;
 using llvm::MemoryBufferRef;
 using llvm::raw_ostream;
 using llvm::SmallString;
 using llvm::SmallVector;

lld/test/wasm/wrap.ll

+; RUN: llc -filetype=obj %s -o %t.o
+; RUN: wasm-ld -wrap nosuchsym -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
+}
+
+define i32 @__wrap_foo() {
+  ret i32 2
+}
+
+; CHECK:      - Type:            CODE
+; CHECK-NEXT:   Relocations:     
+; CHECK-NEXT:     - Type:            R_WASM_FUNCTION_INDEX_LEB
+; CHECK-NEXT:       Index:           1
+; CHECK-NEXT:       Offset:          0x00000004
+
+; CHECK:        FunctionNames:
+; CHECK-NEXT:      - Index:           0
+; CHECK-NEXT:        Name:            _start
+; CHECK-NEXT:      - Index:           1
+; CHECK-NEXT:        Name:            __wrap_foo

lld/wasm/Driver.cpp

case OPT_o:
       break;
     default:
       OS << toString(*Arg) << "\n";
     }
   }
   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;
+  }
+  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));
 
   // Handle --help
   if (Args.hasArg(OPT_help)) {
     Parser.PrintHelp(outs(),
                      (std::string(ArgsArr[0]) + " [options] file...").c_str(),
void LinkerDriver::link(ArrayRef<const char *> ArgsArr) {
   if (errorCount())
     return;
 
   // Handle the `--export <sym>` options
   // This works like --undefined but also exports the symbol if its found
   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();
   if (errorCount())
     return;
 
   // Resolve any variant symbols that were created due to signature
   // mismatchs.
   Symtab->handleSymbolVariants();
   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())
       Sym->ForceExport = true;
     else if (!Config->AllowUndefined)
       error(Twine("symbol exported via --export not found: ") +
             Arg->getValue());
   }

lld/wasm/InputFiles.h

public:
 
   Kind kind() const { return FileKind; }
 
   // An archive file name if this file is created from an archive.
   StringRef ArchiveName;
 
   ArrayRef<Symbol *> getSymbols() const { return Symbols; }
 
+  MutableArrayRef<Symbol *> getMutableSymbols() { return Symbols; }
+
 protected:
   InputFile(Kind K, MemoryBufferRef M) : MB(M), FileKind(K) {}
   MemoryBufferRef MB;
 
   // List of all symbols referenced or defined by this file.
   std::vector<Symbol *> Symbols;
 
 private:

lld/wasm/LTO.cpp

for (const lto::InputFile::Symbol &ObjSym : Obj.symbols()) {
     // flags an undefined in IR with a definition in ASM as prevailing.
     // Once IRObjectFile is fixed to report only one symbol this hack can
     // be removed.
     R.Prevailing = !ObjSym.isUndefined() && Sym->getFile() == &F;
     R.VisibleToRegularObj = Config->Relocatable || Sym->IsUsedInRegularObj ||
                             (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));
 }
 
 // Merge all the bitcode files we have seen, codegen the result
 // and return the resulting objects.
 std::vector<StringRef> BitcodeCompiler::compile() {
   unsigned MaxTasks = LTOObj->getMaxTasks();

lld/wasm/Options.td

 
 def verbose: F<"verbose">, HelpText<"Verbose mode">;
 
 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">,
   HelpText<"Allow undefined symbols in linked binary">;
 
 def allow_undefined_file: J<"allow-undefined-file=">,
   HelpText<"Allow symbols listed in <file> to be undefined in linked binary">;
 

lld/wasm/SymbolTable.h

 // conflicts. For example, obviously, a defined symbol is better than
 // an undefined symbol. Or, if there's a conflict between a lazy and a
 // undefined, it'll read an archive member to read a real definition
 // to replace the lazy symbol. The logic is implemented in the
 // add*() functions, which are called by input files as they are parsed.
 // There is one add* function per symbol type.
 class SymbolTable {
 public:
+  void wrap(Symbol *Sym, Symbol *Real, Symbol *Wrap);
+
   void addFile(InputFile *File);
-  void addCombinedLTOObject();
 
-  std::vector<ObjFile *> ObjectFiles;
-  std::vector<SharedFile *> SharedFiles;
-  std::vector<BitcodeFile *> BitcodeFiles;
-  std::vector<InputFunction *> SyntheticFunctions;
-  std::vector<InputGlobal *> SyntheticGlobals;
+  void addCombinedLTOObject();
 
   void reportRemainingUndefines();
 
   ArrayRef<Symbol *> getSymbols() const { return SymVector; }
 
   Symbol *find(StringRef Name);
 
   void replace(StringRef Name, Symbol* Sym);
public:
   DefinedFunction *addSyntheticFunction(StringRef Name, uint32_t Flags,
                                         InputFunction *Function);
   DefinedData *addOptionalDataSymbol(StringRef Name, uint32_t Value,
                                      uint32_t Flags);
 
   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);
 
   bool getFunctionVariant(Symbol* Sym, const WasmSignature *Sig,
                           const InputFile *File, Symbol **Out);
   InputFunction *replaceWithUnreachable(Symbol *Sym, const WasmSignature &Sig,
                                         StringRef DebugName);

lld/wasm/SymbolTable.cpp

void SymbolTable::addCombinedLTOObject() {
   for (StringRef Filename : LTO->compile()) {
     auto *Obj = make<ObjFile>(MemoryBufferRef(Filename, "lto.tmp"), "");
     Obj->parse(true);
     ObjectFiles.push_back(Obj);
   }
 }
 
 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)
       continue;
     if (!Sym->IsUsedInRegularObj)
       continue;
     error(toString(Sym->getFile()) + ": undefined symbol: " + toString(*Sym));
   }
if (SymIndex == -1) {
     IsNew = true;
   }
 
   if (!IsNew)
     return {SymVector[SymIndex], false};
 
   Symbol *Sym = reinterpret_cast<Symbol *>(make<SymbolUnion>());
   Sym->IsUsedInRegularObj = false;
+  Sym->CanInline = true;
   Sym->Traced = Trace;
   SymVector.emplace_back(Sym);
   return {Sym, true};
 }
 
 std::pair<Symbol *, bool> SymbolTable::insert(StringRef Name,
                                               const InputFile *File) {
   Symbol *S;

 }
 
 // Set a flag for --trace-symbol so that we can print out a log message
 // if a new symbol with the same name is inserted into the symbol table.
 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 &Idx1 = SymMap[CachedHashStringRef(Sym->getName())];
+  int &Idx2 = SymMap[CachedHashStringRef(Real->getName())];
+  int &Idx3 = SymMap[CachedHashStringRef(Wrap->getName())];
+  dbgs() << "wrapping " << Sym->getName() << "\n";
+
+  Idx2 = Idx1;
+  Idx1 = Idx3;
+}
+
 static const uint8_t UnreachableFn[] = {
     0x03 /* ULEB length */, 0x00 /* ULEB num locals */,
     0x00 /* opcode unreachable */, 0x0b /* opcode end */
 };
 
 // Replace the given symbol body with an unreachable function.
 // This is used by handleWeakUndefines in order to generate a callable
 // equivalent of an undefined function and also handleSymbolVariants for

lld/wasm/Symbols.h

public:
   // unreferenced DSO symbols, lazy (archive) symbols, and bitcode symbols that
   // are unreferenced except by other bitcode objects.
   unsigned IsUsedInRegularObj : 1;
 
   // True if ths symbol is explicity marked for export (i.e. via the -e/--export
   // 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;
 
   const WasmSignature* getSignature() const;
 
   bool isInGOT() const { return GOTIndex != INVALID_INDEX; }
 
   uint32_t getGOTIndex() const {
     assert(GOTIndex != INVALID_INDEX);
     return GOTIndex;
   }
 
   void setGOTIndex(uint32_t Index);
   bool hasGOTIndex() const { return GOTIndex != INVALID_INDEX; }
 
 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;
   Kind SymbolKind;
   uint32_t Flags;
   InputFile *File;
   uint32_t OutputSymbolIndex = INVALID_INDEX;
   uint32_t GOTIndex = INVALID_INDEX;
T *replaceSymbol(Symbol *S, ArgT &&... Arg) {
   assert(static_cast<Symbol *>(static_cast<T *>(nullptr)) == nullptr &&
          "Not a Symbol");
 
   Symbol SymCopy = *S;
 
   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.
   // This is for debugging.
   if (S2->Traced)
     printTraceSymbol(S2);
 
   return S2;