[WebAssembly] Add support for custom sections

Copy user-defined custom sections into the output, concatenating
sections with the same name.

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

llvm-svn: 329717

9 files changed, 140 insertions, 1 deletions

diff --git a/lld/test/wasm/Inputs/custom.ll b/lld/test/wasm/Inputs/custom.ll
new file mode 100644
index 00000000000..7a20e2f9b0f
--- /dev/null
+++ b/lld/test/wasm/Inputs/custom.ll
@@ -0,0 +1,6 @@
+target triple = "wasm32-unknown-unknown-wasm"
+
+!0 = !{ !"red", !"foo" }
+!1 = !{ !"green", !"bar" }
+!2 = !{ !"green", !"qux" }
+!wasm.custom_sections = !{ !0, !1, !2 }
diff --git a/lld/test/wasm/custom-sections.ll b/lld/test/wasm/custom-sections.ll
new file mode 100644
index 00000000000..0b0bd5b4700
--- /dev/null
+++ b/lld/test/wasm/custom-sections.ll
@@ -0,0 +1,22 @@
+; RUN: llc -filetype=obj %s -o %t1.o
+; RUN: llc -filetype=obj  %S/Inputs/custom.ll -o %t2.o
+; RUN: wasm-ld --check-signatures --relocatable -o %t.wasm %t1.o %t2.o
+; RUN: obj2yaml %t.wasm | FileCheck %s
+
+target triple = "wasm32-unknown-unknown-wasm"
+
+define i32 @_start() local_unnamed_addr {
+entry:
+  %retval = alloca i32, align 4
+  ret i32 0
+}
+
+!0 = !{ !"red", !"extra" }
+!wasm.custom_sections = !{ !0 }
+
+; CHECK:        - Type:            CUSTOM
+; CHECK-NEXT:     Name:            green
+; CHECK-NEXT:     Payload:         '05677265656E626172717578'
+; CHECK-NEXT:   - Type:            CUSTOM
+; CHECK-NEXT:     Name:            red
+; CHECK-NEXT:     Payload:         037265646578747261666F6F
diff --git a/lld/wasm/InputChunks.cpp b/lld/wasm/InputChunks.cpp
index 0cac19b3d4e..a1b3426888f 100644
--- a/lld/wasm/InputChunks.cpp
+++ b/lld/wasm/InputChunks.cpp
@@ -143,3 +143,13 @@ void InputFunction::setTableIndex(uint32_t Index) {
   assert(!hasTableIndex());
   TableIndex = Index;
 }
+
+InputSection::InputSection(const WasmSection &S, ObjFile *F)
+    : InputChunk(F, InputChunk::Section), Section(S) {
+  assert(Section.Type == llvm::wasm::WASM_SEC_CUSTOM);
+  // TODO check LEB errors
+  unsigned Count;
+  uint64_t NameSize = llvm::decodeULEB128(Section.Content.data(), &Count);
+  uint32_t PayloadOffset = Count + NameSize;
+  Payload = Section.Content.slice(PayloadOffset);
+}
diff --git a/lld/wasm/InputChunks.h b/lld/wasm/InputChunks.h
index 3d6ea23680f..5b9b885f3c9 100644
--- a/lld/wasm/InputChunks.h
+++ b/lld/wasm/InputChunks.h
@@ -44,7 +44,7 @@ class OutputSegment;
 
 class InputChunk {
 public:
-  enum Kind { DataSegment, Function, SyntheticFunction };
+  enum Kind { DataSegment, Function, SyntheticFunction, Section };
 
   Kind kind() const { return SectionKind; }
 
@@ -173,6 +173,25 @@ protected:
   ArrayRef<uint8_t> Body;
 };
 
+// Represents a single Wasm Section within an input file.
+class InputSection : public InputChunk {
+public:
+  InputSection(const WasmSection &S, ObjFile *F);
+
+  StringRef getName() const override { return Section.Name; }
+  uint32_t getComdat() const override { return UINT32_MAX; }
+
+protected:
+  ArrayRef<uint8_t> data() const override { return Payload; }
+
+  // Offset within the input section.  This is only zero since this chunk
+  // type represents an entire input section, not part of one.
+  uint32_t getInputSectionOffset() const override { return 0; }
+
+  const WasmSection &Section;
+  ArrayRef<uint8_t> Payload;
+};
+
 } // namespace wasm
 
 std::string toString(const wasm::InputChunk *);
diff --git a/lld/wasm/InputFiles.cpp b/lld/wasm/InputFiles.cpp
index e7436863461..40a18ff6cf1 100644
--- a/lld/wasm/InputFiles.cpp
+++ b/lld/wasm/InputFiles.cpp
@@ -153,6 +153,8 @@ void ObjFile::parse() {
       CodeSection = &Section;
     else if (Section.Type == WASM_SEC_DATA)
       DataSection = &Section;
+    else if (Section.Type == WASM_SEC_CUSTOM)
+      CustomSections.emplace_back(make<InputSection>(Section, this));
   }
 
   TypeMap.resize(getWasmObj()->types().size());
diff --git a/lld/wasm/InputFiles.h b/lld/wasm/InputFiles.h
index f0fd267d910..8343fb98cb0 100644
--- a/lld/wasm/InputFiles.h
+++ b/lld/wasm/InputFiles.h
@@ -35,6 +35,7 @@ class InputChunk;
 class InputFunction;
 class InputSegment;
 class InputGlobal;
