[lld][WebAssembly] Do not emit initialization for .bss segments

[tlively]

Summary

This patch fixes a bug where initialization code for .bss segments was
emitted in the memory initialization function even though the .bss
segments were discounted in the datacount section and omitted in the
data section. This was producing invalid binaries due to out-of-bounds
segment indices on the memory.init and data.drop instructions that
were trying to operate on the nonexistent .bss segments.

[tlively]

lld/test/wasm/data-segments.ll

65

There used to be 3 data.drop (FC09) instructions at the end, but now you can see that there are only two, as expected.

[sbc100]

lld/test/wasm/data-segments.ll

65

So we doulbly dropping a segment before?

lld/wasm/Writer.cpp

832

So was the crash was happening when there was only a bss segment? And somehow having more than one segment was masking the crash?

Should we add a specific test case for that? An example that would have previously caused and invalid binary would be good.

I'm stil now clear why we didn't see this is more of our emscripten test cases,

[tlively]

lld/test/wasm/data-segments.ll

65

No, the previous dropping code was "FC0900FC0901FC0902", so it was trying to drop segments 0, 1, and 2. But segment 2 did not exist, so this test was previously producing an invalid binary.

lld/wasm/Writer.cpp

832

The crash was happening when there were any bss segments and the memory was exported, since .bss segments are not isBss if the memory is imported. That's why we never saw this in Emscripten. The existing test case was already invalid, so I don't think we need another one.

[sbc100]

LGTM!

Its shame we don't have a validator here in llvm to checkout the linker output!

lld/test/wasm/data-segments.ll

 ; ACTIVE-NEXT:  - Type:            CUSTOM
 ; ACTIVE-NEXT:    Name:            name
 ; ACTIVE-NEXT:    FunctionNames:
 ; ACTIVE-NEXT:      - Index:           0
 ; ACTIVE-NEXT:        Name:            __wasm_call_ctors
 
 ; PASSIVE-LABEL: - Type:            START
 ; PASSIVE-NEXT:    StartFunction:   1
+; PASSIVE-LABEL: - Type:            DATACOUNT
+; PASSIVE-NEXT:    Count:           2
 ; PASSIVE-LABEL: - Type:            CODE
 ; PASSIVE-NEXT:    Functions:
 ; PASSIVE-NEXT:      - Index:           0
 ; PASSIVE-NEXT:        Locals:          []
 ; PASSIVE-NEXT:        Body:            0B
 ; PASSIVE-NEXT:      - Index:           1
 ; PASSIVE-NEXT:        Locals:          []
-; PASSIVE-NEXT:        Body:            41B4D60041004101FE480200044041B4D6004101427FFE0102001A054180084100410DFC08000041900841004114FC08010041A40841004190CE00FC08020041B4D6004102FE17020041B4D600417FFE0002001A0BFC0900FC0901FC09020B
+; PASSIVE-NEXT:        Body:            41B4D60041004101FE480200044041B4D6004101427FFE0102001A054180084100410DFC08000041900841004114FC08010041B4D6004102FE17020041B4D600417FFE0002001A0BFC0900FC09010B
+
 ; PASSIVE-NEXT:  - Index:           2
 ; PASSIVE-NEXT:    Locals:          []
 ; PASSIVE-NEXT:    Body:            0B
 ; PASSIVE-NEXT:  - Type:            DATA
 ; PASSIVE-NEXT:    Segments:
 ; PASSIVE-NEXT:      - SectionOffset:   3
 ; PASSIVE-NEXT:        InitFlags:       1
 ; PASSIVE-NEXT:        Content:         636F6E7374616E74000000002B

lld/wasm/SyntheticSections.h

