[WebAssembly] Add COMDAT support

This adds COMDAT support to the Wasm object-file format.
Spec: https://github.com/WebAssembly/tool-conventions/pull/31

Corresponding LLD change:
https://bugs.llvm.org/show_bug.cgi?id=35533, and D40845

Patch by Nicholas Wilson

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

llvm-svn: 322135

4 files changed, 141 insertions, 14 deletions

diff --git a/llvm/lib/CodeGen/TargetLoweringObjectFileImpl.cpp b/llvm/lib/CodeGen/TargetLoweringObjectFileImpl.cpp
index 24d4baa31e1..bdc4bc0302b 100644
--- a/llvm/lib/CodeGen/TargetLoweringObjectFileImpl.cpp
+++ b/llvm/lib/CodeGen/TargetLoweringObjectFileImpl.cpp
@@ -1254,15 +1254,17 @@ void TargetLoweringObjectFileCOFF::emitLinkerFlagsForGlobal(
 //                                  Wasm
 //===----------------------------------------------------------------------===//
 
-static void checkWasmComdat(const GlobalValue *GV) {
+static const Comdat *getWasmComdat(const GlobalValue *GV) {
   const Comdat *C = GV->getComdat();
   if (!C)
-    return;
+    return nullptr;
 
-  // TODO(sbc): At some point we may need COMDAT support but currently
-  // they are not supported.
-  report_fatal_error("WebAssembly doesn't support COMDATs, '" + C->getName() +
-                     "' cannot be lowered.");
+  if (C->getSelectionKind() != Comdat::Any)
+    report_fatal_error("WebAssembly COMDATs only support "
+                       "SelectionKind::Any, '" + C->getName() + "' cannot be "
+                       "lowered.");
+
+  return C;
 }
 
 static SectionKind getWasmKindForNamedSection(StringRef Name, SectionKind K) {
@@ -1278,16 +1280,25 @@ static SectionKind getWasmKindForNamedSection(StringRef Name, SectionKind K) {
 MCSection *TargetLoweringObjectFileWasm::getExplicitSectionGlobal(
     const GlobalObject *GO, SectionKind Kind, const TargetMachine &TM) const {
   StringRef Name = GO->getSection();
-  checkWasmComdat(GO);
+
   Kind = getWasmKindForNamedSection(Name, Kind);
-  return getContext().getWasmSection(Name, Kind);
+
+  StringRef Group = "";
+  if (const Comdat *C = getWasmComdat(GO)) {
+    Group = C->getName();
+  }
+
+  return getContext().getWasmSection(Name, Kind, Group,
+                                     MCContext::GenericSectionID);
 }
 
 static MCSectionWasm *selectWasmSectionForGlobal(
     MCContext &Ctx, const GlobalObject *GO, SectionKind Kind, Mangler &Mang,
     const TargetMachine &TM, bool EmitUniqueSection, unsigned *NextUniqueID) {
   StringRef Group = "";
-  checkWasmComdat(GO);
+  if (const Comdat *C = getWasmComdat(GO)) {
+    Group = C->getName();
+  }
 
   bool UniqueSectionNames = TM.getUniqueSectionNames();
   SmallString<128> Name = getSectionPrefixForGlobal(Kind);
diff --git a/llvm/lib/MC/WasmObjectWriter.cpp b/llvm/lib/MC/WasmObjectWriter.cpp
index 37b12c63be4..e5099a83b6f 100644
--- a/llvm/lib/MC/WasmObjectWriter.cpp
+++ b/llvm/lib/MC/WasmObjectWriter.cpp
@@ -138,6 +138,14 @@ struct WasmGlobal {
   uint32_t ImportIndex;
 };
 
+// Information about a single item which is part of a COMDAT.  For each data
+// segment or function which is in the COMDAT, there is a corresponding
+// WasmComdatEntry.
+struct WasmComdatEntry {
+  unsigned Kind;
+  uint32_t Index;
+};
+
 // Information about a single relocation.
 struct WasmRelocationEntry {
   uint64_t Offset;                  // Where is the relocation.
@@ -284,8 +292,9 @@ private:
   void writeDataRelocSection();
   void writeLinkingMetaDataSection(
       ArrayRef<WasmDataSegment> Segments, uint32_t DataSize,
-      const SmallVector<std::pair<StringRef, uint32_t>, 4> &SymbolFlags,
-      const SmallVector<std::pair<uint16_t, uint32_t>, 2> &InitFuncs);
+      ArrayRef<std::pair<StringRef, uint32_t>> SymbolFlags,
+      ArrayRef<std::pair<uint16_t, uint32_t>> InitFuncs,
+      const std::map<StringRef, std::vector<WasmComdatEntry>>& Comdats);
 
   uint32_t getProvisionalValue(const WasmRelocationEntry &RelEntry);
   void applyRelocations(ArrayRef<WasmRelocationEntry> Relocations,
@@ -913,8 +922,9 @@ void WasmObjectWriter::writeDataRelocSection() {
 
 void WasmObjectWriter::writeLinkingMetaDataSection(
     ArrayRef<WasmDataSegment> Segments, uint32_t DataSize,
-    const SmallVector<std::pair<StringRef, uint32_t>, 4> &SymbolFlags,
-    const SmallVector<std::pair<uint16_t, uint32_t>, 2> &InitFuncs) {
+    ArrayRef<std::pair<StringRef, uint32_t>> SymbolFlags,
+    ArrayRef<std::pair<uint16_t, uint32_t>> InitFuncs,
+    const std::map<StringRef, std::vector<WasmComdatEntry>>& Comdats) {
   SectionBookkeeping Section;
   startSection(Section, wasm::WASM_SEC_CUSTOM, "linking");
   SectionBookkeeping SubSection;
@@ -956,6 +966,21 @@ void WasmObjectWriter::writeLinkingMetaDataSection(
     endSection(SubSection);
   }
 
+  if (Comdats.size()) {
+    startSection(SubSection, wasm::WASM_COMDAT_INFO);
+    encodeULEB128(Comdats.size(), getStream());
+    for (const auto &C : Comdats) {
+      writeString(C.first);
+      encodeULEB128(0, getStream()); // flags for future use
+      encodeULEB128(C.second.size(), getStream());
+      for (const WasmComdatEntry &Entry : C.second) {
+        encodeULEB128(Entry.Kind, getStream());
+        encodeULEB128(Entry.Index, getStream());
+      }
+    }
+    endSection(SubSection);
+  }
+
   endSection(Section);
 }
 
@@ -997,6 +1022,7 @@ void WasmObjectWriter::writeObject(MCAssembler &Asm,
   SmallVector<WasmExport, 4> Exports;
   SmallVector<std::pair<StringRef, uint32_t>, 4> SymbolFlags;
   SmallVector<std::pair<uint16_t, uint32_t>, 2> InitFuncs;
+  std::map<StringRef, std::vector<WasmComdatEntry>> Comdats;
   unsigned NumFuncImports = 0;
   SmallVector<WasmDataSegment, 4> DataSegments;
   uint32_t DataSize = 0;
@@ -1143,6 +1169,12 @@ void WasmObjectWriter::writeObject(MCAssembler &Asm,
     Segment.Flags = 0;
     DataSize += Segment.Data.size();
     Section.setMemoryOffset(Segment.Offset);
+
+    if (const MCSymbolWasm *C = Section.getGroup()) {
+      Comdats[C->getName()].emplace_back(
+          WasmComdatEntry{wasm::WASM_COMDAT_DATA,
+                          static_cast<uint32_t>(DataSegments.size()) - 1});
+    }
   }
 
   // Handle regular defined and undefined symbols.
@@ -1216,6 +1248,7 @@ void WasmObjectWriter::writeObject(MCAssembler &Asm,
       // address.  For externals these will also be named exports.
       Index = NumGlobalImports + Globals.size();
       auto &DataSection = static_cast<MCSectionWasm &>(WS.getSection());
+      assert(DataSection.isWasmData());
 
       WasmGlobal Global;
       Global.Type = PtrType;
@@ -1239,8 +1272,16 @@ void WasmObjectWriter::writeObject(MCAssembler &Asm,
         Export.Kind = wasm::WASM_EXTERNAL_GLOBAL;
       DEBUG(dbgs() << "  -> export " << Exports.size() << "\n");
       Exports.push_back(Export);
+
       if (!WS.isExternal())
         SymbolFlags.emplace_back(WS.getName(), wasm::WASM_SYMBOL_BINDING_LOCAL);
+
+      if (WS.isFunction()) {
+        auto &Section = static_cast<MCSectionWasm &>(WS.getSection(false));
+        if (const MCSymbolWasm *C = Section.getGroup())
+          Comdats[C->getName()].emplace_back(
+              WasmComdatEntry{wasm::WASM_COMDAT_FUNCTION, Index});
+      }
     }
   }
 
@@ -1372,7 +1413,7 @@ void WasmObjectWriter::writeObject(MCAssembler &Asm,
   writeCodeRelocSection();
   writeDataRelocSection();
   writeLinkingMetaDataSection(DataSegments, DataSize, SymbolFlags,
-                              InitFuncs);
+                              InitFuncs, Comdats);
 
   // TODO: Translate the .comment section to the output.
   // TODO: Translate debug sections to the output.
diff --git a/llvm/lib/Object/WasmObjectFile.cpp b/llvm/lib/Object/WasmObjectFile.cpp
index c33f4f68694..ec2e2bfbb4d 100644
--- a/llvm/lib/Object/WasmObjectFile.cpp
+++ b/llvm/lib/Object/WasmObjectFile.cpp
@@ -10,6 +10,7 @@
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/BinaryFormat/Wasm.h"
 #include "llvm/MC/SubtargetFeature.h"
@@ -422,6 +423,10 @@ Error WasmObjectFile::parseLinkingSection(const uint8_t *Ptr,
       }
       break;
     }
+    case wasm::WASM_COMDAT_INFO:
+      if (Error Err = parseLinkingSectionComdat(Ptr, SubSectionEnd))
+        return Err;
+      break;
     default:
       Ptr += Size;
       break;
@@ -436,6 +441,55 @@ Error WasmObjectFile::parseLinkingSection(const uint8_t *Ptr,
   return Error::success();
 }
 
+Error WasmObjectFile::parseLinkingSectionComdat(const uint8_t *&Ptr,
+                                                const uint8_t *End)
+{
+  uint32_t ComdatCount = readVaruint32(Ptr);
+  StringSet<> ComdatSet;
+  while (ComdatCount--) {
+    StringRef Name = readString(Ptr);
+    if (Name.empty() || !ComdatSet.insert(Name).second)
+      return make_error<GenericBinaryError>("Bad/duplicate COMDAT name " + Twine(Name),
+                                            object_error::parse_failed);
+    Comdats.emplace_back(Name);
+    uint32_t Flags = readVaruint32(Ptr);
+    if (Flags != 0)
+      return make_error<GenericBinaryError>("Unsupported COMDAT flags",
+                                            object_error::parse_failed);
+
+    uint32_t EntryCount = readVaruint32(Ptr);
+    while (EntryCount--) {
+      unsigned Kind = readVaruint32(Ptr);
+      unsigned Index = readVaruint32(Ptr);
+      switch (Kind) {
+      default:
+        return make_error<GenericBinaryError>("Invalid COMDAT entry type",
+                                              object_error::parse_failed);
+      case wasm::WASM_COMDAT_DATA:
+        if (Index >= DataSegments.size())
+          return make_error<GenericBinaryError>("COMDAT data index out of range",
+                                                object_error::parse_failed);
+        if (!DataSegments[Index].Data.Comdat.empty())
+          return make_error<GenericBinaryError>("Data segment in two COMDATs",
+                                                object_error::parse_failed);
+        DataSegments[Index].Data.Comdat = Name;
+        break;
+      case wasm::WASM_COMDAT_FUNCTION:
+        if (Index < NumImportedFunctions || !isValidFunctionIndex(Index))
+          return make_error<GenericBinaryError>("COMDAT function index out of range",
+                                                object_error::parse_failed);
+        Index -= NumImportedFunctions;
+        if (!Functions[Index].Comdat.empty())
+          return make_error<GenericBinaryError>("Function in two COMDATs",
+                                                object_error::parse_failed);
+        Functions[Index].Comdat = Name;
+        break;
+      }
+    }
+  }
+  return Error::success();
+}
+
 WasmSection* WasmObjectFile::findCustomSectionByName(StringRef Name) {
   for (WasmSection& Section : Sections) {
     if (Section.Type == wasm::WASM_SEC_CUSTOM && Section.Name == Name)
diff --git a/llvm/lib/ObjectYAML/WasmYAML.cpp b/llvm/lib/ObjectYAML/WasmYAML.cpp
index 70721e006da..4ae6dccccb1 100644
--- a/llvm/lib/ObjectYAML/WasmYAML.cpp
+++ b/llvm/lib/ObjectYAML/WasmYAML.cpp
@@ -61,6 +61,7 @@ static void sectionMapping(IO &IO, WasmYAML::LinkingSection &Section) {
   IO.mapOptional("SymbolInfo", Section.SymbolInfos);
   IO.mapOptional("SegmentInfo", Section.SegmentInfos);
   IO.mapOptional("InitFunctions", Section.InitFunctions);
+  IO.mapOptional("Comdats", Section.Comdats);
 }
 
 static void sectionMapping(IO &IO, WasmYAML::CustomSection &Section) {
@@ -368,6 +369,26 @@ void MappingTraits<WasmYAML::InitFunction>::mapping(
   IO.mapRequired("FunctionIndex", Init.FunctionIndex);
 }
 
+void ScalarEnumerationTraits<WasmYAML::ComdatKind>::enumeration(
+    IO &IO, WasmYAML::ComdatKind &Kind) {
+#define ECase(X) IO.enumCase(Kind, #X, wasm::WASM_COMDAT_##X);
+  ECase(FUNCTION);
+  ECase(DATA);
+#undef ECase
+}
+
+void MappingTraits<WasmYAML::ComdatEntry>::mapping(
+    IO &IO, WasmYAML::ComdatEntry &ComdatEntry) {
+  IO.mapRequired("Kind", ComdatEntry.Kind);
+  IO.mapRequired("Index", ComdatEntry.Index);
+}
+
+void MappingTraits<WasmYAML::Comdat>::mapping(
+    IO &IO, WasmYAML::Comdat &Comdat) {
+  IO.mapRequired("Name", Comdat.Name);
+  IO.mapRequired("Entries", Comdat.Entries);
+}
+
 void MappingTraits<WasmYAML::SymbolInfo>::mapping(IO &IO,
                                                   WasmYAML::SymbolInfo &Info) {
   IO.mapRequired("Name", Info.Name);