[WebAssembly] Fix --export of LTO symbols

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

llvm-svn: 335881

4 files changed, 53 insertions, 13 deletions

diff --git a/lld/test/wasm/lto/export.ll b/lld/test/wasm/lto/export.ll
new file mode 100644
index 00000000000..44ded6f147f
--- /dev/null
+++ b/lld/test/wasm/lto/export.ll
@@ -0,0 +1,38 @@
+; RUN: llvm-as -o %t.bc %s
+; RUN: not wasm-ld --export=missing -o %t.wasm %t.bc 2>&1 | FileCheck -check-prefix=CHECK-ERROR %s
+; RUN: wasm-ld --export=hidden_function -o %t.wasm %t.bc
+; RUN: obj2yaml %t.wasm | FileCheck %s
+
+target datalayout = "e-m:e-p:32:32-i64:64-n32:64-S128"
+target triple = "wasm32-unknown-unknown"
+
+define hidden i32 @hidden_function() local_unnamed_addr {
+entry:
+  ret i32 0
+}
+
+define void @_start() local_unnamed_addr {
+entry:
+  ret void
+}
+
+; CHECK-ERROR: error: symbol exported via --export not found: missing
+
+; CHECK:        - Type:            EXPORT
+; CHECK-NEXT:     Exports:
+; CHECK-NEXT:       - Name:            memory
+; CHECK-NEXT:         Kind:            MEMORY
+; CHECK-NEXT:         Index:           0
+; CHECK-NEXT:       - Name:            __heap_base
+; CHECK-NEXT:         Kind:            GLOBAL
+; CHECK-NEXT:         Index:           1
+; CHECK-NEXT:       - Name:            __data_end
+; CHECK-NEXT:         Kind:            GLOBAL
+; CHECK-NEXT:         Index:           2
+; CHECK-NEXT:       - Name:            _start
+; CHECK-NEXT:         Kind:            FUNCTION
+; CHECK-NEXT:         Index:           2
+; CHECK-NEXT:       - Name:            hidden_function
+; CHECK-NEXT:         Kind:            FUNCTION
+; CHECK-NEXT:         Index:           1
+; CHECK-NEXT:   - Type:            CODE
diff --git a/lld/wasm/Driver.cpp b/lld/wasm/Driver.cpp
index f77f20f4232..06aa437d869 100644
--- a/lld/wasm/Driver.cpp
+++ b/lld/wasm/Driver.cpp
@@ -439,6 +439,16 @@ void LinkerDriver::link(ArrayRef<const char *> ArgsArr) {
   if (!Config->Relocatable)
     handleWeakUndefines();
 
+  // Handle --export.
+  for (auto *Arg : Args.filtered(OPT_export)) {
+    StringRef Name = Arg->getValue();
+    Symbol *Sym = Symtab->find(Name);
+    if (Sym && Sym->isDefined())
+      Sym->ForceExport = true;
+    else if (!Config->AllowUndefined)
+      error("symbol exported via --export not found: " + Name);
+  }
+
   // Do link-time optimization if given files are LLVM bitcode files.
   // This compiles bitcode files into real object files.
   Symtab->addCombinedLTOObject();
@@ -466,16 +476,6 @@ void LinkerDriver::link(ArrayRef<const char *> ArgsArr) {
   if (errorCount())
     return;
 
-  // Handle --export.
-  for (auto *Arg : Args.filtered(OPT_export)) {
-    StringRef Name = Arg->getValue();
-    Symbol *Sym = Symtab->find(Name);
-    if (Sym && Sym->isDefined())
-      Sym->setHidden(false);
-    else if (!Config->AllowUndefined)
-      error("symbol exported via --export not found: " + Name);
-  }
-
   if (EntrySym)
     EntrySym->setHidden(false);
 
diff --git a/lld/wasm/Symbols.cpp b/lld/wasm/Symbols.cpp
index ef135b35b1a..a11081cbcf7 100644
--- a/lld/wasm/Symbols.cpp
+++ b/lld/wasm/Symbols.cpp
@@ -102,7 +102,7 @@ bool Symbol::isExported() const {
   if (!isDefined() || isLocal())
     return false;
 
-  if (Config->ExportAll)
+  if (ForceExport || Config->ExportAll)
     return true;
 
   return !isHidden();
diff --git a/lld/wasm/Symbols.h b/lld/wasm/Symbols.h
index 794a65df673..815cc97d22d 100644
--- a/lld/wasm/Symbols.h
+++ b/lld/wasm/Symbols.h
@@ -94,11 +94,12 @@ public:
 
   // True if this symbol was referenced by a regular (non-bitcode) object.
   unsigned IsUsedInRegularObj : 1;
+  unsigned ForceExport : 1;
 
 protected:
   Symbol(StringRef Name, Kind K, uint32_t Flags, InputFile *F)
-      : IsUsedInRegularObj(false), Name(Name), SymbolKind(K), Flags(Flags),
-        File(F), Referenced(!Config->GcSections) {}
+      : IsUsedInRegularObj(false), ForceExport(false), Name(Name),
+        SymbolKind(K), Flags(Flags), File(F), Referenced(!Config->GcSections) {}
 
   StringRef Name;
   Kind SymbolKind;
@@ -339,6 +340,7 @@ T *replaceSymbol(Symbol *S, ArgT &&... Arg) {
 
   T *S2 = new (S) T(std::forward<ArgT>(Arg)...);
   S2->IsUsedInRegularObj = SymCopy.IsUsedInRegularObj;
+  S2->ForceExport = SymCopy.ForceExport;
   return S2;
 }