[NFC] Encapsulate source map reader state (#7132)

Move all state relevant to reading source maps out of WasmBinaryReader
and into a new utility, SourceMapReader. This is a prerequisite for
parallelizing the parsing of function bodies, since the source map
reader state is different at the beginning of each function.

Also take the opportunity to simplify the way we read source maps, for
example by deferring the reading of anything but the position of a debug
location until it will be used and by using `std::optional` instead of
singleton `std::set`s to store function prologue and epilogue debug
locations.

Author: tlively

src/ir/module-utils.cpp

@@ -26,32 +26,20 @@ namespace wasm::ModuleUtils {
 
 // Update the file name indices when moving a set of debug locations from one
 // module to another.
-static void updateLocationSet(std::set<Function::DebugLocation>& locations,
-                              std::vector<Index>& fileIndexMap) {
-  std::set<Function::DebugLocation> updatedLocations;
-
-  for (auto iter : locations) {
-    iter.fileIndex = fileIndexMap[iter.fileIndex];
-    updatedLocations.insert(iter);
+static void updateLocation(std::optional<Function::DebugLocation>& location,
+                           std::vector<Index>& fileIndexMap) {
+  if (location) {
+    location->fileIndex = fileIndexMap[location->fileIndex];
   }
-  locations.clear();
-  std::swap(locations, updatedLocations);
 }
 
 // Update the symbol name indices when moving a set of debug locations from one
 // module to another.
-static void updateSymbolSet(std::set<Function::DebugLocation>& locations,
-                            std::vector<Index>& symbolIndexMap) {
-  std::set<Function::DebugLocation> updatedLocations;
-
-  for (auto iter : locations) {
-    if (iter.symbolNameIndex) {
-      iter.symbolNameIndex = symbolIndexMap[*iter.symbolNameIndex];
-    }
-    updatedLocations.insert(iter);
+static void updateSymbol(std::optional<Function::DebugLocation>& location,
+                         std::vector<Index>& symbolIndexMap) {
+  if (location && location->symbolNameIndex) {
+    location->symbolNameIndex = symbolIndexMap[*location->symbolNameIndex];
   }
-  locations.clear();
-  std::swap(locations, updatedLocations);
 }
 
 // Copies a function into a module. If newName is provided it is used as the
@@ -94,8 +82,8 @@ copyFunctionWithoutAdd(Function* func,
         iter.second->fileIndex = (*fileIndexMap)[iter.second->fileIndex];
       }
     }
-    updateLocationSet(ret->prologLocation, *fileIndexMap);
-    updateLocationSet(ret->epilogLocation, *fileIndexMap);
+    updateLocation(ret->prologLocation, *fileIndexMap);
+    updateLocation(ret->epilogLocation, *fileIndexMap);
   }
   if (symbolNameIndexMap) {
     for (auto& iter : ret->debugLocations) {
@@ -105,8 +93,8 @@ copyFunctionWithoutAdd(Function* func,
             (*symbolNameIndexMap)[*(iter.second->symbolNameIndex)];
         }
       }
-      updateSymbolSet(ret->prologLocation, *symbolNameIndexMap);
-      updateSymbolSet(ret->epilogLocation, *symbolNameIndexMap);
+      updateSymbol(ret->prologLocation, *symbolNameIndexMap);
+      updateSymbol(ret->epilogLocation, *symbolNameIndexMap);
     }
   }
   ret->module = func->module;

src/parsing.h

@@ -43,15 +43,6 @@ struct ParseException {
   void dump(std::ostream& o) const;
 };
 
-struct MapParseException {
-  std::string text;
-
-  MapParseException() : text("unknown parse error") {}
-  MapParseException(std::string text) : text(text) {}
-
-  void dump(std::ostream& o) const;
-};
-
 // Helper for parsers that may not have unique label names. This transforms
 // the names into unique ones, as required by Binaryen IR.
 struct UniqueNameMapper {

src/passes/DebugLocationPropagation.cpp

@@ -64,10 +64,10 @@ struct DebugLocationPropagation
         if (auto it = locs.find(previous); it != locs.end()) {
           locs[curr] = it->second;
         }
-      } else if (self->getFunction()->prologLocation.size()) {
+      } else if (self->getFunction()->prologLocation) {
         // Instructions may inherit their locations from the function
         // prolog.
-        locs[curr] = *self->getFunction()->prologLocation.begin();
+        locs[curr] = *self->getFunction()->prologLocation;
       }
     }
     expressionStack.push_back(curr);

src/passes/Print.cpp

@@ -3060,8 +3060,8 @@ void PrintSExpression::visitDefinedFunction(Function* curr) {
   currFunction = curr;
   lastPrintedLocation = std::nullopt;
   lastPrintIndent = 0;
-  if (currFunction->prologLocation.size()) {
-    printDebugLocation(*currFunction->prologLocation.begin());
+  if (currFunction->prologLocation) {
+    printDebugLocation(*currFunction->prologLocation);
   }
   handleSignature(curr, true);
   incIndent();
@@ -3095,14 +3095,14 @@ void PrintSExpression::visitDefinedFunction(Function* curr) {
     }
     assert(controlFlowDepth == 0);
   }
-  if (currFunction->epilogLocation.size()) {
+  if (currFunction->epilogLocation) {
     // Print last debug location: mix of decIndent and printDebugLocation
     // logic.
     doIndent(o, indent);
     if (!minify) {
       indent--;
     }
-    printDebugLocation(*currFunction->epilogLocation.begin());
+    printDebugLocation(*currFunction->epilogLocation);
     o << ')';
   } else {
     decIndent();

src/source-map.h

@@ -0,0 +1,101 @@
+/*
+ * Copyright 2024 WebAssembly Community Group participants
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef wasm_source_map_h
+#define wasm_source_map_h
+
+#include <optional>
+#include <unordered_map>
+
+#include "wasm.h"
+
+namespace wasm {
+
+struct MapParseException {
+  std::string text;
+
+  MapParseException(std::string text) : text(text) {}
+  void dump(std::ostream& o) const;
+};
+
+class SourceMapReader {
+  const std::vector<char>& buffer;
+
+  // Current position in the source map buffer.
+  size_t pos = 0;
+
+  // The location in the binary the next debug location will correspond to. 0
+  // iff there are no more debug locations.
+  size_t location = 0;
+
+  // The file index, line, column, and symbol index the next debug location will
+  // be offset from.
+  uint32_t file = 0;
+  uint32_t line = 1;
+  uint32_t col = 0;
+  uint32_t symbol = 0;
+
+  // Whether the last read record had position and symbol information.
+  bool hasInfo = false;
+  bool hasSymbol = false;
+
+public:
+  SourceMapReader(const std::vector<char>& buffer) : buffer(buffer) {}
+
+  void readHeader(Module& wasm);
+
+  std::optional<Function::DebugLocation>
+  readDebugLocationAt(size_t currLocation);
+
+  // Do not reuse debug info across function boundaries.
+  void finishFunction() { hasInfo = false; }
+
+private:
+  char peek() {
+    if (pos >= buffer.size()) {
+      throw MapParseException("unexpected end of source map");
+    }
+    return buffer[pos];
+  }
+
+  char get() {
+    char c = peek();
+    ++pos;
+    return c;
+  }
+
+  bool maybeGet(char c) {
+    if (pos < buffer.size() && peek() == c) {
+      ++pos;
+      return true;
+    }
+    return false;
+  }
+
+  void expect(char c) {
+    using namespace std::string_literals;
+    char got = get();
+    if (got != c) {
+      throw MapParseException("expected '"s + c + "', got '" + got + "'");
+    }
+  }
+
+  int32_t readBase64VLQ();
+};
+
+} // namespace wasm
+
+#endif // wasm_source_map_h

src/tools/wasm-dis.cpp

@@ -18,6 +18,7 @@
 // wasm2asm console tool
 //
 
+#include "source-map.h"
 #include "support/colors.h"
 #include "support/file.h"
 #include "wasm-io.h"

src/wasm-binary.h

@@ -28,6 +28,7 @@
 #include "ir/import-utils.h"
 #include "ir/module-utils.h"
 #include "parsing.h"
+#include "source-map.h"
 #include "wasm-builder.h"
 #include "wasm-ir-builder.h"
 #include "wasm-traversal.h"
@@ -1403,41 +1404,13 @@ class WasmBinaryWriter {
   void prepare();
 };
 
+extern std::vector<char> defaultEmptySourceMap;
+
 class WasmBinaryReader {
   Module& wasm;
   MixedArena& allocator;
   const std::vector<char>& input;
 
-  // Source map debugging support.
-
-  std::istream* sourceMap;
-
-  // The binary position that the next debug location refers to. That is, this
-  // is the first item in a source map entry that we have read (the "column", in
-  // source map terms, which for wasm means the offset in the binary). We have
-  // read this entry, but have not used it yet (we use it when we read the
-  // expression at this binary offset).
-  //
-  // This is set to 0 as an invalid value if we reach the end of the source map
-  // and there is nothing left to read.
-  size_t nextDebugPos;
-
-  // The debug location (file:line:col) corresponding to |nextDebugPos|. That
-  // is, this is the next 3 fields in a source map entry that we have read, but
-  // not used yet.
-  //
-  // If that location has no debug info (it lacks those 3 fields), then this
-  // contains the info from the previous one, because in a source map, these
-  // fields are relative to their last appearance, so we cannot forget them (we
-  // can't just do something like std::optional<DebugLocation> or such); for
-  // example, if we have line number 100, then no debug info, and then line
-  // number 500, then when we get to 500 we will see "+400" which is relative to
-  // the last existing line number (we "skip" over a place without debug info).
-  Function::DebugLocation nextDebugLocation;
-
-  // Whether debug info is present on |nextDebugPos| (see comment there).
-  bool nextDebugLocationHasDebugInfo;
-
   // Settings.
 
   bool debugInfo = true;
@@ -1448,17 +1421,20 @@ class WasmBinaryReader {
 
   size_t pos = 0;
   Index startIndex = -1;
-  std::set<Function::DebugLocation> debugLocation;
   size_t codeSectionLocation;
   std::unordered_set<uint8_t> seenSections;
 
+  IRBuilder builder;
+  SourceMapReader sourceMapReader;
+
   // All types defined in the type section
   std::vector<HeapType> types;
 
 public:
   WasmBinaryReader(Module& wasm,
                    FeatureSet features,
-                   const std::vector<char>& input);
+                   const std::vector<char>& input,
+                   const std::vector<char>& sourceMap = defaultEmptySourceMap);
 
   void setDebugInfo(bool value) { debugInfo = value; }
   void setDWARF(bool value) { DWARF = value; }
@@ -1584,8 +1560,6 @@ class WasmBinaryReader {
   Expression* readExpression();
   void readGlobals();
 
-  IRBuilder builder;
-
   // validations that cannot be performed on the Module
   void validateBinary();
 
@@ -1607,13 +1581,6 @@ class WasmBinaryReader {
   void readDylink(size_t);
   void readDylink0(size_t);
 
-  // Debug information reading helpers
-  void setDebugLocations(std::istream* sourceMap_) { sourceMap = sourceMap_; }
-  std::unordered_map<std::string, Index> debugInfoFileIndices;
-  std::unordered_map<std::string, Index> debugInfoSymbolNameIndices;
-  void readNextDebugLocation();
-  void readSourceMapHeader();
-
   Index readMemoryAccess(Address& alignment, Address& offset);
   std::tuple<Name, Address, Address> getMemarg();
 

src/wasm-stack.h

@@ -456,8 +456,8 @@ class BinaryenIRToBinaryWriter
 
   void emit(Expression* curr) { writer.visit(curr); }
   void emitHeader() {
-    if (func->prologLocation.size()) {
-      parent.writeDebugLocation(*func->prologLocation.begin());
+    if (func->prologLocation) {
+      parent.writeDebugLocation(*func->prologLocation);
     }
     writer.mapLocalsAndEmitHeader();
   }
@@ -469,8 +469,8 @@ class BinaryenIRToBinaryWriter
   void emitFunctionEnd() {
     // Indicate the debug location corresponding to the end opcode
     // that terminates the function code.
-    if (func->epilogLocation.size()) {
-      parent.writeDebugLocation(*func->epilogLocation.begin());
+    if (func->epilogLocation) {
+      parent.writeDebugLocation(*func->epilogLocation);
     } else {
       // The end opcode has no debug location.
       parent.writeNoDebugLocation();

src/wasm.h

@@ -2095,8 +2095,8 @@ class Function : public Importable {
   // One can explicitly set the debug location of an expression to
   // nullopt to stop the propagation of debug locations.
   std::unordered_map<Expression*, std::optional<DebugLocation>> debugLocations;
-  std::set<DebugLocation> prologLocation;
-  std::set<DebugLocation> epilogLocation;
+  std::optional<DebugLocation> prologLocation;
+  std::optional<DebugLocation> epilogLocation;
 
   // General debugging info support: track instructions and the function itself.
   std::unordered_map<Expression*, BinaryLocations::Span> expressionLocations;

src/wasm/CMakeLists.txt

@@ -2,6 +2,7 @@ file(GLOB wasm_HEADERS ../*.h)
 set(wasm_SOURCES
   literal.cpp
   parsing.cpp
+  source-map.cpp
   wasm.cpp
   wasm-binary.cpp
   wasm-debug.cpp