rewrite SimplifyLocals to use FastExecutionWalker

1 files changed, 71 insertions, 12 deletions

diff --git a/src/passes/SimplifyLocals.cpp b/src/passes/SimplifyLocals.cpp
index cbfc0dd66..98806410f 100644
--- a/src/passes/SimplifyLocals.cpp
+++ b/src/passes/SimplifyLocals.cpp
@@ -15,26 +15,85 @@
  */
 
 //
-// Miscellaneous locals-related optimizations
+// Locals-related optimizations
 //
+// This "sinks" set_locals, pushing them to the next get_local where possible
 
 #include <wasm.h>
+#include <wasm-traversal.h>
 #include <pass.h>
+#include <ast_utils.h>
 
 namespace wasm {
 
-struct SimplifyLocals : public WalkerPass<WasmWalker<SimplifyLocals>> {
+struct SimplifyLocals : public WalkerPass<FastExecutionWalker<SimplifyLocals>> {
+  struct SinkableInfo {
+    Expression** item;
+    EffectAnalyzer effects;
+
+    SinkableInfo(Expression** item) : item(item) {
+      effects.walk(*item);
+    }
+  };
+
+  // locals in current linear execution trace, which we try to sink
+  std::map<Name, SinkableInfo> sinkables;
+
+  void noteNonLinear() {
+    sinkables.clear();
+  }
+
   void visitBlock(Block *curr) {
-    // look for pairs of setlocal-getlocal, which can be just a setlocal (since it returns a value)
-    if (curr->list.size() == 0) return;
-    for (size_t i = 0; i < curr->list.size() - 1; i++) {
-      auto set = curr->list[i]->dyn_cast<SetLocal>();
-      if (!set) continue;
-      auto get = curr->list[i + 1]->dyn_cast<GetLocal>();
-      if (!get) continue;
-      if (set->name != get->name) continue;
-      curr->list.erase(curr->list.begin() + i + 1);
-      i -= 1;
+    // note locals, we can sink them from here TODO sink from elsewhere?
+    ExpressionList& list = curr->list;
+    for (size_t z = 0; z < list.size(); z++) {
+      walk(list[z]);
+      auto* item = list[z];
+      if (item->is<SetLocal>()) {
+        Name name = item->cast<SetLocal>()->name;
+        assert(sinkables.count(name) == 0);
+        sinkables.emplace(std::make_pair(name, SinkableInfo(&list[z])));
+      }
+    }
+  }
+
+  void visitGetLocal(GetLocal *curr) {
+    auto found = sinkables.find(curr->name);
+    if (found != sinkables.end()) {
+      // sink it, and nop the origin TODO: clean up nops
+      replaceCurrent(*found->second.item);
+      // reuse the getlocal that is dying
+      *found->second.item = curr;
+      ExpressionManipulator::nop(curr);
+      sinkables.erase(found);
+    }
+  }
+
+  void visitSetLocal(SetLocal *curr) {
+    walk(curr->value);
+    // if we are a potentially-sinkable thing, forget it - this
+    // write overrides the last TODO: optimizable
+    // TODO: if no get_locals left, can remove the set as well (== expressionizer in emscripten optimizer)
+    auto found = sinkables.find(curr->name);
+    if (found != sinkables.end()) {
+      sinkables.erase(found);
+    }
+  }
+
+  void walk(Expression*& curr) override {
+    FastExecutionWalker::walk(curr);
+
+    // invalidations TODO: this is O(n^2) in sinkables
+    EffectAnalyzer effects;
+    effects.walk(curr); // TODO: this, accumulated, is O(n * nesting) <= O(n^2)
+    std::vector<Name> invalidated;
+    for (auto& sinkable : sinkables) {
+      if (effects.invalidates(sinkable.second.effects)) {
+        invalidated.push_back(sinkable.first);
+      }
+    }
+    for (auto name : invalidated) {
+      sinkables.erase(name);
     }
   }
 };