Add debug information locations to the function prolog/epilog (#1674)

The current patch:
* Preserves the debug locations from function prolog and epilog
* Preserves the debug locations of the nested blocks

Author: yurydelendik

src/passes/Print.cpp

@@ -437,20 +437,29 @@ struct PrintSExpression : public Visitor<PrintSExpression> {
     if (!full) full = isFullForced();
   }
 
-  void visit(Expression* curr) {
+  void printDebugLocation(const Function::DebugLocation &location) {
+    if (lastPrintedLocation == location) {
+      return;
+    }
+    lastPrintedLocation = location;
+    auto fileName = currModule->debugInfoFileNames[location.fileIndex];
+    o << ";;@ " << fileName << ":" << location.lineNumber << ":" << location.columnNumber << '\n';
+    doIndent(o, indent);
+  }
+
+  void printDebugLocation(Expression* curr) {
     if (currFunction) {
       // show an annotation, if there is one
       auto& debugLocations = currFunction->debugLocations;
       auto iter = debugLocations.find(curr);
       if (iter != debugLocations.end()) {
-        auto fileName = currModule->debugInfoFileNames[iter->second.fileIndex];
-        if (lastPrintedLocation != iter->second) {
-          lastPrintedLocation = iter->second;
-          o << ";;@ " << fileName << ":" << iter->second.lineNumber << ":" << iter->second.columnNumber << '\n';
-          doIndent(o, indent);
-        }
+        printDebugLocation(iter->second);
       }
     }
+  }
+
+  void visit(Expression* curr) {
+    printDebugLocation(curr);
     Visitor<PrintSExpression>::visit(curr);
   }
 
@@ -487,7 +496,10 @@ struct PrintSExpression : public Visitor<PrintSExpression> {
     // special-case Block, because Block nesting (in their first element) can be incredibly deep
     std::vector<Block*> stack;
     while (1) {
-      if (stack.size() > 0) doIndent(o, indent);
+      if (stack.size() > 0) {
+        doIndent(o, indent);
+        printDebugLocation(curr);
+      }
       stack.push_back(curr);
       if (full) {
         o << "[" << printType(curr->type) << "] ";
@@ -849,6 +861,9 @@ struct PrintSExpression : public Visitor<PrintSExpression> {
   void visitFunction(Function* curr) {
     currFunction = curr;
     lastPrintedLocation = { 0, 0, 0 };
+    if (currFunction->prologLocation.size()) {
+      printDebugLocation(*currFunction->prologLocation.begin());
+    }
     o << '(';
     printMajor(o, "func ");
     printName(curr->name, o);
@@ -901,7 +916,17 @@ struct PrintSExpression : public Visitor<PrintSExpression> {
       // Print the stack IR.
       WasmPrinter::printStackIR(curr->stackIR.get(), o, curr);
     }
-    decIndent();
+    if (currFunction->epilogLocation.size() && lastPrintedLocation != *currFunction->epilogLocation.begin()) {
+      // Print last debug location: mix of decIndent and printDebugLocation logic.
+      doIndent(o, indent);
+      if (!minify) {
+        indent--;
+      }
+      printDebugLocation(*currFunction->epilogLocation.begin());
+      o << ')';
+    } else {
+      decIndent();
+    }
   }
   void printTableHeader(Table* curr) {
     o << '(';

src/wasm-binary.h

@@ -722,6 +722,7 @@ class WasmBinaryWriter {
 
   void writeSourceMapProlog();
   void writeSourceMapEpilog();
+  void writeDebugLocation(const Function::DebugLocation& loc);
   void writeDebugLocation(Expression* curr, Function* func);
 
   // helpers
@@ -775,7 +776,7 @@ class WasmBinaryBuilder {
 
   size_t pos = 0;
   Index startIndex = -1;
-  bool useDebugLocation;
+  std::set<Function::DebugLocation> debugLocation;
 
   std::set<BinaryConsts::Section> seenSections;
 
@@ -787,7 +788,7 @@ class WasmBinaryBuilder {
       debug(debug),
       sourceMap(nullptr),
       nextDebugLocation(0, { 0, 0, 0 }),
-      useDebugLocation(false) {}
+      debugLocation() {}
 
   void read();
   void readUserSection(size_t payloadLen);
@@ -903,7 +904,6 @@ class WasmBinaryBuilder {
   void setDebugLocations(std::istream* sourceMap_) {
       sourceMap = sourceMap_;
   }
-  Function::DebugLocation debugLocation;
   std::unordered_map<std::string, Index> debugInfoFileIndices;
   void readNextDebugLocation();
   void readSourceMapHeader();

src/wasm-s-parser.h

@@ -51,15 +51,17 @@ class Element {
   bool quoted_;
 
 public:
-  Element(MixedArena& allocator) : isList_(true), list_(allocator), line(-1), col(-1), loc(nullptr) {}
+  Element(MixedArena& allocator) : isList_(true), list_(allocator), line(-1), col(-1), startLoc(nullptr), endLoc(nullptr) {}
 
   bool isList() const { return isList_; }
   bool isStr() const { return !isList_; }
   bool dollared() const { return isStr() && dollared_; }
   bool quoted() const { return isStr() && quoted_; }
 
   size_t line, col;
-  SourceLocation* loc;
+  // original locations at the start/end of the S-Expression list
+  SourceLocation* startLoc;
+  SourceLocation* endLoc;
 
   // list methods
   List& list();
@@ -72,7 +74,7 @@ class Element {
   cashew::IString str() const;
   const char* c_str() const;
   Element* setString(cashew::IString str__, bool dollared__, bool quoted__);
-  Element* setMetadata(size_t line_, size_t col_, SourceLocation* loc_);
+  Element* setMetadata(size_t line_, size_t col_, SourceLocation* startLoc_);
 
   // printing
   friend std::ostream& operator<<(std::ostream& o, Element& e);
@@ -218,6 +220,8 @@ class SExpressionWasmBuilder {
   void parseElem(Element& s);
   void parseInnerElem(Element& s, Index i = 1, Expression* offset = nullptr);
   void parseType(Element& s);
+
+  Function::DebugLocation getDebugLocation(const SourceLocation& loc);
 };
 
 } // namespace wasm

src/wasm-stack.h

@@ -257,6 +257,9 @@ class StackIRFunctionStackWriter : StackWriter<StackWriterMode::Stack2Binary, Pa
 
 template<StackWriterMode Mode, typename Parent>
 void StackWriter<Mode, Parent>::mapLocalsAndEmitHeader() {
+  if (func->prologLocation.size()) {
+    parent.writeDebugLocation(*func->prologLocation.begin());
+  }
   // Map them
   for (Index i = 0; i < func->getNumParams(); i++) {
     size_t curr = mappedLocals.size();
@@ -1166,6 +1169,9 @@ bool StackWriter<Mode, Parent>::justAddToStack(Expression* curr) {
 
 template<StackWriterMode Mode, typename Parent>
 void StackWriter<Mode, Parent>::finishFunctionBody() {
+  if (func->epilogLocation.size()) {
+    parent.writeDebugLocation(*func->epilogLocation.begin());
+  }
   o << int8_t(BinaryConsts::End);
 }
 

src/wasm.h

@@ -618,10 +618,16 @@ class Function {
     uint32_t fileIndex, lineNumber, columnNumber;
     bool operator==(const DebugLocation& other) const { return fileIndex == other.fileIndex && lineNumber == other.lineNumber && columnNumber == other.columnNumber; }
     bool operator!=(const DebugLocation& other) const { return !(*this == other); }
+    bool operator<(const DebugLocation& other) const {
+      return fileIndex != other.fileIndex ? fileIndex < other.fileIndex :
+        lineNumber != other.lineNumber ? lineNumber < other.lineNumber : columnNumber < other.columnNumber;
+    }
   };
   std::unordered_map<Expression*, DebugLocation> debugLocations;
+  std::set<DebugLocation> prologLocation;
+  std::set<DebugLocation> epilogLocation;
 
-  Function() : result(none) {}
+  Function() : result(none), prologLocation(), epilogLocation() {}
 
   size_t getNumParams();
   size_t getNumVars();

src/wasm/wasm-binary.cpp

@@ -565,14 +565,20 @@ void WasmBinaryWriter::writeUserSection(const UserSection& section) {
   finishSection(start);
 }
 
+void WasmBinaryWriter::writeDebugLocation(const Function::DebugLocation& loc) {
+  if (loc == lastDebugLocation) {
+    return;
+  }
+  auto offset = o.size();
+  sourceMapLocations.emplace_back(offset, &loc);
+  lastDebugLocation = loc;
+}
+
 void WasmBinaryWriter::writeDebugLocation(Expression* curr, Function* func) {
   auto& debugLocations = func->debugLocations;
   auto iter = debugLocations.find(curr);
-  if (iter != debugLocations.end() && iter->second != lastDebugLocation) {
-    auto offset = o.size();
-    auto& loc = iter->second;
-    sourceMapLocations.emplace_back(offset, &loc);
-    lastDebugLocation = loc;
+  if (iter != debugLocations.end()) {
+    writeDebugLocation(iter->second);
   }
 }
 
@@ -1058,11 +1064,19 @@ void WasmBinaryBuilder::readFunctions() {
       throwError("empty function size");
     }
     endOfFunction = pos + size;
+
+    Function *func = new Function;
+    func->name = Name::fromInt(i);
+    currFunction = func;
+
+    readNextDebugLocation();
+
     auto type = functionTypes[i];
     if (debug) std::cerr << "reading " << i << std::endl;
-    std::vector<Type> params, vars;
+    func->type = type->name;
+    func->result = type->result;
     for (size_t j = 0; j < type->params.size(); j++) {
-      params.emplace_back(type->params[j]);
+      func->params.emplace_back(type->params[j]);
     }
     size_t numLocalTypes = getU32LEB();
     for (size_t t = 0; t < numLocalTypes; t++) {
@@ -1077,23 +1091,16 @@ void WasmBinaryBuilder::readFunctions() {
         throwError("too many locals, wasm VMs would not accept this binary");
       }
       while (num > 0) {
-        vars.emplace_back(type);
+        func->vars.push_back(type);
         num--;
       }
     }
-    auto func = Builder(wasm).makeFunction(
-      Name::fromInt(i),
-      std::move(params),
-      type->result,
-      std::move(vars)
-    );
-    func->type = type->name;
-    currFunction = func;
+    std::swap(func->prologLocation, debugLocation);
     {
       // process the function body
       if (debug) std::cerr << "processing function: " << i << std::endl;
       nextLabel = 0;
-      useDebugLocation = false;
+      debugLocation.clear();
       willBeIgnored = false;
       // process body
       assert(breakTargetNames.size() == 0);
@@ -1111,8 +1118,9 @@ void WasmBinaryBuilder::readFunctions() {
         throwError("binary offset at function exit not at expected location");
       }
     }
+    std::swap(func->epilogLocation, debugLocation);
     currFunction = nullptr;
-    useDebugLocation = false;
+    debugLocation.clear();
     functions.push_back(func);
   }
   if (debug) std::cerr << " end function bodies" << std::endl;
@@ -1268,26 +1276,38 @@ void WasmBinaryBuilder::readSourceMapHeader() {
 void WasmBinaryBuilder::readNextDebugLocation() {
   if (!sourceMap) return;
 
-  char ch;
-  *sourceMap >> ch;
-  if (ch == '\"') { // end of records
-    nextDebugLocation.first = 0;
-    return;
-  }
-  if (ch != ',') {
-    throw MapParseException("Unexpected delimiter");
-  }
+  while (nextDebugLocation.first && nextDebugLocation.first <= pos) {
+    if (nextDebugLocation.first < pos) {
+      std::cerr << "skipping debug location info for 0x";
+      std::cerr << std::hex << nextDebugLocation.first << std::dec << std::endl;
+    }
+    debugLocation.clear();
+    // use debugLocation only for function expressions
+    if (currFunction) {
+      debugLocation.insert(nextDebugLocation.second);
+    }
 
-  int32_t positionDelta = readBase64VLQ(*sourceMap);
-  uint32_t position = nextDebugLocation.first + positionDelta;
-  int32_t fileIndexDelta = readBase64VLQ(*sourceMap);
-  uint32_t fileIndex = nextDebugLocation.second.fileIndex + fileIndexDelta;
-  int32_t lineNumberDelta = readBase64VLQ(*sourceMap);
-  uint32_t lineNumber = nextDebugLocation.second.lineNumber + lineNumberDelta;
-  int32_t columnNumberDelta = readBase64VLQ(*sourceMap);
-  uint32_t columnNumber = nextDebugLocation.second.columnNumber + columnNumberDelta;
+    char ch;
+    *sourceMap >> ch;
+    if (ch == '\"') { // end of records
+      nextDebugLocation.first = 0;
+      break;
+    }
+    if (ch != ',') {
+      throw MapParseException("Unexpected delimiter");
+    }
 
-  nextDebugLocation = { position, { fileIndex, lineNumber, columnNumber } };
+    int32_t positionDelta = readBase64VLQ(*sourceMap);
+    uint32_t position = nextDebugLocation.first + positionDelta;
+    int32_t fileIndexDelta = readBase64VLQ(*sourceMap);
+    uint32_t fileIndex = nextDebugLocation.second.fileIndex + fileIndexDelta;
+    int32_t lineNumberDelta = readBase64VLQ(*sourceMap);
+    uint32_t lineNumber = nextDebugLocation.second.lineNumber + lineNumberDelta;
+    int32_t columnNumberDelta = readBase64VLQ(*sourceMap);
+    uint32_t columnNumber = nextDebugLocation.second.columnNumber + columnNumberDelta;
+
+    nextDebugLocation = { position, { fileIndex, lineNumber, columnNumber } };
+  }
 }
 
 Expression* WasmBinaryBuilder::readExpression() {
@@ -1345,6 +1365,7 @@ void WasmBinaryBuilder::processExpressions() {
       auto peek = input[pos];
       if (peek == BinaryConsts::End || peek == BinaryConsts::Else) {
         if (debug) std::cerr << "== processExpressions finished with unreachable" << std::endl;
+        readNextDebugLocation();
         lastSeparator = BinaryConsts::ASTNodes(peek);
         pos++;
         return;
@@ -1654,15 +1675,10 @@ BinaryConsts::ASTNodes WasmBinaryBuilder::readExpression(Expression*& curr) {
     throwError("Reached function end without seeing End opcode");
   }
   if (debug) std::cerr << "zz recurse into " << ++depth << " at " << pos << std::endl;
-  if (nextDebugLocation.first) {
-    while (nextDebugLocation.first && nextDebugLocation.first <= pos) {
-      if (nextDebugLocation.first < pos) {
-        std::cerr << "skipping debug location info for " << nextDebugLocation.first << std::endl;
-      }
-      debugLocation = nextDebugLocation.second;
-      useDebugLocation = currFunction != NULL; // using only for function expressions
-      readNextDebugLocation();
-    }
+  readNextDebugLocation();
+  std::set<Function::DebugLocation> currDebugLocation;
+  if (debugLocation.size()) {
+    currDebugLocation.insert(*debugLocation.begin());
   }
   uint8_t code = getInt8();
   if (debug) std::cerr << "readExpression seeing " << (int)code << std::endl;
@@ -1710,8 +1726,8 @@ BinaryConsts::ASTNodes WasmBinaryBuilder::readExpression(Expression*& curr) {
       break;
     }
   }
-  if (useDebugLocation && curr) {
-    currFunction->debugLocations[curr] = debugLocation;
+  if (curr && currDebugLocation.size()) {
+    currFunction->debugLocations[curr] = *currDebugLocation.begin();
   }
   if (debug) std::cerr << "zz recurse from " << depth-- << " at " << pos << std::endl;
   return BinaryConsts::ASTNodes(code);
@@ -1764,13 +1780,18 @@ void WasmBinaryBuilder::visitBlock(Block* curr) {
     curr->name = getNextLabel();
     breakStack.push_back({curr->name, curr->type != none});
     stack.push_back(curr);
-    if (getInt8() == BinaryConsts::Block) {
+    auto peek = input[pos];
+    if (peek == BinaryConsts::Block) {
       // a recursion
+      readNextDebugLocation();
       curr = allocator.alloc<Block>();
+      pos++;
+      if (debugLocation.size()) {
+        currFunction->debugLocations[curr] = *debugLocation.begin();
+      }
       continue;
     } else {
       // end of recursion
-      ungetInt8();
       break;
     }
   }