+class InputSection;
 
 class InputFile {
 public:
@@ -108,6 +109,7 @@ public:
   std::vector<InputSegment *> Segments;
   std::vector<InputFunction *> Functions;
   std::vector<InputGlobal *> Globals;
+  std::vector<InputSection *> CustomSections;
 
   ArrayRef<Symbol *> getSymbols() const { return Symbols; }
   Symbol *getSymbol(uint32_t Index) const { return Symbols[Index]; }
diff --git a/lld/wasm/OutputSections.cpp b/lld/wasm/OutputSections.cpp
index 29b1c43f06c..cf3d8b332dd 100644
--- a/lld/wasm/OutputSections.cpp
+++ b/lld/wasm/OutputSections.cpp
@@ -189,3 +189,38 @@ void DataSection::writeRelocations(raw_ostream &OS) const {
     for (const InputChunk *C : Seg->InputSegments)
       C->writeRelocations(OS);
 }
+
+CustomSection::CustomSection(std::string Name,
+                             ArrayRef<InputSection *> InputSections)
+    : OutputSection(WASM_SEC_CUSTOM, Name), PayloadSize(0),
+      InputSections(InputSections) {
+  raw_string_ostream OS(NameData);
+  encodeULEB128(Name.size(), OS);
+  OS << Name;
+  OS.flush();
+
+  for (InputSection *Section : InputSections) {
+    Section->OutputOffset = PayloadSize;
+    PayloadSize += Section->getSize();
+  }
+
+  createHeader(PayloadSize + NameData.size());
+}
+
+void CustomSection::writeTo(uint8_t *Buf) {
+  log("writing " + toString(*this) + " size=" + Twine(getSize()) +
+      " chunks=" + Twine(InputSections.size()));
+
+  assert(Offset);
+  Buf += Offset;
+
+  // Write section header
+  memcpy(Buf, Header.data(), Header.size());
+  Buf += Header.size();
+  memcpy(Buf, NameData.data(), NameData.size());
+  Buf += NameData.size();
+
+  // Write custom sections payload
+  parallelForEach(InputSections,
+                  [&](const InputSection *Section) { Section->writeTo(Buf); });
+}
diff --git a/lld/wasm/OutputSections.h b/lld/wasm/OutputSections.h
index 2ecee951757..3cb7be1653c 100644
--- a/lld/wasm/OutputSections.h
+++ b/lld/wasm/OutputSections.h
@@ -113,6 +113,27 @@ protected:
   size_t BodySize = 0;
 };
 
+// Represents a custom section in the output file.  Wasm custom sections are 
+// used for storing user-defined metadata.  Unlike the core sections types
+// they are identified by their string name.
+// The linker combines custom sections that have the same name by simply
+// concatenating them.
+// Note that some custom sections such as "name" and "linking" are handled
+// separately and are instead synthesized by the linker.
+class CustomSection : public OutputSection {
+public:
+  CustomSection(std::string Name, ArrayRef<InputSection *> InputSections);
+  size_t getSize() const override {
+    return Header.size() + NameData.size() + PayloadSize;
+  }
+  void writeTo(uint8_t *Buf) override;
+
+protected:
+  size_t PayloadSize;
+  ArrayRef<InputSection *> InputSections;
+  std::string NameData;
+};
+
 } // namespace wasm
 } // namespace lld
 
diff --git a/lld/wasm/Writer.cpp b/lld/wasm/Writer.cpp
index a98cfc6b7c9..18d5d4a8d52 100644
--- a/lld/wasm/Writer.cpp
+++ b/lld/wasm/Writer.cpp
@@ -20,6 +20,7 @@
 #include "lld/Common/Strings.h"
 #include "lld/Common/Threads.h"
 #include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/StringMap.h"
 #include "llvm/BinaryFormat/Wasm.h"
 #include "llvm/Object/WasmTraits.h"
 #include "llvm/Support/FileOutputBuffer.h"
@@ -85,6 +86,7 @@ private:
   void createElemSection();
   void createCodeSection();
   void createDataSection();
+  void createCustomSections();
 
   // Custom sections
   void createRelocSections();
@@ -111,6 +113,8 @@ private:
   std::vector<const Symbol *> SymtabEntries;
   std::vector<WasmInitEntry> InitFunctions;
 
+  llvm::StringMap<std::vector<InputSection *>> CustomSectionMapping;
+
   // Elements that are used to construct the final output
   std::string Header;
   std::vector<OutputSection *> OutputSections;
@@ -295,6 +299,23 @@ void Writer::createExportSection() {
   }
 }
 
+void Writer::createCustomSections() {
+  log("createCustomSections");
+  for (ObjFile *File : Symtab->ObjectFiles)
+    for (InputSection *Section : File->CustomSections)
+      CustomSectionMapping[Section->getName()].push_back(Section);
+
+  for (auto &Pair : CustomSectionMapping) {
+    StringRef Name = Pair.first();
+    // These custom sections are known the linker and synthesized rather than
+    // blindly copied
+    if (Name == "linking" || Name == "name" || Name.startswith("reloc."))
+      continue;
+    DEBUG(dbgs() << "createCustomSection: " << Name << "\n");
+    OutputSections.push_back(make<CustomSection>(Name, Pair.second));
+  }
+}
+
 void Writer::createElemSection() {
   if (IndirectFunctions.empty())
     return;
@@ -647,6 +668,7 @@ void Writer::createSections() {
   createElemSection();
   createCodeSection();
   createDataSection();
+  createCustomSections();
 
   // Custom sections
   if (Config->Relocatable) {