[lld][WebAssembly] Process stub libraries before performing LTO

There are cases where stub library processing can trigger new exports
which might require them to be included at LTO time.

Specifically `processStubLibraries` marks symbols as `forceExports`
which even effect the LTO process.

And since the LTO process can generate new undefined symbols
(specifically libcall function) we need to also process the stub
libraries after LTO.

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

7 files changed, 114 insertions, 4 deletions

diff --git a/lld/test/wasm/lto/Inputs/foo.ll b/lld/test/wasm/lto/Inputs/foo.ll
new file mode 100644
index 000000000000..5fd7651b49cb
--- /dev/null
+++ b/lld/test/wasm/lto/Inputs/foo.ll
@@ -0,0 +1,7 @@
+target datalayout = "e-m:e-p:32:32-p10:8:8-p20:8:8-i64:64-n32:64-S128"
+target triple = "wasm32-unknown-unknown"
+
+define void @foo() local_unnamed_addr {
+entry:
+  ret void
+}
diff --git a/lld/test/wasm/lto/Inputs/libcall.ll b/lld/test/wasm/lto/Inputs/libcall.ll
new file mode 100644
index 000000000000..23eb1f756284
--- /dev/null
+++ b/lld/test/wasm/lto/Inputs/libcall.ll
@@ -0,0 +1,12 @@
+target datalayout = "e-m:e-p:32:32-p10:8:8-p20:8:8-i64:64-n32:64-S128"
+target triple = "wasm32-unknown-unknown"
+
+; This function, when compiled will generate a new undefined reference to
+; memcpy, that is not present in the bitcode object
+define void @func_with_libcall(ptr %a, ptr %b) {
+entry:
+  call void @llvm.memcpy.p0.p0.i64(ptr %a, ptr %b, i64 1024, i1 false)
+  ret void
+}
+
+declare void @llvm.memcpy.p0.p0.i64(ptr nocapture, ptr nocapture, i64, i1)
diff --git a/lld/test/wasm/lto/Inputs/stub.so b/lld/test/wasm/lto/Inputs/stub.so
new file mode 100644
index 000000000000..e76c890dbb43
--- /dev/null
+++ b/lld/test/wasm/lto/Inputs/stub.so
@@ -0,0 +1,3 @@
+#STUB
+bar: foo
+memcpy: foo
diff --git a/lld/test/wasm/lto/stub-library-libcall.s b/lld/test/wasm/lto/stub-library-libcall.s
new file mode 100644
index 000000000000..3989e204f0f2
--- /dev/null
+++ b/lld/test/wasm/lto/stub-library-libcall.s
@@ -0,0 +1,49 @@
+# 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/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
+
+# The function `func_with_libcall` will generate an undefined reference to
+# `memcpy` at LTO time.  `memcpy` itself also declared in stub.so and depends
+# on `foo`
+
+# If %t_foo.o is not included in the link we get an undefined symbol reported
+# to the dependency of memcpy on the foo export:
+
+# RUN: not wasm-ld %t_main.o %t_libcall.o %p/Inputs/stub.so -o %t.wasm 2>&1 | FileCheck --check-prefix=MISSING %s
+# MISSING: stub.so: undefined symbol: foo. Required by memcpy
+
+#--- main.s
+.functype func_with_libcall (i32, i32) -> ()
+.globl _start
+_start:
+    .functype _start () -> ()
+    i32.const 0
+    i32.const 0
+    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
+# CHECK-NEXT:        Kind:            FUNCTION
+# CHECK-NEXT:        SigIndex:        0
+
+# CHECK:         Exports:
+# CHECK-NEXT:       - Name:            memory
+# CHECK-NEXT:         Kind:            MEMORY
+# CHECK-NEXT:         Index:           0
+# CHECK-NEXT:       - Name:            _start
+# CHECK-NEXT:         Kind:            FUNCTION
+# CHECK-NEXT:         Index:           1
+# CHECK-NEXT:       - Name:            foo
+# CHECK-NEXT:         Kind:            FUNCTION
+# CHECK-NEXT:         Index:           2
diff --git a/lld/test/wasm/lto/stub-library.s b/lld/test/wasm/lto/stub-library.s
new file mode 100644
index 000000000000..20e2a6221141
--- /dev/null
+++ b/lld/test/wasm/lto/stub-library.s
@@ -0,0 +1,33 @@
+# RUN: llvm-mc -filetype=obj -triple=wasm32-unknown-unknown -o %t.o %s
+# RUN: llvm-as %S/Inputs/foo.ll -o %t1.o
+# RUN: wasm-ld %t.o %t1.o %p/Inputs/stub.so -o %t.wasm
+# RUN: obj2yaml %t.wasm | FileCheck %s
+
+# The function `bar` is declared in stub.so and depends on `foo`, which happens
+# be in an LTO object.
+# This verifies that stub library dependencies (required exports) can be defined
+# in LTO objects.
+.functype bar () -> ()
+
+.globl _start
+_start:
+    .functype _start () -> ()
+    call bar
+    end_function
+
+# CHECK:         Imports:
+# CHECK-NEXT:      - Module:          env
+# CHECK-NEXT:        Field:           bar
+# CHECK-NEXT:        Kind:            FUNCTION
+# CHECK-NEXT:        SigIndex:        0
+
+# CHECK:         Exports:
+# CHECK-NEXT:       - Name:            memory
+# CHECK-NEXT:         Kind:            MEMORY
+# CHECK-NEXT:         Index:           0
+# CHECK-NEXT:       - Name:            _start
+# CHECK-NEXT:         Kind:            FUNCTION
+# CHECK-NEXT:         Index:           1
+# CHECK-NEXT:       - Name:            foo
+# CHECK-NEXT:         Kind:            FUNCTION
+# CHECK-NEXT:         Index:           2
diff --git a/lld/test/wasm/stub_library.s b/lld/test/wasm/stub-library.s
index 9cbf2505ea9e..9cbf2505ea9e 100644
--- a/lld/test/wasm/stub_library.s
+++ b/lld/test/wasm/stub-library.s
diff --git a/lld/wasm/Driver.cpp b/lld/wasm/Driver.cpp
index 68cd8cabbd7f..7ce221e1f663 100644
--- a/lld/wasm/Driver.cpp
+++ b/lld/wasm/Driver.cpp
@@ -881,8 +881,7 @@ 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->isUsedInRegularObj ||
-          sym->forceImport) {
+      if (!sym || !sym->isUndefined() || sym->forceImport) {
         LLVM_DEBUG(llvm::dbgs() << "stub not in needed: " << name << "\n");
         continue;
       }
@@ -907,7 +906,6 @@ static void processStubLibraries() {
           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())
@@ -1211,7 +1209,9 @@ void LinkerDriver::linkerMain(ArrayRef<const char *> argsArr) {
   if (errorCount())
     return;
 
-  writeWhyExtract();
+  // processStubLibraries must happen before LTO because it can trigger the
+  // export of arbirary symbols that might themselves be defined in LTO objects.
+  processStubLibraries();
 
   // Do link-time optimization if given files are LLVM bitcode files.
   // This compiles bitcode files into real object files.
@@ -1219,8 +1219,14 @@ void LinkerDriver::linkerMain(ArrayRef<const char *> argsArr) {
   if (errorCount())
     return;
 
+  // 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.
   processStubLibraries();
 
+  writeWhyExtract();
+
   createOptionalSymbols();
 
   // Resolve any variant symbols that were created due to signature