[WebAssembly] Add support for user-defined custom sections

[sbc100]

Summary

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

[sbc100]

• cleanup

[ruiu]

I think overall it is looking good.

wasm/InputChunks.h

184

If you move this to .cpp, can you remove #include "LEB128.h"?

Please add a comment what this code is doing.

195

It may make sense to do this in the ctor and store an ArrayRef instead of an size_t offset as a member.

197–199

What is this function? Please add a comment to InputChunk class.

wasm/OutputSections.h

116

Please add a class comment.

118

I wouldn't add explicit because it is hard to accidentally trigger implicit type conversion for this class.

wasm/Writer.cpp

314

Are you sure that these section names don't start with "."?

330

You can eliminate this variable.

[sbc100]

• feedback

[sbc100]

wasm/Writer.cpp

314

Yes. These sections are all created explicitly with these names

[yurydelendik]

Looks good to me (from D45118 point of view)

[sbc100]

@ruiu , good to go now?

[ruiu]

wasm/OutputSections.h

116

Can you expand the comment so that those who don't really now what "custom section" is can understand this comment? Currently, in order to understand this comment, you need to know what custom section is before reading this comment.

You need to explain that in wasm, text, data, and some other sections are special sections that are identified by type, and unlike other file formats, they don't have name. However, you can still add sections that don't have special type and have name just like normal sections in other file formats, and such sections are called "custom sections" in wasm.

wasm/Writer.cpp

304–305

nit: we omit { if the content is only one line.

[sbc100]

• feedback

[ruiu]

LGTM

test/wasm/custom-sections.ll

+; RUN: llc -filetype=obj %s -o %t.o
+; RUN: wasm-ld --check-signatures --relocatable -o %t.wasm %t.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", !"foo" }
+!1 = !{ !"green", !"bar" }
+!2 = !{ !"green", !"qux" }
+!wasm.custom_sections = !{ !0, !1, !2 }
+
+; CHECK:        - Type:            CUSTOM
+; CHECK_NEXT:     Name:            green
+; CHECK_NEXT:     Payload:         '05677265656E626172717578'
+; CHECK_NEXT:   - Type:            CUSTOM
+; CHECK_NEXT:     Name:            red
+; CHECK_NEXT:     Payload:         03726564666F6F

wasm/InputChunks.h

 namespace lld {
 namespace wasm {
 
 class ObjFile;
 class OutputSegment;
 
 class InputChunk {
 public:
-  enum Kind { DataSegment, Function, SyntheticFunction };
+  enum Kind { DataSegment, Function, SyntheticFunction, Section };
 
   Kind kind() const { return SectionKind; }
 
   uint32_t getSize() const { return data().size(); }
 
   void copyRelocations(const WasmSection &Section);
 
   void writeTo(uint8_t *SectionStart) const;

 
 protected:
   ArrayRef<uint8_t> data() const override { return Body; }
 
   StringRef Name;
   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 *);
 } // namespace lld
 
 #endif // LLD_WASM_INPUT_CHUNKS_H

wasm/InputChunks.cpp

 }
 
 void InputFunction::setTableIndex(uint32_t Index) {
   DEBUG(dbgs() << "InputFunction::setTableIndex: " << getName() << " -> "
                << Index << "\n");
   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);
+}

