[NFC] Cleanly separate expression tracking from source maps (#7571)

We track debug info in two ways: for source maps, and for DWARF. DWARF needs
more information, not just the start of each expression's location but also the end
and delimiters (else for an if) as well. The existing code slightly mixed the two
together, which is annoying because code annotations require the same tracking
DWARF does. To improve that, this PR refactors the code to cleanly separate the
two forms of tracking:

* writeSourceMapLocation is now the only method that source maps use.
* trackExpressionStart|End|Delimiter is now used by DWARF (and soon code
  annotations).

As a result,

* BinaryInstWriter no longer needs a sourceMap param. It was using !sourceMap
  in the sense of "maybe DWARF", but given custom annotations we'll need more
  anyhow. Simplify the code by letting the parent decide what to do.
  Replace DWARF scanning ahead of sections with a "preScan" method. This will
  be extended for code annotations later.

Author: kripken

src/wasm-binary.h

@@ -1351,9 +1351,12 @@ class WasmBinaryWriter {
   void writeSourceMapEpilog();
   void writeDebugLocation(const Function::DebugLocation& loc);
   void writeNoDebugLocation();
-  void writeDebugLocation(Expression* curr, Function* func);
-  void writeDebugLocationEnd(Expression* curr, Function* func);
-  void writeExtraDebugLocation(Expression* curr, Function* func, size_t id);
+  void writeSourceMapLocation(Expression* curr, Function* func);
+
+  // Track where expressions go in the binary format.
+  void trackExpressionStart(Expression* curr, Function* func);
+  void trackExpressionEnd(Expression* curr, Function* func);
+  void trackExpressionDelimiter(Expression* curr, Function* func, size_t id);
 
   // helpers
   void writeInlineString(std::string_view name);
@@ -1614,9 +1617,9 @@ class WasmBinaryReader {
 
   static Name escape(Name name);
   void findAndReadNames();
-  void readFeatures(size_t);
-  void readDylink(size_t);
-  void readDylink0(size_t);
+  void readFeatures(size_t payloadLen);
+  void readDylink(size_t payloadLen);
+  void readDylink0(size_t payloadLen);
 
   Index readMemoryAccess(Address& alignment, Address& offset);
   std::tuple<Name, Address, Address> getMemarg();
@@ -1627,7 +1630,14 @@ class WasmBinaryReader {
   }
 
 private:
-  bool hasDWARFSections();
+  // In certain modes we need to note the locations of expressions, to match
+  // them against sections like DWARF or custom annotations. As this incurs
+  // overhead, we only note locations when we actually need to.
+  bool needCodeLocations = false;
+
+  // Scans ahead in the binary to check certain conditions like
+  // needCodeLocations.
+  void preScan();
 };
 
 } // namespace wasm

src/wasm-stack.h

@@ -93,17 +93,16 @@ class BinaryInstWriter : public OverriddenVisitor<BinaryInstWriter> {
   BinaryInstWriter(WasmBinaryWriter& parent,
                    BufferWithRandomAccess& o,
                    Function* func,
-                   bool sourceMap,
                    bool DWARF)
-    : parent(parent), o(o), func(func), sourceMap(sourceMap), DWARF(DWARF) {}
+    : parent(parent), o(o), func(func), DWARF(DWARF) {}
 
   void visit(Expression* curr) {
-    if (func && !sourceMap) {
-      parent.writeDebugLocation(curr, func);
+    if (func) {
+      parent.trackExpressionStart(curr, func);
     }
     OverriddenVisitor<BinaryInstWriter>::visit(curr);
-    if (func && !sourceMap) {
-      parent.writeDebugLocationEnd(curr, func);
+    if (func) {
+      parent.trackExpressionEnd(curr, func);
     }
   }
 
@@ -136,7 +135,6 @@ class BinaryInstWriter : public OverriddenVisitor<BinaryInstWriter> {
   WasmBinaryWriter& parent;
   BufferWithRandomAccess& o;
   Function* func = nullptr;
-  bool sourceMap;
   bool DWARF;
 
   std::vector<Name> breakStack;
@@ -452,7 +450,7 @@ class BinaryenIRToBinaryWriter
                            bool sourceMap = false,
                            bool DWARF = false)
     : BinaryenIRWriter<BinaryenIRToBinaryWriter>(func), parent(parent),
-      writer(parent, o, func, sourceMap, DWARF), sourceMap(sourceMap) {}
+      writer(parent, o, func, DWARF), sourceMap(sourceMap) {}
 
   void emit(Expression* curr) { writer.visit(curr); }
   void emitHeader() {
@@ -480,7 +478,7 @@ class BinaryenIRToBinaryWriter
   void emitUnreachable() { writer.emitUnreachable(); }
   void emitDebugLocation(Expression* curr) {
     if (sourceMap) {
-      parent.writeDebugLocation(curr, func);
+      parent.writeSourceMapLocation(curr, func);
     }
   }
 
@@ -521,7 +519,7 @@ class StackIRToBinaryWriter {
                         StackIR& stackIR,
                         bool sourceMap = false,
                         bool DWARF = false)
-    : parent(parent), writer(parent, o, func, sourceMap, DWARF), func(func),
+    : parent(parent), writer(parent, o, func, DWARF), func(func),
       stackIR(stackIR), sourceMap(sourceMap) {}
 
   void write();

src/wasm/wasm-binary.cpp

@@ -1489,37 +1489,43 @@ void WasmBinaryWriter::writeNoDebugLocation() {
   }
 }
 
-void WasmBinaryWriter::writeDebugLocation(Expression* curr, Function* func) {
-  if (sourceMap) {
-    auto& debugLocations = func->debugLocations;
-    auto iter = debugLocations.find(curr);
-    if (iter != debugLocations.end() && iter->second) {
-      // There is debug information here, write it out.
-      writeDebugLocation(*(iter->second));
-    } else {
-      // This expression has no debug location.
-      writeNoDebugLocation();
-    }
+void WasmBinaryWriter::writeSourceMapLocation(Expression* curr,
+                                              Function* func) {
+  assert(sourceMap);
+
+  auto& debugLocations = func->debugLocations;
+  auto iter = debugLocations.find(curr);
+  if (iter != debugLocations.end() && iter->second) {
+    // There is debug information here, write it out.
+    writeDebugLocation(*(iter->second));
+  } else {
+    // This expression has no debug location.
+    writeNoDebugLocation();
   }
+}
+
+void WasmBinaryWriter::trackExpressionStart(Expression* curr, Function* func) {
   // If this is an instruction in a function, and if the original wasm had
-  // binary locations tracked, then track it in the output as well.
+  // binary locations tracked, then track it in the output as well. We also
+  // track locations of instructions that have code annotations, as their binary
+  // location goes in the custom section.
   if (func && !func->expressionLocations.empty()) {
     binaryLocations.expressions[curr] =
       BinaryLocations::Span{BinaryLocation(o.size()), 0};
     binaryLocationTrackedExpressionsForFunc.push_back(curr);
   }
 }
 
-void WasmBinaryWriter::writeDebugLocationEnd(Expression* curr, Function* func) {
+void WasmBinaryWriter::trackExpressionEnd(Expression* curr, Function* func) {
   if (func && !func->expressionLocations.empty()) {
     auto& span = binaryLocations.expressions.at(curr);
     span.end = o.size();
   }
 }
 
-void WasmBinaryWriter::writeExtraDebugLocation(Expression* curr,
-                                               Function* func,
-                                               size_t id) {
+void WasmBinaryWriter::trackExpressionDelimiter(Expression* curr,
+                                                Function* func,
+                                                size_t id) {
   if (func && !func->expressionLocations.empty()) {
     binaryLocations.delimiters[curr][id] = o.size();
   }
@@ -1799,11 +1805,19 @@ WasmBinaryReader::WasmBinaryReader(Module& wasm,
   wasm.features = features;
 }
 
-bool WasmBinaryReader::hasDWARFSections() {
+void WasmBinaryReader::preScan() {
+  // TODO: Once we support code annotations here, we will need to always scan,
+  //       but for now, DWARF is the only reason.
+  if (!DWARF) {
+    return;
+  }
+
   assert(pos == 0);
   getInt32(); // magic
   getInt32(); // version
-  bool has = false;
+
+  bool foundDWARF = false;
+
   while (more()) {
     uint8_t sectionCode = getInt8();
     uint32_t payloadLen = getU32LEB();
@@ -1813,31 +1827,32 @@ bool WasmBinaryReader::hasDWARFSections() {
     auto oldPos = pos;
     if (sectionCode == BinaryConsts::Section::Custom) {
       auto sectionName = getInlineString();
-      if (Debug::isDWARFSection(sectionName)) {
-        has = true;
+      // DWARF sections contain code offsets.
+      if (DWARF && Debug::isDWARFSection(sectionName)) {
+        needCodeLocations = true;
+        foundDWARF = true;
         break;
       }
     }
     pos = oldPos + payloadLen;
   }
+
+  if (DWARF && !foundDWARF) {
+    // The user asked for DWARF, but no DWARF sections exist in practice, so
+    // disable the support.
+    DWARF = false;
+  }
+
+  // Reset.
   pos = 0;
-  return has;
 }
 
 void WasmBinaryReader::read() {
-  if (DWARF) {
-    // In order to update dwarf, we must store info about each IR node's
-    // binary position. This has noticeable memory overhead, so we don't do it
-    // by default: the user must request it by setting "DWARF", and even if so
-    // we scan ahead to see that there actually *are* DWARF sections, so that
-    // we don't do unnecessary work.
-    if (!hasDWARFSections()) {
-      DWARF = false;
-    }
-  }
+  preScan();
 
   // Skip ahead and read the name section so we know what names to use when we
   // construct module elements.
+  // TODO: Combine this pre-scan with the one in preScan().
   if (debugInfo) {
     findAndReadNames();
   }
@@ -1879,7 +1894,7 @@ void WasmBinaryReader::read() {
         readFunctionSignatures();
         break;
       case BinaryConsts::Section::Code:
-        if (DWARF) {
+        if (needCodeLocations) {
           codeSectionLocation = pos;
         }
         readFunctions();
@@ -2871,7 +2886,7 @@ void WasmBinaryReader::readFunctions() {
   if (numFuncBodies + numFuncImports != wasm.functions.size()) {
     throwError("invalid function section size, must equal types");
   }
-  if (DWARF) {
+  if (needCodeLocations) {
     builder.setBinaryLocation(&pos, codeSectionLocation);
   }
   for (size_t i = 0; i < numFuncBodies; i++) {
@@ -2885,7 +2900,7 @@ void WasmBinaryReader::readFunctions() {
     auto& func = wasm.functions[numFuncImports + i];
     currFunction = func.get();
 
-    if (DWARF) {
+    if (needCodeLocations) {
       func->funcLocation = BinaryLocations::FunctionLocations{
         BinaryLocation(sizePos - codeSectionLocation),
         BinaryLocation(pos - codeSectionLocation),

src/wasm/wasm-stack.cpp

@@ -50,8 +50,8 @@ void BinaryInstWriter::visitIf(If* curr) {
 }
 
 void BinaryInstWriter::emitIfElse(If* curr) {
-  if (func && !sourceMap) {
-    parent.writeExtraDebugLocation(curr, func, BinaryLocations::Else);
+  if (func) {
+    parent.trackExpressionDelimiter(curr, func, BinaryLocations::Else);
   }
   o << int8_t(BinaryConsts::Else);
 }
@@ -2156,16 +2156,16 @@ void BinaryInstWriter::visitTryTable(TryTable* curr) {
 }
 
 void BinaryInstWriter::emitCatch(Try* curr, Index i) {
-  if (func && !sourceMap) {
-    parent.writeExtraDebugLocation(curr, func, i);
+  if (func) {
+    parent.trackExpressionDelimiter(curr, func, i);
   }
   o << int8_t(BinaryConsts::Catch_Legacy)
     << U32LEB(parent.getTagIndex(curr->catchTags[i]));
 }
 
 void BinaryInstWriter::emitCatchAll(Try* curr) {
-  if (func && !sourceMap) {
-    parent.writeExtraDebugLocation(curr, func, curr->catchBodies.size());
+  if (func) {
+    parent.trackExpressionDelimiter(curr, func, curr->catchBodies.size());
   }
   o << int8_t(BinaryConsts::CatchAll_Legacy);
 }
@@ -2736,8 +2736,8 @@ void BinaryInstWriter::emitScopeEnd(Expression* curr) {
   assert(!breakStack.empty());
   breakStack.pop_back();
   o << int8_t(BinaryConsts::End);
-  if (func && !sourceMap) {
-    parent.writeDebugLocationEnd(curr, func);
+  if (func) {
+    parent.trackExpressionEnd(curr, func);
   }
 }
 
@@ -3213,12 +3213,9 @@ void StackIRToBinaryWriter::write() {
       case StackInst::LoopBegin:
       case StackInst::TryTableBegin: {
         if (sourceMap) {
-          parent.writeDebugLocation(inst->origin, func);
+          parent.writeSourceMapLocation(inst->origin, func);
         }
         writer.visit(inst->origin);
-        if (sourceMap) {
-          parent.writeDebugLocationEnd(inst->origin, func);
-        }
         break;
       }
       case StackInst::TryEnd: