5 files changed, 106 insertions, 45 deletions

diff --git a/src/parser/wat-parser.cpp b/src/parser/wat-parser.cpp
index 5503a8295..7f6dd2975 100644
--- a/src/parser/wat-parser.cpp
+++ b/src/parser/wat-parser.cpp
@@ -94,52 +94,8 @@ Result<> parseDefs(Ctx& ctx,
 void propagateDebugLocations(Module& wasm) {
   // Copy debug locations from parents or previous siblings to expressions that
   // do not already have their own debug locations.
-  struct Propagator : WalkerPass<ExpressionStackWalker<Propagator>> {
-    using Super = WalkerPass<ExpressionStackWalker<Propagator>>;
-    bool isFunctionParallel() override { return true; }
-    bool modifiesBinaryenIR() override { return false; }
-    bool requiresNonNullableLocalFixups() override { return false; }
-    void runOnFunction(Module* module, Function* func) override {
-      if (!func->debugLocations.empty()) {
-        Super::runOnFunction(module, func);
-      }
-    }
-
-    // Unannotated instructions inherit either their previous sibling's location
-    // or their parent's location. Look up whichever is current for a given
-    // parent.
-    std::unordered_map<Expression*, Function::DebugLocation> parentDefaults;
-
-    static void doPreVisit(Propagator* self, Expression** currp) {
-      Super::doPreVisit(self, currp);
-      auto* curr = *currp;
-      auto& locs = self->getFunction()->debugLocations;
-      auto& parentDefaults = self->parentDefaults;
-      if (auto it = locs.find(curr); it != locs.end()) {
-        // Children will inherit this location.
-        parentDefaults[curr] = it->second;
-        if (auto* parent = self->getParent()) {
-          // Subsequent siblings will inherit this location.
-          parentDefaults[parent] = it->second;
-        }
-      } else {
-        // No annotation, see if we should inherit one.
-        if (auto* parent = self->getParent()) {
-          if (auto defaultIt = parentDefaults.find(parent);
-              defaultIt != parentDefaults.end()) {
-            // We have a default to inherit. Our children will inherit it, too.
-            locs[curr] = parentDefaults[curr] = defaultIt->second;
-          }
-        }
-      }
-    }
-
-    std::unique_ptr<Pass> create() override {
-      return std::make_unique<Propagator>();
-    }
-  };
   PassRunner runner(&wasm);
-  runner.add(std::make_unique<Propagator>());
+  runner.add("propagate-debug-locs");
   runner.run();
 }
 
diff --git a/src/passes/CMakeLists.txt b/src/passes/CMakeLists.txt
index 4cb696c21..1e5540ab0 100644
--- a/src/passes/CMakeLists.txt
+++ b/src/passes/CMakeLists.txt
@@ -29,6 +29,7 @@ set(passes_SOURCES
   DeadArgumentElimination.cpp
   DeadCodeElimination.cpp
   DeAlign.cpp
+  DebugLocationPropagation.cpp
   DeNaN.cpp
   Directize.cpp
   DuplicateImportElimination.cpp
diff --git a/src/passes/DebugLocationPropagation.cpp b/src/passes/DebugLocationPropagation.cpp
new file mode 100644
index 000000000..07ae53faa
--- /dev/null
+++ b/src/passes/DebugLocationPropagation.cpp
@@ -0,0 +1,99 @@
+/*
+ * 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.
+ */
+
+//
+// DebugLocationPropagation aim to pass debug location from parents or
+// previous siblings to expression which has no debug location. This is
+// useful for compilers that use Binaryen API to generate WebAssembly modules.
+//
+
+#include "pass.h"
+#include "wasm-traversal.h"
+#include "wasm.h"
+#include <cassert>
+#include <unordered_map>
+
+namespace wasm {
+
+struct DebugLocationPropagation
+  : WalkerPass<PostWalker<DebugLocationPropagation>> {
+
+  // The top element of this stack is the previous sibling or parent of the
+  // current expression in `doPrevisit`. To maintain this invariant, expressions
+  // are only popped once we are done visiting their parents.
+  ExpressionStack expressionStack;
+
+  using Super = WalkerPass<PostWalker<DebugLocationPropagation>>;
+  bool isFunctionParallel() override { return true; }
+  bool modifiesBinaryenIR() override { return false; }
+  bool requiresNonNullableLocalFixups() override { return false; }
+  void runOnFunction(Module* module, Function* func) override {
+    if (!func->debugLocations.empty()) {
+      Super::runOnFunction(module, func);
+    }
+  }
+
+  Expression* getPrevious() {
+    if (expressionStack.empty()) {
+      return nullptr;
+    }
+    assert(expressionStack.size() >= 1);
+    return expressionStack[expressionStack.size() - 1];
+  }
+
+  static void doPreVisit(DebugLocationPropagation* self, Expression** currp) {
+    auto* curr = *currp;
+    auto& locs = self->getFunction()->debugLocations;
+    auto& expressionStack = self->expressionStack;
+    if (locs.find(curr) == locs.end()) {
+      // No debug location, see if we should inherit one.
+      if (auto* previous = self->getPrevious()) {
+        if (auto it = locs.find(previous); it != locs.end()) {
+          locs[curr] = it->second;
+        }
+      }
+    }
+    expressionStack.push_back(curr);
+  }
+
+  static void doPostVisit(DebugLocationPropagation* self, Expression** currp) {
+    auto& exprStack = self->expressionStack;
+    while (exprStack.back() != *currp) {
+      // pop all the child expressions and keep current expression in stack.
+      exprStack.pop_back();
+    }
+    // the stack should never be empty
+    assert(!exprStack.empty());
+  }
+
+  static void scan(DebugLocationPropagation* self, Expression** currp) {
+    self->pushTask(DebugLocationPropagation::doPostVisit, currp);
+
+    PostWalker<DebugLocationPropagation>::scan(self, currp);
+
+    self->pushTask(DebugLocationPropagation::doPreVisit, currp);
+  }
+
+  std::unique_ptr<Pass> create() override {
+    return std::make_unique<DebugLocationPropagation>();
+  }
+};
+
+Pass* createDebugLocationPropagationPass() {
+  return new DebugLocationPropagation();
+}
+
+} // namespace wasm
diff --git a/src/passes/pass.cpp b/src/passes/pass.cpp
index 19ddaf2d4..4a4a9c558 100644
--- a/src/passes/pass.cpp
+++ b/src/passes/pass.cpp
@@ -126,6 +126,10 @@ void PassRegistry::registerPasses() {
   registerPass("dealign",
                "forces all loads and stores to have alignment 1",
                createDeAlignPass);
+  registerPass(
+    "propagate-debug-locs",
+    "propagate debug location from parents or previous siblings to child nodes",
+    createDebugLocationPropagationPass);
   registerPass("denan",
                "instrument the wasm to convert NaNs into 0 at runtime",
                createDeNaNPass);
diff --git a/src/passes/passes.h b/src/passes/passes.h
index 23a9ea70b..8695539b3 100644
--- a/src/passes/passes.h
+++ b/src/passes/passes.h
@@ -38,6 +38,7 @@ Pass* createDataFlowOptsPass();
 Pass* createDeadCodeEliminationPass();
 Pass* createDeNaNPass();
 Pass* createDeAlignPass();
+Pass* createDebugLocationPropagationPass();
 Pass* createDirectizePass();
 Pass* createDiscardGlobalEffectsPass();
 Pass* createDWARFDumpPass();