[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

185

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

Please add a comment what this code is doing.

196

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

198–200

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

315

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

331

You can eliminate this variable.

[sbc100]

• feedback

[sbc100]

wasm/Writer.cpp

315

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

305–306

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

 
 #ifndef LLD_WASM_INPUT_CHUNKS_H
 #define LLD_WASM_INPUT_CHUNKS_H
 
 #include "Config.h"
 #include "InputFiles.h"
 #include "lld/Common/ErrorHandler.h"
 #include "llvm/Object/Wasm.h"
+#include "llvm/Support/LEB128.h"
 
 using llvm::object::WasmSection;
 using llvm::object::WasmSegment;
 using llvm::wasm::WasmFunction;
 using llvm::wasm::WasmRelocation;
 using llvm::wasm::WasmSignature;
 
 namespace llvm {
 class raw_ostream;
 }
 
 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)
+      : 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);
+    PayloadOffset = Count + NameSize;
+  }
+
+  const WasmSection &getSection() { return Section; }
+
+  StringRef getName() const override { return Section.Name; }
+  uint32_t getComdat() const override { return UINT32_MAX; }
+
+protected:
+  ArrayRef<uint8_t> data() const override {
+    return Section.Content.slice(PayloadOffset);
+  }
+  uint32_t getInputSectionOffset() const override {
+    return 0; // TODO if we decide to make absolute offset: Section.Offset
+  }
+
+  const WasmSection &Section;
+  size_t PayloadOffset;
+};
+
 } // namespace wasm
 
 std::string toString(const wasm::InputChunk *);
 } // namespace lld
 
 #endif // LLD_WASM_INPUT_CHUNKS_H

wasm/InputChunks.cpp

void InputChunk::copyRelocations(const WasmSection &Section) {
   for (const WasmRelocation &R : Section.Relocations)
     if (R.Offset >= Start && R.Offset < Start + Size)
       Relocations.push_back(R);
 }
 
 // Copy this input chunk to an mmap'ed output file and apply relocations.
 void InputChunk::writeTo(uint8_t *Buf) const {
   // Copy contents
+  DEBUG(dbgs() << "writing chunk offset = " << OutputOffset << "\n");
+  DEBUG(dbgs() << "writing chunk " << this << "\n");
+  DEBUG(dbgs() << "writing chunk " << getInputSectionOffset() << "\n");
   memcpy(Buf + OutputOffset, data().data(), data().size());
 
   // Apply relocations
   if (Relocations.empty())
     return;
 
   DEBUG(dbgs() << "applying relocations: " << getName()
                << " count=" << Relocations.size() << "\n");

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) {
+      DEBUG(dbgs() << " --------- custom section " << Section.Name << "\n");
+      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;
 };
 
+class CustomSection : public OutputSection {
+public:
+  explicit CustomSection(std::string Name, size_t Size,
+                         ArrayRef<InputSection *> InputSections);
+  size_t getSize() const override {
+    return Header.size() + NameData.size() + Size;
+  }
+  void writeTo(uint8_t *Buf) override;
+
+protected:
+  size_t Size;
+  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, size_t Size,
+                             ArrayRef<InputSection *> InputSections)
+    : OutputSection(WASM_SEC_CUSTOM, Name), Size(Size),
+      InputSections(InputSections) {
+  raw_string_ostream OS(NameData);
+  encodeULEB128(Name.size(), OS);
+  OS << Name;
+  OS.flush();
+
+  createHeader(Size + 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();
 
   uint64_t FileSize = 0;
   uint32_t NumMemoryPages = 0;
   uint32_t MaxMemoryPages = 0;
 
   std::vector<const WasmSignature *> Types;
   DenseMap<WasmSignature, int32_t> TypeIndices;
   std::vector<const Symbol *> ImportedSymbols;
   unsigned NumImportedFunctions = 0;
   unsigned NumImportedGlobals = 0;
   std::vector<Symbol *> ExportedSymbols;
+  std::vector<const CustomSection *> CustomSections;
   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;
+    ArrayRef<InputSection*> Sections = Pair.second;
+    DEBUG(dbgs() << "custom section: " << Name << "\n");
+    uint32_t Offset = 0;
+    for (InputSection *Section : Sections) {
+      Section->OutputOffset = Offset;
+      DEBUG(dbgs() << " input section: " << Section->getName() << "\n");
+      DEBUG(dbgs() << " input section: " << Section << "\n");
+      DEBUG(dbgs() << " input offset: " << Offset << "\n");
+      DEBUG(dbgs() << " input size: " << Section->getSize() << "\n");
+      Offset += Section->getSize();
+    }
+
+    DEBUG(dbgs() << "createCustomSection: " << Name << " size=" << Offset
+                 << "\n");
+    auto Section = make<CustomSection>(Name, Offset, Sections);
+    OutputSections.push_back(Section);
+  }
+}
+
 void Writer::createElemSection() {
   if (IndirectFunctions.empty())
     return;
 
   SyntheticSection *Section = createSyntheticSection(WASM_SEC_ELEM);
   raw_ostream &OS = Section->getStream();
 
   writeUleb128(OS, 1, "segment count");
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();