[WebAssembly] Fix for discarded init functions

When a function is excluded via comdat we shouldn't add it to the
final list of init functions.

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

llvm-svn: 362769

6 files changed, 46 insertions, 17 deletions

diff --git a/lld/wasm/InputChunks.h b/lld/wasm/InputChunks.h
index f0e762a22c8..ba2bb6f0e54 100644
--- a/lld/wasm/InputChunks.h
+++ b/lld/wasm/InputChunks.h
@@ -64,9 +64,12 @@ public:
   // If GC is disabled, all sections start out as live by default.
   unsigned Live : 1;
 
+  // Signals the chunk was discarded by COMDAT handling.
+  unsigned Discarded : 1;
+
 protected:
   InputChunk(ObjFile *F, Kind K)
-      : File(F), Live(!Config->GcSections), SectionKind(K) {}
+      : File(F), Live(!Config->GcSections), Discarded(false), SectionKind(K) {}
   virtual ~InputChunk() = default;
   virtual ArrayRef<uint8_t> data() const = 0;
 
diff --git a/lld/wasm/InputFiles.cpp b/lld/wasm/InputFiles.cpp
index 1d49f6385ba..10a5e22c720 100644
--- a/lld/wasm/InputFiles.cpp
+++ b/lld/wasm/InputFiles.cpp
@@ -306,14 +306,17 @@ void ObjFile::parse(bool IgnoreComdats) {
   TypeIsUsed.resize(getWasmObj()->types().size(), false);
 
   ArrayRef<StringRef> Comdats = WasmObj->linkingData().Comdats;
-  for (unsigned I = 0; I < Comdats.size(); ++I) {
-    bool IsNew = IgnoreComdats || Symtab->addComdat(Comdats[I]);
+  for (StringRef Comdat : Comdats) {
+    bool IsNew = IgnoreComdats || Symtab->addComdat(Comdat);
     KeptComdats.push_back(IsNew);
   }
 
   // Populate `Segments`.
-  for (const WasmSegment &S : WasmObj->dataSegments())
-    Segments.emplace_back(make<InputSegment>(S, this));
+  for (const WasmSegment &S : WasmObj->dataSegments()) {
+    auto* Seg = make<InputSegment>(S, this);
+    Seg->Discarded = isExcludedByComdat(Seg);
+    Segments.emplace_back(Seg);
+  }
   setRelocs(Segments, DataSection);
 
   // Populate `Functions`.
@@ -322,9 +325,11 @@ void ObjFile::parse(bool IgnoreComdats) {
   ArrayRef<WasmSignature> Types = WasmObj->types();
   Functions.reserve(Funcs.size());
 
-  for (size_t I = 0, E = Funcs.size(); I != E; ++I)
-    Functions.emplace_back(
-        make<InputFunction>(Types[FuncTypes[I]], &Funcs[I], this));
+  for (size_t I = 0, E = Funcs.size(); I != E; ++I) {
+    auto* Func = make<InputFunction>(Types[FuncTypes[I]], &Funcs[I], this);
+    Func->Discarded = isExcludedByComdat(Func);
+    Functions.emplace_back(Func);
+  }
   setRelocs(Functions, CodeSection);
 
   // Populate `Globals`.
@@ -387,21 +392,16 @@ Symbol *ObjFile::createDefined(const WasmSymbol &Sym) {
   case WASM_SYMBOL_TYPE_FUNCTION: {
     InputFunction *Func =
         Functions[Sym.Info.ElementIndex - WasmObj->getNumImportedFunctions()];
-    if (isExcludedByComdat(Func)) {
-      Func->Live = false;
+    if (Func->Discarded)
       return nullptr;
-    }
-
     if (Sym.isBindingLocal())
       return make<DefinedFunction>(Name, Flags, this, Func);
     return Symtab->addDefinedFunction(Name, Flags, this, Func);
   }
   case WASM_SYMBOL_TYPE_DATA: {
     InputSegment *Seg = Segments[Sym.Info.DataRef.Segment];
-    if (isExcludedByComdat(Seg)) {
-      Seg->Live = false;
+    if (Seg->Discarded)
       return nullptr;
-    }
 
     uint32_t Offset = Sym.Info.DataRef.Offset;
     uint32_t Size = Sym.Info.DataRef.Size;
@@ -439,12 +439,22 @@ Symbol *ObjFile::createUndefined(const WasmSymbol &Sym, bool IsCalledDirectly) {
 
   switch (Sym.Info.Kind) {
   case WASM_SYMBOL_TYPE_FUNCTION:
+    if (Sym.isBindingLocal())
+      return make<UndefinedFunction>(Name, Sym.Info.ImportName,
+                                     Sym.Info.ImportModule, Flags, this,
+                                     Sym.Signature, IsCalledDirectly);
     return Symtab->addUndefinedFunction(Name, Sym.Info.ImportName,
                                         Sym.Info.ImportModule, Flags, this,
                                         Sym.Signature, IsCalledDirectly);
   case WASM_SYMBOL_TYPE_DATA:
+    if (Sym.isBindingLocal())
+      return make<UndefinedData>(Name, Flags, this);
     return Symtab->addUndefinedData(Name, Flags, this);
   case WASM_SYMBOL_TYPE_GLOBAL:
+    if (Sym.isBindingLocal())
+      return make<UndefinedGlobal>(Name, Sym.Info.ImportName,
+                                   Sym.Info.ImportModule, Flags, this,
+                                   Sym.GlobalType);
     return Symtab->addUndefinedGlobal(Name, Sym.Info.ImportName,
                                       Sym.Info.ImportModule, Flags, this,
                                       Sym.GlobalType);
diff --git a/lld/wasm/MarkLive.cpp b/lld/wasm/MarkLive.cpp
index a53c4727311..b3dce23bcaa 100644
--- a/lld/wasm/MarkLive.cpp
+++ b/lld/wasm/MarkLive.cpp
@@ -52,8 +52,11 @@ void lld::wasm::markLive() {
     if (Sym == WasmSym::CallCtors) {
       for (const ObjFile *Obj : Symtab->ObjectFiles) {
         const WasmLinkingData &L = Obj->getWasmObj()->linkingData();
-        for (const WasmInitFunc &F : L.InitFunctions)
-          Enqueue(Obj->getFunctionSymbol(F.Symbol));
+        for (const WasmInitFunc &F : L.InitFunctions) {
+          auto* InitSym = Obj->getFunctionSymbol(F.Symbol);
+          if (!InitSym->isDiscarded())
+            Enqueue(InitSym);
+        }
       }
     }
   };
diff --git a/lld/wasm/Symbols.cpp b/lld/wasm/Symbols.cpp
index 07b895a2e67..284fafddfe5 100644
--- a/lld/wasm/Symbols.cpp
+++ b/lld/wasm/Symbols.cpp
@@ -63,6 +63,12 @@ InputChunk *Symbol::getChunk() const {
   return nullptr;
 }
 
+bool Symbol::isDiscarded() const {
+  if (InputChunk *C = getChunk())
+    return C->Discarded;
+  return false;
+}
+
 bool Symbol::isLive() const {
   if (auto *G = dyn_cast<DefinedGlobal>(this))
     return G->Global->Live;
@@ -74,6 +80,7 @@ bool Symbol::isLive() const {
 }
 
 void Symbol::markLive() {
+  assert(!isDiscarded());
   if (auto *G = dyn_cast<DefinedGlobal>(this))
     G->Global->Live = true;
   if (auto *E = dyn_cast<DefinedEvent>(this))
diff --git a/lld/wasm/Symbols.h b/lld/wasm/Symbols.h
index 98ff155fc5a..5c0f4b2ac46 100644
--- a/lld/wasm/Symbols.h
+++ b/lld/wasm/Symbols.h
@@ -69,6 +69,9 @@ public:
   bool isWeak() const;
   bool isHidden() const;
 
+  // Returns true if this symbol exists in a discarded (due to COMDAT) section
+  bool isDiscarded() const;
+
   // True if this is an undefined weak symbol. This only works once
   // all input files have been added.
   bool isUndefWeak() const {
diff --git a/lld/wasm/Writer.cpp b/lld/wasm/Writer.cpp
index 0aa1180cb2b..43a8c8de637 100644
--- a/lld/wasm/Writer.cpp
+++ b/lld/wasm/Writer.cpp
@@ -701,6 +701,9 @@ void Writer::calculateInitFunctions() {
     const WasmLinkingData &L = File->getWasmObj()->linkingData();
     for (const WasmInitFunc &F : L.InitFunctions) {
       FunctionSymbol *Sym = File->getFunctionSymbol(F.Symbol);
+      // comdat exclusions can cause init functions be discarded.
+      if (Sym->isDiscarded())
+        continue;
       assert(Sym->isLive());
       if (*Sym->Signature != WasmSignature{{}, {}})
         error("invalid signature for init func: " + toString(*Sym));