wasm/InputFiles.h

 
 namespace lld {
 namespace wasm {
 
 class InputChunk;
 class InputFunction;
 class InputSegment;
 class InputGlobal;
+class InputSection;
 
 class InputFile {
 public:
   enum Kind {
     ObjectKind,
     ArchiveKind,
   };
 
public:
   std::vector<uint32_t> TypeMap;
   std::vector<bool> TypeIsUsed;
   // Maps function indices to table indices
   std::vector<uint32_t> TableEntries;
   std::vector<bool> UsedComdats;
   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]; }
   FunctionSymbol *getFunctionSymbol(uint32_t Index) const;
   DataSymbol *getDataSymbol(uint32_t Index) const;
   GlobalSymbol *getGlobalSymbol(uint32_t Index) const;
 
 private:

wasm/InputFiles.cpp

void ObjFile::parse() {
   // Find the code and data sections.  Wasm objects can have at most one code
   // and one data section.
   for (const SectionRef &Sec : WasmObj->sections()) {
     const WasmSection &Section = WasmObj->getWasmSection(Sec);
     if (Section.Type == WASM_SEC_CODE)
       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());
   TypeIsUsed.resize(getWasmObj()->types().size(), false);
 
   ArrayRef<StringRef> Comdats = WasmObj->linkingData().Comdats;
   UsedComdats.resize(Comdats.size());
   for (unsigned I = 0; I < Comdats.size(); ++I)

wasm/OutputSections.h

public:
   void writeRelocations(raw_ostream &OS) const override;
 
 protected:
   ArrayRef<OutputSegment *> Segments;
   std::string DataSectionHeader;
   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
 
 #endif // LLD_WASM_OUTPUT_SECTIONS_H

wasm/OutputSections.cpp

uint32_t DataSection::numRelocations() const {
   return Count;
 }
 
 void DataSection::writeRelocations(raw_ostream &OS) const {
   for (const OutputSegment *Seg : Segments)
     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); });
+}

wasm/Writer.cpp

 #include "OutputSegment.h"
 #include "SymbolTable.h"
 #include "WriterUtils.h"
 #include "lld/Common/ErrorHandler.h"
 #include "lld/Common/Memory.h"
 #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"
 #include "llvm/Support/Format.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/LEB128.h"
 
 #include <cstdarg>
private:
   void createTableSection();
   void createGlobalSection();
   void createExportSection();
   void createImportSection();
   void createMemorySection();
   void createElemSection();
   void createCodeSection();
   void createDataSection();
+  void createCustomSections();
 
   // Custom sections
   void createRelocSections();
   void createLinkingSection();
   void createNameSection();
 
   void writeHeader();
   void writeSections();
private:
   std::vector<Symbol *> ExportedSymbols;
   std::vector<const DefinedData *> DefinedFakeGlobals;
   std::vector<InputGlobal *> InputGlobals;
   std::vector<InputFunction *> InputFunctions;
   std::vector<const FunctionSymbol *> IndirectFunctions;
   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;
 
   std::unique_ptr<FileOutputBuffer> Buffer;
 
   std::vector<OutputSegment *> Segments;
   llvm::SmallDenseMap<StringRef, OutputSegment *> SegmentMap;
for (const Symbol *Sym : ExportedSymbols) {
     else if (isa<DefinedData>(Sym))
       Export = {Name, WASM_EXTERNAL_GLOBAL, FakeGlobalIndex++};
     else
       llvm_unreachable("unexpected symbol type");
     writeExport(OS, Export);
   }
 }
 
+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;
 
   SyntheticSection *Section = createSyntheticSection(WASM_SEC_ELEM);
   raw_ostream &OS = Section->getStream();
 
   writeUleb128(OS, 1, "segment count");
   writeUleb128(OS, 0, "table index");
   WasmInitExpr InitExpr;
   InitExpr.Opcode = WASM_OPCODE_I32_CONST;
   InitExpr.Value.Int32 = kInitialTableOffset;
   writeInitExpr(OS, InitExpr);
   writeUleb128(OS, IndirectFunctions.size(), "elem count");
 
   uint32_t TableIndex = kInitialTableOffset;
   for (const FunctionSymbol *Sym : IndirectFunctions) {
     assert(Sym->getTableIndex() == TableIndex);
     writeUleb128(OS, Sym->getFunctionIndex(), "function index");
     ++TableIndex;
void Writer::createSections() {
   createFunctionSection();
   createTableSection();
   createMemorySection();
   createGlobalSection();
   createExportSection();
   createElemSection();
   createCodeSection();
   createDataSection();
+  createCustomSections();
 
   // Custom sections
   if (Config->Relocatable) {
     createLinkingSection();
     createRelocSections();
   }
   if (!Config->StripDebug && !Config->StripAll)
     createNameSection();