public:
   bool isNeeded() const override { return exports.size() > 0; }
   void writeBody() override;
 
   std::vector<llvm::wasm::WasmExport> exports;
 };
 
 class StartSection : public SyntheticSection {
 public:
-  StartSection(uint32_t numSegments)
-      : SyntheticSection(llvm::wasm::WASM_SEC_START), numSegments(numSegments) {
-  }
+  StartSection(bool hasInitializedSegments)
+      : SyntheticSection(llvm::wasm::WASM_SEC_START),
+        hasInitializedSegments(hasInitializedSegments) {}
   bool isNeeded() const override;
   void writeBody() override;
 
 protected:
-  uint32_t numSegments;
+  bool hasInitializedSegments;
 };
 
 class ElemSection : public SyntheticSection {
 public:
   ElemSection()
       : SyntheticSection(llvm::wasm::WASM_SEC_ELEM) {}
   bool isNeeded() const override { return indirectFunctions.size() > 0; };
   void writeBody() override;

lld/wasm/SyntheticSections.cpp

void ExportSection::writeBody() {
   raw_ostream &os = bodyOutputStream;
 
   writeUleb128(os, exports.size(), "export count");
   for (const WasmExport &export_ : exports)
     writeExport(os, export_);
 }
 
 bool StartSection::isNeeded() const {
-  return !config->relocatable && numSegments && config->sharedMemory;
+  return !config->relocatable && hasInitializedSegments && config->sharedMemory;
 }
 
 void StartSection::writeBody() {
   raw_ostream &os = bodyOutputStream;
   writeUleb128(os, WasmSym::initMemory->getFunctionIndex(), "function index");
 }
 
 void ElemSection::addEntry(FunctionSymbol *sym) {

lld/wasm/Writer.cpp

 // The writer writes a SymbolTable result to a file.
 class Writer {
 public:
   void run();
 
 private:
   void openFile();
 
+  bool needsPassiveInitialization(const OutputSegment *segment);
+  bool hasPassiveInitializedSegments();
+
   void createInitMemoryFunction();
   void createApplyRelocationsFunction();
   void createCallCtorsFunction();
   void createInitTLSFunction();
 
   void assignIndexes();
   void populateSymtab();
   void populateProducers();
std::string functionBody;
     raw_string_ostream os(functionBody);
     writeUleb128(os, bodyContent.size(), "function size");
     os << bodyContent;
   }
   ArrayRef<uint8_t> body = arrayRefFromStringRef(saver.save(functionBody));
   cast<SyntheticFunction>(func->function)->setBody(body);
 }
 
+bool Writer::needsPassiveInitialization(const OutputSegment *segment) {
+  return segment->initFlags & WASM_SEGMENT_IS_PASSIVE &&
+         segment->name != ".tdata" && !segment->isBss;
+}
+
+bool Writer::hasPassiveInitializedSegments() {
+  return std::find_if(segments.begin(), segments.end(),
+                      [this](const OutputSegment *s) {
+                        return this->needsPassiveInitialization(s);
+                      }) != segments.end();
+}
+
 void Writer::createInitMemoryFunction() {
   LLVM_DEBUG(dbgs() << "createInitMemoryFunction\n");
   assert(WasmSym::initMemoryFlag);
   uint32_t flagAddress = WasmSym::initMemoryFlag->getVirtualAddress();
   std::string bodyContent;
   {
     raw_string_ostream os(bodyContent);
     writeUleb128(os, 0, "num locals");
 
-    if (segments.size()) {
+    if (hasPassiveInitializedSegments()) {
       // Initialize memory in a thread-safe manner. The thread that successfully
       // increments the flag from 0 to 1 is is responsible for performing the
       // memory initialization. Other threads go sleep on the flag until the
       // first thread finishing initializing memory, increments the flag to 2,
       // and wakes all the other threads. Once the flag has been set to 2,
       // subsequently started threads will skip the sleep. All threads
       // unconditionally drop their passive data segments once memory has been
       // initialized. The generated code is as follows:
if (hasPassiveInitializedSegments()) {
       writeUleb128(os, WASM_OPCODE_I32_ATOMIC_WAIT, "i32.atomic.wait");
       writeMemArg(os, 2, 0);
       writeU8(os, WASM_OPCODE_DROP, "drop");
 
       writeU8(os, WASM_OPCODE_ELSE, "ELSE");
 
       // Did increment 0, so conditionally initialize passive data segments
       for (const OutputSegment *s : segments) {
-        if (s->initFlags & WASM_SEGMENT_IS_PASSIVE && s->name != ".tdata") {
+        if (needsPassiveInitialization(s)) {
           // destination address
           writeI32Const(os, s->startVA, "destination address");
           // source segment offset
           writeI32Const(os, 0, "segment offset");
           // memory region size
           writeI32Const(os, s->size, "memory region size");
           // memory.init instruction
           writeU8(os, WASM_OPCODE_MISC_PREFIX, "bulk-memory prefix");
           writeUleb128(os, WASM_OPCODE_MEMORY_INIT, "memory.init");
           writeUleb128(os, s->index, "segment index immediate");
           writeU8(os, 0, "memory index immediate");
         }
       }
 
       // Set flag to 2 to mark end of initialization
       writeI32Const(os, flagAddress, "flag address");
       writeI32Const(os, 2, "flag value");
       writeU8(os, WASM_OPCODE_ATOMICS_PREFIX, "atomics prefix");
       writeUleb128(os, WASM_OPCODE_I32_ATOMIC_STORE, "i32.atomic.store");
       writeMemArg(os, 2, 0);
 
       // Notify any waiters that memory initialization is complete
       writeI32Const(os, flagAddress, "flag address");
       writeI32Const(os, -1, "number of waiters");
       writeU8(os, WASM_OPCODE_ATOMICS_PREFIX, "atomics prefix");
       writeUleb128(os, WASM_OPCODE_ATOMIC_NOTIFY, "atomic.notify");
       writeMemArg(os, 2, 0);
       writeU8(os, WASM_OPCODE_DROP, "drop");
 
       writeU8(os, WASM_OPCODE_END, "END");
 
       // Unconditionally drop passive data segments
       for (const OutputSegment *s : segments) {
-        if (s->initFlags & WASM_SEGMENT_IS_PASSIVE && s->name != ".tdata") {
+        if (needsPassiveInitialization(s)) {
           // data.drop instruction
           writeU8(os, WASM_OPCODE_MISC_PREFIX, "bulk-memory prefix");
           writeUleb128(os, WASM_OPCODE_DATA_DROP, "data.drop");
           writeUleb128(os, s->index, "segment index immediate");
         }
       }
     }
     writeU8(os, WASM_OPCODE_END, "END");
void Writer::createSyntheticSections() {
   out.typeSec = make<TypeSection>();
   out.importSec = make<ImportSection>();
   out.functionSec = make<FunctionSection>();
   out.tableSec = make<TableSection>();
   out.memorySec = make<MemorySection>();
   out.eventSec = make<EventSection>();
   out.globalSec = make<GlobalSection>();
   out.exportSec = make<ExportSection>();
-  out.startSec = make<StartSection>(segments.size());
+  out.startSec = make<StartSection>(hasPassiveInitializedSegments());
   out.elemSec = make<ElemSection>();
   out.dataCountSec = make<DataCountSection>(segments);
   out.linkingSec = make<LinkingSection>(initFunctions, segments);
   out.nameSec = make<NameSection>();
   out.producersSec = make<ProducersSection>();
   out.targetFeaturesSec = make<TargetFeaturesSection>();
 }