[lld][WebAssembly] stub objects: Fix handling of LTO libcall dependencies

This actually simplifies the code by performs a pre-pass of the stub
objects prior to LTO.

This should be the final change needed before we can make the switch
on the emscripten side: https://github.com/emscripten-core/emscripten/pull/18905

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

2 files changed, 33 insertions, 15 deletions

diff --git a/lld/test/wasm/lto/stub-library-libcall.s b/lld/test/wasm/lto/stub-library-libcall.s
index 3989e204f0f2..ce88a32dd99d 100644
--- a/lld/test/wasm/lto/stub-library-libcall.s
+++ b/lld/test/wasm/lto/stub-library-libcall.s
@@ -1,6 +1,6 @@
 # RUN: split-file %s %t
 # RUN: llvm-mc -filetype=obj -triple=wasm32-unknown-unknown -o %t_main.o %t/main.s
-# RUN: llvm-mc -filetype=obj -triple=wasm32-unknown-unknown -o %t_foo.o %t/foo.s
+# RUN: llvm-as %S/Inputs/foo.ll -o %t_foo.o
 # RUN: llvm-as %S/Inputs/libcall.ll -o %t_libcall.o
 # RUN: wasm-ld %t_main.o %t_libcall.o %t_foo.o %p/Inputs/stub.so -o %t.wasm
 # RUN: obj2yaml %t.wasm | FileCheck %s
@@ -25,12 +25,6 @@ _start:
     call func_with_libcall
     end_function
 
-#--- foo.s
-.globl foo
-foo:
-  .functype foo () -> ()
-  end_function
-
 # CHECK:         Imports:
 # CHECK-NEXT:      - Module:          env
 # CHECK-NEXT:        Field:           memcpy
@@ -46,4 +40,4 @@ foo:
 # CHECK-NEXT:         Index:           1
 # CHECK-NEXT:       - Name:            foo
 # CHECK-NEXT:         Kind:            FUNCTION
-# CHECK-NEXT:         Index:           2
+# CHECK-NEXT:         Index:           3
diff --git a/lld/wasm/Driver.cpp b/lld/wasm/Driver.cpp
index 8b02c9a85e26..baca9395f5b0 100644
--- a/lld/wasm/Driver.cpp
+++ b/lld/wasm/Driver.cpp
@@ -874,6 +874,28 @@ static void createOptionalSymbols() {
     WasmSym::tlsBase = createOptionalGlobal("__tls_base", false);
 }
 
+static void processStubLibrariesPreLTO() {
+  log("-- processStubLibrariesPreLTO");
+  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 the symbol is not present at all (yet), or if it is present but
+      // undefined, then mark the dependent symbols as used by a regular
+      // object so they will be preserved and exported by the LTO process.
+      if (!sym || sym->isUndefined()) {
+        for (const auto dep : deps) {
+          auto* needed = symtab->find(dep);
+          if (needed ) {
+            needed->isUsedInRegularObj = true;
+          }
+        }
+      }
+    }
+  }
+}
+
 static void processStubLibraries() {
   log("-- processStubLibraries");
   for (auto &stub_file : symtab->stubFiles) {
@@ -881,12 +903,14 @@ static void processStubLibraries() {
                << "processing stub file: " << stub_file->getName() << "\n");
     for (auto [name, deps]: stub_file->symbolDependencies) {
       auto* sym = symtab->find(name);
-      if (!sym || !sym->isUndefined() || sym->forceImport) {
-        LLVM_DEBUG(llvm::dbgs() << "stub not in needed: " << name << "\n");
+      if (!sym || !sym->isUndefined()) {
+        LLVM_DEBUG(llvm::dbgs() << "stub symbol not needed: " << name << "\n");
         continue;
       }
       // The first stub library to define a given symbol sets this and
       // definitions in later stub libraries are ignored.
+      if (sym->forceImport)
+        continue;  // Already handled
       sym->forceImport = true;
       if (sym->traced)
         message(toString(stub_file) + ": importing " + name);
@@ -1213,9 +1237,9 @@ void LinkerDriver::linkerMain(ArrayRef<const char *> argsArr) {
   if (errorCount())
     return;
 
-  // processStubLibraries must happen before LTO because it can trigger the
-  // export of arbirary symbols that might themselves be defined in LTO objects.
-  processStubLibraries();
+  // We process the stub libraries once beofore LTO to ensure that any possible
+  // required exports are preserved by the LTO process.
+  processStubLibrariesPreLTO();
 
   // Do link-time optimization if given files are LLVM bitcode files.
   // This compiles bitcode files into real object files.
@@ -1225,8 +1249,8 @@ void LinkerDriver::linkerMain(ArrayRef<const char *> argsArr) {
 
   // The LTO process can generate new undefined symbols, specifically libcall
   // functions.  Because those symbols might be declared in a stub library we
-  // need the process the stub libraries once again after LTO to handle any
-  // newly undefined symbols.
+  // need the process the stub libraries once again after LTO to handle all
+  // undefined symbols, including ones that didn't exist prior to LTO.
   processStubLibraries();
 
   writeWhyExtract();