1 files changed, 31 insertions, 18 deletions

diff --git a/src/wasm/wasm-binary.cpp b/src/wasm/wasm-binary.cpp
index 4b748a5af..9b1f24a0f 100644
--- a/src/wasm/wasm-binary.cpp
+++ b/src/wasm/wasm-binary.cpp
@@ -117,7 +117,7 @@ template<typename T> int32_t WasmBinaryWriter::startSection(T code) {
   if (sourceMap) {
     sourceMapLocationsSizeAtSectionStart = sourceMapLocations.size();
   }
-  binaryLocationsSizeAtSectionStart = binaryLocations.size();
+  binaryLocationsSizeAtSectionStart = binaryLocations.expressions.size();
   return writeU32LEBPlaceholder(); // section size to be filled in later
 }
 
@@ -145,7 +145,7 @@ void WasmBinaryWriter::finishSection(int32_t start) {
     }
   }
 
-  if (binaryLocationsSizeAtSectionStart != binaryLocations.size()) {
+  if (binaryLocationsSizeAtSectionStart != binaryLocations.expressions.size()) {
     // We added the binary locations, adjust them: they must be relative
     // to the code section.
     assert(binaryLocationsSizeAtSectionStart == 0);
@@ -153,14 +153,18 @@ void WasmBinaryWriter::finishSection(int32_t start) {
     // offsets that are relative to the body, which is after that section type
     // byte and the the size LEB.
     auto body = start + sizeFieldSize;
-    for (auto& pair : binaryLocations) {
-      // Offsets are relative to the body of the code section: after the
-      // section type byte and the size.
-      // Everything was moved by the adjustment, track that. After this,
-      // we are at the right absolute address.
-      pair.second -= adjustmentForLEBShrinking;
-      // We are relative to the section start.
-      pair.second -= body;
+    // Offsets are relative to the body of the code section: after the
+    // section type byte and the size.
+    // Everything was moved by the adjustment, track that. After this,
+    // we are at the right absolute address.
+    // We are relative to the section start.
+    auto totalAdjustment = adjustmentForLEBShrinking + body;
+    for (auto& pair : binaryLocations.expressions) {
+      pair.second -= totalAdjustment;
+    }
+    for (auto& pair : binaryLocations.functions) {
+      pair.second.start -= totalAdjustment;
+      pair.second.end -= totalAdjustment;
     }
   }
 }
@@ -335,9 +339,13 @@ void WasmBinaryWriter::writeFunctions() {
       for (auto* curr : binaryLocationTrackedExpressionsForFunc) {
         // We added the binary locations, adjust them: they must be relative
         // to the code section.
-        binaryLocations[curr] -= adjustmentForLEBShrinking;
+        binaryLocations.expressions[curr] -= adjustmentForLEBShrinking;
       }
     }
+    if (!binaryLocationTrackedExpressionsForFunc.empty()) {
+      binaryLocations.functions[func] = BinaryLocations::Span{
+        uint32_t(start - adjustmentForLEBShrinking), uint32_t(o.size())};
+    }
     tableOfContents.functionBodies.emplace_back(
       func->name, sizePos + sizeFieldSize, size);
     binaryLocationTrackedExpressionsForFunc.clear();
@@ -697,11 +705,10 @@ void WasmBinaryWriter::writeDebugLocation(Expression* curr, Function* func) {
       writeDebugLocation(iter->second);
     }
   }
-  // TODO: remove source map debugging support and refactor this method
-  // to something that directly thinks about DWARF, instead of indirectly
-  // looking at func->binaryLocations as a proxy for that etc.
-  if (func && !func->binaryLocations.empty()) {
-    binaryLocations[curr] = o.size();
+  // 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.
+  if (func && !func->expressionLocations.empty()) {
+    binaryLocations.expressions[curr] = o.size();
     binaryLocationTrackedExpressionsForFunc.push_back(curr);
   }
 }
@@ -1357,11 +1364,17 @@ void WasmBinaryBuilder::readFunctions() {
     }
     endOfFunction = pos + size;
 
-    Function* func = new Function;
+    auto* func = new Function;
     func->name = Name::fromInt(i);
     func->sig = functionSignatures[i];
     currFunction = func;
 
+    if (DWARF) {
+      func->funcLocation =
+        BinaryLocations::Span{uint32_t(pos - codeSectionLocation),
+                              uint32_t(pos - codeSectionLocation + size)};
+    }
+
     readNextDebugLocation();
 
     BYN_TRACE("reading " << i << std::endl);
@@ -2280,7 +2293,7 @@ BinaryConsts::ASTNodes WasmBinaryBuilder::readExpression(Expression*& curr) {
       currFunction->debugLocations[curr] = *currDebugLocation.begin();
     }
     if (DWARF && currFunction) {
-      currFunction->binaryLocations[curr] = startPos - codeSectionLocation;
+      currFunction->expressionLocations[curr] = startPos - codeSectionLocation;
     }
   }
   BYN_TRACE("zz recurse from " << depth-- << " at " << pos << std::endl);