2 files changed, 36 insertions, 8 deletions

diff --git a/src/ast_utils.h b/src/ast_utils.h
index bf42d10e1..5ab427178 100644
--- a/src/ast_utils.h
+++ b/src/ast_utils.h
@@ -53,6 +53,7 @@ struct EffectAnalyzer : public WasmWalker<EffectAnalyzer> {
   bool accessesLocal() { return readsLocal || writesLocal; }
   bool accessesMemory() { return calls || readsMemory || writesMemory; }
   bool hasSideEffects() { return calls || writesLocal || writesMemory; }
+  bool hasAnything() { return branches || calls || readsLocal || writesLocal || readsMemory || writesMemory; }
 
   // checks if these effects would invalidate another set (e.g., if we write, we invalidate someone that reads, they can't be moved past us)
   bool invalidates(EffectAnalyzer& other) {
@@ -61,7 +62,23 @@ struct EffectAnalyzer : public WasmWalker<EffectAnalyzer> {
                     || (accessesMemory() && (other.writesMemory || other.calls)) || (accessesLocal() && other.writesLocal);
   }
 
+  // the checks above happen after the node's children were processed, in the order of execution
+  // we must also check for control flow that happens before the children, i.e., loops
+  bool checkPre(Expression* curr) {
+    if (curr->is<Loop>()) {
+      branches = true;
+      return true;
+    }
+    return false;
+  }
+
+  bool checkPost(Expression* curr) {
+    visit(curr);
+    return hasAnything();
+  }
+
   void visitBlock(Block *curr) { branches = true; }
+  void visitLoop(Loop *curr) { branches = true; }
   void visitIf(If *curr) { branches = true; }
   void visitBreak(Break *curr) { branches = true; }
   void visitSwitch(Switch *curr) { branches = true; }
diff --git a/src/passes/SimplifyLocals.cpp b/src/passes/SimplifyLocals.cpp
index a0f843bbd..cd6cb2ca7 100644
--- a/src/passes/SimplifyLocals.cpp
+++ b/src/passes/SimplifyLocals.cpp
@@ -78,14 +78,7 @@ struct SimplifyLocals : public WalkerPass<FastExecutionWalker<SimplifyLocals>> {
     }
   }
 
-  void walk(Expression*& curr) override {
-    if (!curr) return;
-
-    FastExecutionWalker::walk(curr);
-
-    // invalidations TODO: this is O(n^2) in sinkables
-    EffectAnalyzer effects;
-    effects.visit(curr);
+  void checkInvalidations(EffectAnalyzer& effects) {
     std::vector<Name> invalidated;
     for (auto& sinkable : sinkables) {
       if (effects.invalidates(sinkable.second.effects)) {
@@ -96,6 +89,24 @@ struct SimplifyLocals : public WalkerPass<FastExecutionWalker<SimplifyLocals>> {
       sinkables.erase(name);
     }
   }
+
+  void walk(Expression*& curr) override {
+    if (!curr) return;
+
+    EffectAnalyzer effects;
+
+    if (effects.checkPre(curr)) {
+      checkInvalidations(effects);
+    }
+
+    FastExecutionWalker::walk(curr);
+
+    // TODO: this is O(n^2) in sinkables
+
+    if (effects.checkPost(curr)) {
+      checkInvalidations(effects);
+    }
+  }
 };
 
 static RegisterPass<SimplifyLocals> registerPass("simplify-locals", "miscellaneous locals-related optimizations");