[lld][WebAssembly] Initial support for stub libraries

See the docs in lld/docs/WebAssembly.rst for more on this.

This feature unlocks a lot of simplification in the emscripten toolchain
since we can represent the JS libraries to wasm-ld as stub libraries.

See https://github.com/emscripten-core/emscripten/issues/18875

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

9 files changed, 133 insertions, 5 deletions

diff --git a/lld/wasm/Driver.cpp b/lld/wasm/Driver.cpp
index 310f9df2d5b6..68cd8cabbd7f 100644
--- a/lld/wasm/Driver.cpp
+++ b/lld/wasm/Driver.cpp
@@ -279,6 +279,12 @@ void LinkerDriver::addFile(StringRef path) {
   case file_magic::wasm_object:
     files.push_back(createObjectFile(mbref));
     break;
+  case file_magic::unknown:
+    if (mbref.getBuffer().starts_with("#STUB\n")) {
+      files.push_back(make<StubFile>(mbref));
+      break;
+    }
+    [[fallthrough]];
   default:
     error("unknown file type: " + mbref.getBufferIdentifier());
   }
@@ -868,6 +874,53 @@ static void createOptionalSymbols() {
     WasmSym::tlsBase = createOptionalGlobal("__tls_base", false);
 }
 
+static void processStubLibraries() {
+  log("-- processStubLibraries");
+  for (auto &stub_file : symtab->stubFiles) {
+    LLVM_DEBUG(llvm::dbgs()
+               << "processing stub file: " << stub_file->getName() << "\n");
+    for (auto [name, deps]: stub_file->symbolDependencies) {
+      auto* sym = symtab->find(name);
+      if (!sym || !sym->isUndefined() || !sym->isUsedInRegularObj ||
+          sym->forceImport) {
+        LLVM_DEBUG(llvm::dbgs() << "stub not in needed: " << name << "\n");
+        continue;
+      }
+      // The first stub library to define a given symbol sets this and
+      // definitions in later stub libraries are ignored.
+      sym->forceImport = true;
+      if (sym->traced)
+        message(toString(stub_file) + ": importing " + name);
+      else
+        LLVM_DEBUG(llvm::dbgs()
+                   << toString(stub_file) << ": importing " << name << "\n");
+      for (const auto dep : deps) {
+        auto* needed = symtab->find(dep);
+        if (!needed) {
+          error(toString(stub_file) + ": undefined symbol: " + dep +
+                ". Required by " + toString(*sym));
+        } else if (needed->isUndefined()) {
+          error(toString(stub_file) +
+                ": undefined symbol: " + toString(*needed) +
+                ". Required by " + toString(*sym));
+        } else {
+          LLVM_DEBUG(llvm::dbgs()
+                     << "force export: " << toString(*needed) << "\n");
+          needed->forceExport = true;
+          needed->isUsedInRegularObj = true;
+          if (auto *lazy = dyn_cast<LazySymbol>(needed)) {
+            lazy->fetch();
+            if (!config->whyExtract.empty())
+              config->whyExtractRecords.emplace_back(stub_file->getName(),
+                                                     sym->getFile(), *sym);
+          }
+        }
+      }
+    }
+  }
+  log("-- done processStubLibraries");
+}
+
 // Reconstructs command line arguments so that so that you can re-run
 // the same command with the same inputs. This is for --reproduce.
 static std::string createResponseFile(const opt::InputArgList &args) {
@@ -1166,6 +1219,8 @@ void LinkerDriver::linkerMain(ArrayRef<const char *> argsArr) {
   if (errorCount())
     return;
 
+  processStubLibraries();
+
   createOptionalSymbols();
 
   // Resolve any variant symbols that were created due to signature
diff --git a/lld/wasm/InputFiles.cpp b/lld/wasm/InputFiles.cpp
index 75760293bbaa..2d9768c768f2 100644
--- a/lld/wasm/InputFiles.cpp
+++ b/lld/wasm/InputFiles.cpp
@@ -12,6 +12,7 @@
 #include "InputElement.h"
 #include "OutputSegment.h"
 #include "SymbolTable.h"
+#include "lld/Common/Args.h"
 #include "lld/Common/CommonLinkerContext.h"
 #include "lld/Common/Reproduce.h"
 #include "llvm/Object/Binary.h"
@@ -678,6 +679,48 @@ Symbol *ObjFile::createUndefined(const WasmSymbol &sym, bool isCalledDirectly) {
   llvm_unreachable("unknown symbol kind");
 }
 
+
+StringRef strip(StringRef s) {
+  while (s.starts_with(" ")) {
+    s = s.drop_front();
+  }
+  while (s.ends_with(" ")) {
+    s = s.drop_back();
+  }
+  return s;
+}
+
+void StubFile::parse() {
+  bool first = false;
+
+  for (StringRef line : args::getLines(mb)) {
+    // File must begin with #STUB
+    if (first) {
+      assert(line == "#STUB\n");
+      first = false;
+    }
+
+    // Lines starting with # are considered comments
+    if (line.startswith("#"))
+      continue;
+
+    StringRef sym;
+    StringRef rest;
+    std::tie(sym, rest) = line.split(':');
+    sym = strip(sym);
+    rest = strip(rest);
+
+    symbolDependencies[sym] = {};
+
+    while (rest.size()) {
+      StringRef first;
+      std::tie(first, rest) = rest.split(',');
+      first = strip(first);
+      symbolDependencies[sym].push_back(first);
+    }
+  }
+}
+
 void ArchiveFile::parse() {
   // Parse a MemoryBufferRef as an archive file.
   LLVM_DEBUG(dbgs() << "Parsing library: " << toString(this) << "\n");
diff --git a/lld/wasm/InputFiles.h b/lld/wasm/InputFiles.h
index c72f64cb2bd0..11cee5405b65 100644
--- a/lld/wasm/InputFiles.h
+++ b/lld/wasm/InputFiles.h
@@ -47,6 +47,7 @@ public:
     SharedKind,
     ArchiveKind,
     BitcodeKind,
+    StubKind,
   };
 
   virtual ~InputFile() {}
@@ -183,6 +184,18 @@ public:
   static bool doneLTO;
 };
 
+// Stub libray (See docs/WebAssembly.rst)
+class StubFile : public InputFile {
+public:
+  explicit StubFile(MemoryBufferRef m) : InputFile(StubKind, m) {}
+
+  static bool classof(const InputFile *f) { return f->kind() == StubKind; }
+
+  void parse();
+
+  llvm::DenseMap<StringRef, std::vector<StringRef>> symbolDependencies;
+};
+
 inline bool isBitcode(MemoryBufferRef mb) {
   return identify_magic(mb.getBuffer()) == llvm::file_magic::bitcode;
 }
diff --git a/lld/wasm/Relocations.cpp b/lld/wasm/Relocations.cpp
index 2f6dd6af2d03..ce41cdcb3e07 100644
--- a/lld/wasm/Relocations.cpp
+++ b/lld/wasm/Relocations.cpp
@@ -32,9 +32,9 @@ static bool requiresGOTAccess(const Symbol *sym) {
 }
 
 static bool allowUndefined(const Symbol* sym) {
-  // Symbols with explicit import names are always allowed to be undefined at
+  // Symbols that are explicitly imported are always allowed to be undefined at
   // link time.
-  if (sym->importName)
+  if (sym->isImported())
     return true;
   if (isa<UndefinedFunction>(sym) && config->importUndefined)
     return true;
diff --git a/lld/wasm/SymbolTable.cpp b/lld/wasm/SymbolTable.cpp
index 881b1231ffdf..d33176a0fa54 100644
--- a/lld/wasm/SymbolTable.cpp
+++ b/lld/wasm/SymbolTable.cpp
@@ -38,6 +38,13 @@ void SymbolTable::addFile(InputFile *file) {
     return;
   }
 
+  // stub file
+  if (auto *f = dyn_cast<StubFile>(file)) {
+    f->parse();
+    stubFiles.push_back(f);
+    return;
+  }
+
   if (config->trace)
     message(toString(file));
 
diff --git a/lld/wasm/SymbolTable.h b/lld/wasm/SymbolTable.h
index 5009e6039602..ef2a023b68c4 100644
--- a/lld/wasm/SymbolTable.h
+++ b/lld/wasm/SymbolTable.h
@@ -102,6 +102,7 @@ public:
   DefinedFunction *createUndefinedStub(const WasmSignature &sig);
 
   std::vector<ObjFile *> objectFiles;
+  std::vector<StubFile *> stubFiles;
   std::vector<SharedFile *> sharedFiles;
   std::vector<BitcodeFile *> bitcodeFiles;
   std::vector<InputFunction *> syntheticFunctions;
diff --git a/lld/wasm/Symbols.cpp b/lld/wasm/Symbols.cpp
index 8864e840dd58..567ff49dfa44 100644
--- a/lld/wasm/Symbols.cpp
+++ b/lld/wasm/Symbols.cpp
@@ -221,6 +221,10 @@ void Symbol::setHidden(bool isHidden) {
     flags |= WASM_SYMBOL_VISIBILITY_DEFAULT;
 }
 
+bool Symbol::isImported() const {
+  return isUndefined() && (importName.has_value() || forceImport);
+}
+
 bool Symbol::isExported() const {
   if (!isDefined() || isLocal())
     return false;
diff --git a/lld/wasm/Symbols.h b/lld/wasm/Symbols.h
index 16f1b535876e..34fff4b962bd 100644
--- a/lld/wasm/Symbols.h
+++ b/lld/wasm/Symbols.h
@@ -114,6 +114,7 @@ public:
   void setOutputSymbolIndex(uint32_t index);
 
   WasmSymbolType getWasmType() const;
+  bool isImported() const;
   bool isExported() const;
   bool isExportedExplicit() const;
 
@@ -135,7 +136,8 @@ protected:
   Symbol(StringRef name, Kind k, uint32_t flags, InputFile *f)
       : name(name), file(f), symbolKind(k), referenced(!config->gcSections),
         requiresGOT(false), isUsedInRegularObj(false), forceExport(false),
-        canInline(false), traced(false), isStub(false), flags(flags) {}
+        forceImport(false), canInline(false), traced(false), isStub(false),
+        flags(flags) {}
 
   StringRef name;
   InputFile *file;
@@ -160,6 +162,8 @@ public:
   // -e/--export command line flag)
   bool forceExport : 1;
 
+  bool forceImport : 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.
@@ -661,6 +665,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->forceImport = symCopy.forceImport;
   s2->canInline = symCopy.canInline;
   s2->traced = symCopy.traced;
   s2->referenced = symCopy.referenced;
diff --git a/lld/wasm/Writer.cpp b/lld/wasm/Writer.cpp
index 030ef7468791..d9e87276b31b 100644
--- a/lld/wasm/Writer.cpp
+++ b/lld/wasm/Writer.cpp
@@ -744,7 +744,7 @@ static bool shouldImport(Symbol *sym) {
   if (config->allowUndefinedSymbols.count(sym->getName()) != 0)
     return true;
 
-  return sym->importName.has_value();
+  return sym->isImported();
 }
 
 void Writer::calculateImports() {
@@ -1709,7 +1709,7 @@ void Writer::run() {
       sym->forceExport = true;
   }
 
-  // Delay reporting error about explicit exports until after
+  // Delay reporting errors about explicit exports until after
   // addStartStopSymbols which can create optional symbols.
   for (auto &name : config->requiredExports) {
     Symbol *sym = symtab->find(name);