DWARF: Track the positions of 'end', 'else', 'catch' binary locations (#2603)

Control flow structures have those in addition to the normal span of
(start, end), and we need to track them too.

Tracking them during reading requires us to track control flow
structures while parsing, so that we can know to which structure
an end/else/catch refers to.

We track these locations using a map on the side of instruction
to its "extra" locations. That avoids increasing the size of the
tracking info for the much more common non-control flow
instructions.

Note that there is one more 'end' location, that of the function
(not referring to any instruction). I left that to a later PR to
not increase this one too much.

1 files changed, 58 insertions, 13 deletions

diff --git a/src/wasm/wasm-debug.cpp b/src/wasm/wasm-debug.cpp
index a8b70247f..aa596aa61 100644
--- a/src/wasm/wasm-debug.cpp
+++ b/src/wasm/wasm-debug.cpp
@@ -337,21 +337,28 @@ struct AddrExprMap {
   std::unordered_map<BinaryLocation, Expression*> startMap;
   std::unordered_map<BinaryLocation, Expression*> endMap;
 
+  // Some instructions have extra binary locations, like the else and end in
+  // and if. Track those separately, including their expression and their id
+  // ("else", "end", etc.), as they are rare, and we don't want to
+  // bloat the common case which is represented in the earlier maps.
+  struct ExtraInfo {
+    Expression* expr;
+    BinaryLocations::DelimiterId id;
+  };
+  std::unordered_map<BinaryLocation, ExtraInfo> extraMap;
+
   // Construct the map from the binaryLocations loaded from the wasm.
   AddrExprMap(const Module& wasm) {
     for (auto& func : wasm.functions) {
       for (auto pair : func->expressionLocations) {
         add(pair.first, pair.second);
       }
+      for (auto pair : func->delimiterLocations) {
+        add(pair.first, pair.second);
+      }
     }
   }
 
-  // Construct the map from new binaryLocations just written
-  AddrExprMap(const BinaryLocations& newLocations) {
-    for (auto pair : newLocations.expressions) {
-      add(pair.first, pair.second);
-    }
-  }
 
   Expression* getStart(BinaryLocation addr) const {
     auto iter = startMap.find(addr);
@@ -369,13 +376,30 @@ struct AddrExprMap {
     return nullptr;
   }
 
+  ExtraInfo getExtra(BinaryLocation addr) const {
+    auto iter = extraMap.find(addr);
+    if (iter != extraMap.end()) {
+      return iter->second;
+    }
+    return ExtraInfo{nullptr, BinaryLocations::Invalid};
+  }
+
 private:
-  void add(Expression* expr, BinaryLocations::Span span) {
+  void add(Expression* expr, const BinaryLocations::Span span) {
     assert(startMap.count(span.start) == 0);
     startMap[span.start] = expr;
     assert(endMap.count(span.end) == 0);
     endMap[span.end] = expr;
   }
+
+  void add(Expression* expr, const BinaryLocations::DelimiterLocations& extra) {
+    for (Index i = 0; i < extra.size(); i++) {
+      if (extra[i] != 0) {
+        assert(extraMap.count(extra[i]) == 0);
+        extraMap[extra[i]] = ExtraInfo{expr, BinaryLocations::DelimiterId(i)};
+      }
+    }
+  }
 };
 
 // Represents a mapping of addresses to expressions. Note that we use a single
@@ -423,7 +447,6 @@ struct LocationUpdater {
   const BinaryLocations& newLocations;
 
   AddrExprMap oldExprAddrMap;
-  AddrExprMap newExprAddrMap;
   FuncAddrMap oldFuncAddrMap;
 
   // TODO: for memory efficiency, we may want to do this in a streaming manner,
@@ -431,7 +454,7 @@ struct LocationUpdater {
 
   LocationUpdater(Module& wasm, const BinaryLocations& newLocations)
     : wasm(wasm), newLocations(newLocations), oldExprAddrMap(wasm),
-      newExprAddrMap(newLocations), oldFuncAddrMap(wasm) {}
+      oldFuncAddrMap(wasm) {}
 
   // Updates an expression's address. If there was never an instruction at that
   // address, or if there was but if that instruction no longer exists, return
@@ -448,15 +471,14 @@ struct LocationUpdater {
   }
 
   bool hasOldExprAddr(BinaryLocation oldAddr) const {
-    return oldExprAddrMap.getStart(oldAddr) != nullptr;
+    return oldExprAddrMap.getStart(oldAddr);
   }
 
   BinaryLocation getNewExprEndAddr(BinaryLocation oldAddr) const {
     if (auto* expr = oldExprAddrMap.getEnd(oldAddr)) {
       auto iter = newLocations.expressions.find(expr);
       if (iter != newLocations.expressions.end()) {
-        BinaryLocation newAddr = iter->second.end;
-        return newAddr;
+        return iter->second.end;
       }
     }
     return 0;
@@ -478,6 +500,25 @@ struct LocationUpdater {
     }
     return 0;
   }
+
+  bool hasOldFuncAddr(BinaryLocation oldAddr) const {
+    return oldFuncAddrMap.get(oldAddr);
+  }
+
+  BinaryLocation getNewExtraAddr(BinaryLocation oldAddr) const {
+    auto info = oldExprAddrMap.getExtra(oldAddr);
+    if (info.expr) {
+      auto iter = newLocations.delimiters.find(info.expr);
+      if (iter != newLocations.delimiters.end()) {
+        return iter->second[info.id];
+      }
+    }
+    return 0;
+  }
+
+  bool hasOldExtraAddr(BinaryLocation oldAddr) const {
+    return oldExprAddrMap.getExtra(oldAddr).expr;
+  }
 };
 
 static void updateDebugLines(llvm::DWARFYAML::Data& data,
@@ -511,11 +552,15 @@ static void updateDebugLines(llvm::DWARFYAML::Data& data,
         BinaryLocation newAddr = 0;
         if (locationUpdater.hasOldExprAddr(oldAddr)) {
           newAddr = locationUpdater.getNewExprAddr(oldAddr);
-        } else {
+        } else if (locationUpdater.hasOldFuncAddr(oldAddr)) {
           newAddr = locationUpdater.getNewFuncAddr(oldAddr);
+        } else if (locationUpdater.hasOldExtraAddr(oldAddr)) {
+          newAddr = locationUpdater.getNewExtraAddr(oldAddr);
         }
+        // TODO: last 'end' of a function
         if (newAddr) {
           newAddrs.push_back(newAddr);
+          assert(newAddrInfo.count(newAddr) == 0);
           newAddrInfo.emplace(newAddr, state);
           auto& updatedState = newAddrInfo.at(newAddr);
           // The only difference is the address TODO other stuff?