handle loops in effect analyzer

1 files changed, 19 insertions, 8 deletions

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");