1 files changed, 65 insertions, 15 deletions

diff --git a/src/wasm/wasm-binary.cpp b/src/wasm/wasm-binary.cpp
index 71dd0fc21..fa90532d9 100644
--- a/src/wasm/wasm-binary.cpp
+++ b/src/wasm/wasm-binary.cpp
@@ -72,6 +72,14 @@ void WasmBinaryWriter::write() {
     writeSourceMapEpilog();
   }
 
+#ifdef BUILD_LLVM_DWARF
+  // Update DWARF user sections after writing the data referred to by them
+  // (function bodies), and before writing the user sections themselves.
+  if (Debug::hasDWARFSections(*wasm)) {
+    Debug::writeDWARFSections(*wasm, binaryLocations);
+  }
+#endif
+
   writeLateUserSections();
   writeFeaturesSection();
 
@@ -109,6 +117,7 @@ template<typename T> int32_t WasmBinaryWriter::startSection(T code) {
   if (sourceMap) {
     sourceMapLocationsSizeAtSectionStart = sourceMapLocations.size();
   }
+  binaryLocationsSizeAtSectionStart = binaryLocations.size();
   return writeU32LEBPlaceholder(); // section size to be filled in later
 }
 
@@ -116,22 +125,44 @@ void WasmBinaryWriter::finishSection(int32_t start) {
   // section size does not include the reserved bytes of the size field itself
   int32_t size = o.size() - start - MaxLEB32Bytes;
   auto sizeFieldSize = o.writeAt(start, U32LEB(size));
-  if (sizeFieldSize != MaxLEB32Bytes) {
+  // We can move things back if the actual LEB for the size doesn't use the
+  // maximum 5 bytes. In that case we need to adjust offsets after we move
+  // things backwards.
+  auto adjustmentForLEBShrinking = MaxLEB32Bytes - sizeFieldSize;
+  if (adjustmentForLEBShrinking) {
     // we can save some room, nice
     assert(sizeFieldSize < MaxLEB32Bytes);
     std::move(&o[start] + MaxLEB32Bytes,
               &o[start] + MaxLEB32Bytes + size,
               &o[start] + sizeFieldSize);
-    auto adjustment = MaxLEB32Bytes - sizeFieldSize;
-    o.resize(o.size() - adjustment);
+    o.resize(o.size() - adjustmentForLEBShrinking);
     if (sourceMap) {
       for (auto i = sourceMapLocationsSizeAtSectionStart;
            i < sourceMapLocations.size();
            ++i) {
-        sourceMapLocations[i].first -= adjustment;
+        sourceMapLocations[i].first -= adjustmentForLEBShrinking;
       }
     }
   }
+
+  if (binaryLocationsSizeAtSectionStart != binaryLocations.size()) {
+    // We added the binary locations, adjust them: they must be relative
+    // to the code section.
+    assert(binaryLocationsSizeAtSectionStart == 0);
+    // The section type byte is right before the LEB for the size; we want
+    // 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;
+    }
+  }
 }
 
 int32_t
@@ -266,6 +297,7 @@ void WasmBinaryWriter::writeFunctions() {
   auto start = startSection(BinaryConsts::Section::Code);
   o << U32LEB(importInfo->getNumDefinedFunctions());
   ModuleUtils::iterDefinedFunctions(*wasm, [&](Function* func) {
+    assert(binaryLocationTrackedExpressionsForFunc.empty());
     size_t sourceMapLocationsSizeAtFunctionStart = sourceMapLocations.size();
     BYN_TRACE("write one at" << o.size() << std::endl);
     size_t sizePos = writeU32LEBPlaceholder();
@@ -284,22 +316,31 @@ void WasmBinaryWriter::writeFunctions() {
     BYN_TRACE("body size: " << size << ", writing at " << sizePos
                             << ", next starts at " << o.size() << "\n");
     auto sizeFieldSize = o.writeAt(sizePos, U32LEB(size));
-    if (sizeFieldSize != MaxLEB32Bytes) {
+    // We can move things back if the actual LEB for the size doesn't use the
+    // maximum 5 bytes. In that case we need to adjust offsets after we move
+    // things backwards.
+    auto adjustmentForLEBShrinking = MaxLEB32Bytes - sizeFieldSize;
+    if (adjustmentForLEBShrinking) {
       // we can save some room, nice
       assert(sizeFieldSize < MaxLEB32Bytes);
       std::move(&o[start], &o[start] + size, &o[sizePos] + sizeFieldSize);
-      auto adjustment = MaxLEB32Bytes - sizeFieldSize;
-      o.resize(o.size() - adjustment);
+      o.resize(o.size() - adjustmentForLEBShrinking);
       if (sourceMap) {
         for (auto i = sourceMapLocationsSizeAtFunctionStart;
              i < sourceMapLocations.size();
              ++i) {
-          sourceMapLocations[i].first -= adjustment;
+          sourceMapLocations[i].first -= adjustmentForLEBShrinking;
         }
       }
+      for (auto* curr : binaryLocationTrackedExpressionsForFunc) {
+        // We added the binary locations, adjust them: they must be relative
+        // to the code section.
+        binaryLocations[curr] -= adjustmentForLEBShrinking;
+      }
     }
     tableOfContents.functionBodies.emplace_back(
       func->name, sizePos + sizeFieldSize, size);
+    binaryLocationTrackedExpressionsForFunc.clear();
   });
   finishSection(start);
 }
@@ -649,10 +690,19 @@ void WasmBinaryWriter::writeDebugLocation(const Function::DebugLocation& loc) {
 }
 
 void WasmBinaryWriter::writeDebugLocation(Expression* curr, Function* func) {
-  auto& debugLocations = func->debugLocations;
-  auto iter = debugLocations.find(curr);
-  if (iter != debugLocations.end()) {
-    writeDebugLocation(iter->second);
+  if (sourceMap) {
+    auto& debugLocations = func->debugLocations;
+    auto iter = debugLocations.find(curr);
+    if (iter != debugLocations.end()) {
+      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();
+    binaryLocationTrackedExpressionsForFunc.push_back(curr);
   }
 }
 
@@ -809,6 +859,9 @@ void WasmBinaryBuilder::read() {
         readFunctionSignatures();
         break;
       case BinaryConsts::Section::Code:
+        if (DWARF) {
+          codeSectionLocation = pos;
+        }
         readFunctions();
         break;
       case BinaryConsts::Section::Export:
@@ -1288,9 +1341,6 @@ void WasmBinaryBuilder::readFunctionSignatures() {
 
 void WasmBinaryBuilder::readFunctions() {
   BYN_TRACE("== readFunctions\n");
-  if (DWARF) {
-    codeSectionLocation = pos;
-  }
   size_t total = getU32LEB();
   if (total != functionSignatures.size()) {
     throwError("invalid function section size, must equal types");