[lld][WebAssembly] Initial support merging string data

This change adds support for a new WASM_SEG_FLAG_STRINGS flag in
the object format which works in a similar fashion to SHF_STRINGS
in the ELF world.

Unlike the ELF linker this support is currently limited:
- No support for SHF_MERGE (non-string merging)
- Always do full tail merging ("lo" can be merged with "hello")
- Only support single byte strings (p2align 0)

Like the ELF linker merging is only performed at `-O1` and above.

This fixes part of https://bugs.llvm.org/show_bug.cgi?id=48828,
although crucially it doesn't not currently support debug sections
because they are not represented by data segments (they are custom
sections)

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

lld/test/wasm/merge-string.s

@@ -0,0 +1,65 @@
+// RUN: llvm-mc -filetype=obj -triple=wasm32-unknown-unknown %s -o %t.o
+// RUN: wasm-ld -O2 %t.o -o %t.wasm --no-gc-sections --no-entry
+// RUN: obj2yaml %t.wasm | FileCheck %s --check-prefixes=COMMON,MERGE
+// RUN: wasm-ld -O0 %t.o -o %t2.wasm --no-gc-sections --no-entry
+// RUN: obj2yaml %t2.wasm | FileCheck --check-prefixes=COMMON,NOMERGE %s
+
+        .section .rodata1,"S",@
+        .asciz "abc"
+foo:
+        .ascii "a"
+        .size foo, 1
+bar:
+        .asciz "bc"
+        .asciz "bc"
+        .size bar, 4
+
+.globl foo
+.globl bar
+.export_name    foo, foo
+.export_name    bar, bar
+
+//       COMMON:  - Type:            GLOBAL
+//  COMMON-NEXT:    Globals:
+//  COMMON-NEXT:      - Index:           0
+//  COMMON-NEXT:        Type:            I32
+//  COMMON-NEXT:        Mutable:         true
+//  COMMON-NEXT:        InitExpr:
+//  COMMON-NEXT:          Opcode:          I32_CONST
+//  COMMON-NEXT:          Value:           66576
+//  COMMON-NEXT:      - Index:           1
+//  COMMON-NEXT:        Type:            I32
+//  COMMON-NEXT:        Mutable:         false
+//  COMMON-NEXT:        InitExpr:
+//  COMMON-NEXT:          Opcode:          I32_CONST
+//   MERGE-NEXT:          Value:           1024
+// NOMERGE-NEXT:          Value:           1028
+//  COMMON-NEXT:      - Index:           2
+//  COMMON-NEXT:        Type:            I32
+//  COMMON-NEXT:        Mutable:         false
+//  COMMON-NEXT:        InitExpr:
+//  COMMON-NEXT:          Opcode:          I32_CONST
+//   MERGE-NEXT:          Value:           1025
+// NOMERGE-NEXT:          Value:           1029
+//  COMMON-NEXT:  - Type:            EXPORT
+//  COMMON-NEXT:    Exports:
+//  COMMON-NEXT:      - Name:            memory
+//  COMMON-NEXT:        Kind:            MEMORY
+//  COMMON-NEXT:        Index:           0
+//  COMMON-NEXT:      - Name:            foo
+//  COMMON-NEXT:        Kind:            GLOBAL
+//  COMMON-NEXT:        Index:           1
+//  COMMON-NEXT:      - Name:            bar
+//  COMMON-NEXT:        Kind:            GLOBAL
+//  COMMON-NEXT:        Index:           2
+
+//
+//       COMMON:  - Type:            DATA
+//  COMMON-NEXT:    Segments:
+//  COMMON-NEXT:      - SectionOffset:   7
+//  COMMON-NEXT:        InitFlags:       0
+//  COMMON-NEXT:        Offset:
+//  COMMON-NEXT:          Opcode:          I32_CONST
+//  COMMON-NEXT:          Value:           1024
+//   MERGE-NEXT:          Content:         '61626300'
+// NOMERGE-NEXT:          Content:         '6162630061626300626300'

lld/wasm/CMakeLists.txt

@@ -10,6 +10,7 @@ add_lld_library(lldWasm
   MapFile.cpp
   MarkLive.cpp
   OutputSections.cpp
+  OutputSegment.cpp
   Relocations.cpp
   SymbolTable.cpp
   Symbols.cpp

lld/wasm/Driver.cpp

@@ -385,7 +385,7 @@ static void readConfigs(opt::InputArgList &args) {
                    LLVM_ENABLE_NEW_PASS_MANAGER);
   config->ltoDebugPassManager = args.hasArg(OPT_lto_debug_pass_manager);
   config->mapFile = args.getLastArgValue(OPT_Map);
-  config->optimize = args::getInteger(args, OPT_O, 0);
+  config->optimize = args::getInteger(args, OPT_O, 1);
   config->outputFile = args.getLastArgValue(OPT_o);
   config->relocatable = args.hasArg(OPT_relocatable);
   config->gcSections =
@@ -795,6 +795,18 @@ static void wrapSymbols(ArrayRef<WrappedSymbol> wrapped) {
     symtab->wrap(w.sym, w.real, w.wrap);
 }
 
+static void splitSections() {
+  // splitIntoPieces needs to be called on each MergeInputSection
+  // before calling finalizeContents().
+  LLVM_DEBUG(llvm::dbgs() << "splitSections\n");
+  parallelForEach(symtab->objectFiles, [](ObjFile *file) {
+    for (InputSegment *seg : file->segments) {
+      if (auto *s = dyn_cast<MergeInputSegment>(seg))
+        s->splitIntoPieces();
+    }
+  });
+}
+
 void LinkerDriver::linkerMain(ArrayRef<const char *> argsArr) {
   WasmOptTable parser;
   opt::InputArgList args = parser.parse(argsArr.slice(1));
@@ -981,6 +993,10 @@ void LinkerDriver::linkerMain(ArrayRef<const char *> argsArr) {
   if (errorCount())
     return;
 
+  // Split WASM_SEG_FLAG_STRINGS sections into pieces in preparation for garbage
+  // collection.
+  splitSections();
+
   // Do size optimizations: garbage collection
   markLive();
 

lld/wasm/InputChunks.cpp

@@ -13,6 +13,7 @@
 #include "lld/Common/ErrorHandler.h"
 #include "lld/Common/LLVM.h"
 #include "llvm/Support/LEB128.h"
+#include "llvm/Support/xxhash.h"
 
 #define DEBUG_TYPE "lld"
 
@@ -126,6 +127,10 @@ void InputChunk::writeTo(uint8_t *buf) const {
   memcpy(buf + outSecOff, data().data(), data().size());
 
   // Apply relocations
+  relocate(buf + outSecOff);
+}
+
+void InputChunk::relocate(uint8_t *buf) const {
   if (relocations.empty())
     return;
 
@@ -135,11 +140,11 @@ void InputChunk::writeTo(uint8_t *buf) const {
 
   LLVM_DEBUG(dbgs() << "applying relocations: " << toString(this)
                     << " count=" << relocations.size() << "\n");
-  int32_t off = outSecOff - getInputSectionOffset();
+  int32_t inputSectionOffset = getInputSectionOffset();
   auto tombstone = getTombstone();
 
   for (const WasmRelocation &rel : relocations) {
-    uint8_t *loc = buf + rel.Offset + off;
+    uint8_t *loc = buf + rel.Offset - inputSectionOffset;
     auto value = file->calcNewValue(rel, tombstone, this);
     LLVM_DEBUG(dbgs() << "apply reloc: type=" << relocTypeToString(rel.Type));
     if (rel.Type != R_WASM_TYPE_INDEX_LEB)
@@ -357,8 +362,20 @@ void InputFunction::writeTo(uint8_t *buf) const {
   LLVM_DEBUG(dbgs() << "  total: " << (buf + chunkSize - orig) << "\n");
 }
 
+uint64_t InputSegment::getOffset(uint64_t offset) const {
+  if (const MergeInputSegment *ms = dyn_cast<MergeInputSegment>(this)) {
+    LLVM_DEBUG(dbgs() << "getOffset(merged): " << getName() << "\n");
+    LLVM_DEBUG(dbgs() << "offset: " << offset << "\n");
+    LLVM_DEBUG(dbgs() << "parentOffset: " << ms->getParentOffset(offset)
+                      << "\n");
+    assert(ms->parent);
+    return ms->parent->getOffset(ms->getParentOffset(offset));
+  }
+  return outputSegmentOffset + offset;
+}
+
 uint64_t InputSegment::getVA(uint64_t offset) const {
-  return outputSeg->startVA + outputSegmentOffset + offset;
+  return (outputSeg ? outputSeg->startVA : 0) + getOffset(offset);
 }
 
 // Generate code to apply relocations to the data section at runtime.
@@ -431,6 +448,93 @@ void InputSegment::generateRelocationCode(raw_ostream &os) const {
   }
 }
 
+// Split WASM_SEG_FLAG_STRINGS section. Such a section is a sequence of
+// null-terminated strings.
+void MergeInputSegment::splitStrings(ArrayRef<uint8_t> data) {
+  LLVM_DEBUG(llvm::dbgs() << "splitStrings\n");
+  size_t off = 0;
+  StringRef s = toStringRef(data);
+
+  while (!s.empty()) {
+    size_t end = s.find(0);
+    if (end == StringRef::npos)
+      fatal(toString(this) + ": string is not null terminated");
+    size_t size = end + 1;
+
+    pieces.emplace_back(off, xxHash64(s.substr(0, size)), true);
+    s = s.substr(size);
+    off += size;
+  }
+}
+
+// This function is called after we obtain a complete list of input sections
+// that need to be linked. This is responsible to split section contents
+// into small chunks for further processing.
+//
+// Note that this function is called from parallelForEach. This must be
+// thread-safe (i.e. no memory allocation from the pools).
+void MergeInputSegment::splitIntoPieces() {
+  assert(pieces.empty());
+  // As of now we only support WASM_SEG_FLAG_STRINGS but in the future we
+  // could add other types of splitting (see ELF's splitIntoPieces).
+  assert(segment->Data.LinkingFlags & WASM_SEG_FLAG_STRINGS);
+  splitStrings(data());
+}
+
+SegmentPiece *MergeInputSegment::getSegmentPiece(uint64_t offset) {
+  if (this->data().size() <= offset)
+    fatal(toString(this) + ": offset is outside the section");
+
+  // If Offset is not at beginning of a section piece, it is not in the map.
+  // In that case we need to  do a binary search of the original section piece
+  // vector.
+  auto it = partition_point(
+      pieces, [=](SegmentPiece p) { return p.inputOff <= offset; });
+  return &it[-1];
+}
+
+// Returns the offset in an output section for a given input offset.
+// Because contents of a mergeable section is not contiguous in output,
+// it is not just an addition to a base output offset.
+uint64_t MergeInputSegment::getParentOffset(uint64_t offset) const {
+  // If Offset is not at beginning of a section piece, it is not in the map.
+  // In that case we need to search from the original section piece vector.
+  const SegmentPiece *piece = getSegmentPiece(offset);
+  uint64_t addend = offset - piece->inputOff;
+  return piece->outputOff + addend;
+}
+
+uint32_t SyntheticMergedDataSegment::getSize() const {
+  return builder.getSize();
+}
+
+void SyntheticMergedDataSegment::writeTo(uint8_t *buf) const {
+  builder.write(buf + outSecOff);
+
+  // Apply relocations
+  relocate(buf + outSecOff);
+}
+
+void SyntheticMergedDataSegment::finalizeContents() {
+  // Add all string pieces to the string table builder to create section
+  // contents.
+  for (MergeInputSegment *sec : segments)
+    for (size_t i = 0, e = sec->pieces.size(); i != e; ++i)
+      if (sec->pieces[i].live)
+        builder.add(sec->getData(i));
+
+  // Fix the string table content. After this, the contents will never change.
+  builder.finalize();
+
+  // finalize() fixed tail-optimized strings, so we can now get
+  // offsets of strings. Get an offset for each string and save it
+  // to a corresponding SectionPiece for easy access.
+  for (MergeInputSegment *sec : segments)
+    for (size_t i = 0, e = sec->pieces.size(); i != e; ++i)
+      if (sec->pieces[i].live)
+        sec->pieces[i].outputOff = builder.getOffset(sec->getData(i));
+}
+
 uint64_t InputSection::getTombstoneForSection(StringRef name) {
   // When a function is not live we need to update relocations referring to it.
   // If they occur in DWARF debug symbols, we want to change the pc of the