1 files changed, 164 insertions, 211 deletions

diff --git a/src/ir/LocalGraph.cpp b/src/ir/LocalGraph.cpp
index 53c6a52e2..77f5906f0 100644
--- a/src/ir/LocalGraph.cpp
+++ b/src/ir/LocalGraph.cpp
@@ -20,250 +20,203 @@
 #include <wasm-printing.h>
 #include <ir/find_all.h>
 #include <ir/local-graph.h>
+#include <cfg/cfg-traversal.h>
 
 namespace wasm {
 
-LocalGraph::LocalGraph(Function* func, Module* module) {
-  walkFunctionInModule(func, module);
+namespace LocalGraphInternal {
+
+// A relevant action. Supports a get, a set, or an
+// "other" which can be used for other purposes, to mark
+// their position in a block
+struct Action {
+  enum What {
+    Get = 0,
+    Set = 1
+  };
+  What what;
+  Index index; // the local index read or written
+  Expression* expr; // the expression itself
+
+  Action(What what, Index index, Expression* expr) : what(what), index(index), expr(expr) {
+    if (what == Get) assert(expr->is<GetLocal>());
+    if (what == Set) assert(expr->is<SetLocal>());
+  }
 
-#ifdef LOCAL_GRAPH_DEBUG
-  std::cout << "LocalGraph::dump\n";
-  for (auto& pair : getSetses) {
-    auto* get = pair.first;
-    auto& sets = pair.second;
-    std::cout << "GET\n" << get << " is influenced by\n";
-    for (auto* set : sets) {
-      std::cout << set << '\n';
+  bool isGet() { return what == Get; }
+  bool isSet() { return what == Set; }
+};
+
+// information about a basic block
+struct Info {
+  std::vector<Action> actions; // actions occurring in this block
+  std::vector<SetLocal*> lastSets; // for each index, the last set_local for it
+
+  void dump(Function* func) {
+    if (actions.empty()) return;
+    std::cout << "    actions:\n";
+    for (auto& action : actions) {
+      std::cout << "      " << (action.isGet() ? "get" : "set") << " " << func->getLocalName(action.index) << "\n";
     }
   }
-#endif
-}
+};
 
-void LocalGraph::computeInfluences() {
-  for (auto& pair : locations) {
-    auto* curr = pair.first;
-    if (auto* set = curr->dynCast<SetLocal>()) {
-      FindAll<GetLocal> findAll(set->value);
-      for (auto* get : findAll.list) {
-        getInfluences[get].insert(set);
-      }
-    } else {
-      auto* get = curr->cast<GetLocal>();
-      for (auto* set : getSetses[get]) {
-        setInfluences[set].insert(get);
-      }
-    }
+// flow helper class. flows the gets to their sets
+
+struct Flower : public CFGWalker<Flower, Visitor<Flower>, Info> {
+  LocalGraph::GetSetses& getSetses;
+  LocalGraph::Locations& locations;
+
+  Flower(LocalGraph::GetSetses& getSetses, LocalGraph::Locations& locations, Function* func) : getSetses(getSetses), locations(locations) {
+    setFunction(func);
+    // create the CFG by walking the IR
+    CFGWalker<Flower, Visitor<Flower>, Info>::doWalkFunction(func);
+    // flow gets across blocks
+    flow(func);
   }
-}
 
-void LocalGraph::doWalkFunction(Function* func) {
-  numLocals = func->getNumLocals();
-  if (numLocals == 0) return; // nothing to do
-  // We begin with each param being assigned from the incoming value, and the zero-init for the locals,
-  // so the initial state is the identity permutation
-  currMapping.resize(numLocals);
-  for (auto& set : currMapping) {
-    set = { nullptr };
+  BasicBlock* makeBasicBlock() {
+    auto* ret = new BasicBlock();
+    auto& lastSets = ret->contents.lastSets;
+    lastSets.resize(getFunction()->getNumLocals());
+    std::fill(lastSets.begin(), lastSets.end(), nullptr);
+    return ret;
   }
-  PostWalker<LocalGraph>::walk(func->body);
-}
 
-// control flow
+  // cfg traversal work
 
-void LocalGraph::visitBlock(Block* curr) {
-  if (curr->name.is() && breakMappings.find(curr->name) != breakMappings.end()) {
-    auto& infos = breakMappings[curr->name];
-    infos.emplace_back(std::move(currMapping));
-    currMapping = std::move(merge(infos));
-    breakMappings.erase(curr->name);
+  static void doVisitGetLocal(Flower* self, Expression** currp) {
+    auto* curr = (*currp)->cast<GetLocal>();
+     // if in unreachable code, skip
+    if (!self->currBasicBlock) return;
+    self->currBasicBlock->contents.actions.emplace_back(Action::Get, curr->index, curr);
+    self->locations[curr] = currp;
   }
-}
 
-void LocalGraph::finishIf() {
-  // that's it for this if, merge
-  std::vector<Mapping> breaks;
-  breaks.emplace_back(std::move(currMapping));
-  breaks.emplace_back(std::move(mappingStack.back()));
-  mappingStack.pop_back();
-  currMapping = std::move(merge(breaks));
-}
-
-void LocalGraph::afterIfCondition(LocalGraph* self, Expression** currp) {
-  self->mappingStack.push_back(self->currMapping);
-}
-void LocalGraph::afterIfTrue(LocalGraph* self, Expression** currp) {
-  auto* curr = (*currp)->cast<If>();
-  if (curr->ifFalse) {
-    auto afterCondition = std::move(self->mappingStack.back());
-    self->mappingStack.back() = std::move(self->currMapping);
-    self->currMapping = std::move(afterCondition);
-  } else {
-    self->finishIf();
+  static void doVisitSetLocal(Flower* self, Expression** currp) {
+    auto* curr = (*currp)->cast<SetLocal>();
+    // if in unreachable code, skip
+    if (!self->currBasicBlock) return;
+    self->currBasicBlock->contents.actions.emplace_back(Action::Set, curr->index, curr);
+    self->currBasicBlock->contents.lastSets[curr->index] = curr;
+    self->locations[curr] = currp;
   }
-}
-void LocalGraph::afterIfFalse(LocalGraph* self, Expression** currp) {
-  self->finishIf();
-}
-void LocalGraph::beforeLoop(LocalGraph* self, Expression** currp) {
-  // save the state before entering the loop, for calculation later of the merge at the loop top
-  self->mappingStack.push_back(self->currMapping);
-  self->loopGetStack.push_back({});
-}
-void LocalGraph::visitLoop(Loop* curr) {
-  if (curr->name.is() && breakMappings.find(curr->name) != breakMappings.end()) {
-    auto& infos = breakMappings[curr->name];
-    infos.emplace_back(std::move(mappingStack.back()));
-    auto before = infos.back();
-    auto& merged = merge(infos);
-    // every local we created a phi for requires us to update get_local operations in
-    // the loop - the branch back has means that gets in the loop have potentially
-    // more sets reaching them.
-    // we can detect this as follows: if a get of oldIndex has the same sets
-    // as the sets at the entrance to the loop, then it is affected by the loop
-    // header sets, and we can add to there sets that looped back
-    auto linkLoopTop = [&](Index i, Sets& getSets) {
-      auto& beforeSets = before[i];
-      if (getSets.size() < beforeSets.size()) {
-        // the get trivially has fewer sets, so it overrode the loop entry sets
-        return;
+
+  void flow(Function* func) {
+    auto numLocals = func->getNumLocals();
+    std::vector<std::vector<GetLocal*>> allGets;
+    allGets.resize(numLocals);
+    std::unordered_set<BasicBlock*> seen;
+    std::vector<BasicBlock*> work;
+    for (auto& block : basicBlocks) {
+#ifdef LOCAL_GRAPH_DEBUG
+      std::cout << "basic block " << block.get() << " :\n";
+      for (auto& action : block->contents.actions) {
+        std::cout << "  action: " << action.expr << '\n';
       }
-      if (!std::includes(getSets.begin(), getSets.end(),
-                         beforeSets.begin(), beforeSets.end())) {
-        // the get has not the same sets as in the loop entry
-        return;
+      for (auto* lastSet : block->contents.lastSets) {
+        std::cout << "  last set " << lastSet << '\n';
       }
-      // the get has the entry sets, so add any new ones
-      for (auto* set : merged[i]) {
-        getSets.insert(set);
+#endif
+      // go through the block, finding each get and adding it to its index,
+      // and seeing how sets affect that
+      auto& actions = block->contents.actions;
+      // move towards the front, handling things as we go
+      for (int i = int(actions.size()) - 1; i >= 0; i--) {
+        auto& action = actions[i];
+        auto index = action.index;
+        if (action.isGet()) {
+          allGets[index].push_back(action.expr->cast<GetLocal>());
+        } else {
+          // this set is the only set for all those gets
+          auto* set = action.expr->cast<SetLocal>();
+          auto& gets = allGets[index];
+          for (auto* get : gets) {
+            getSetses[get].insert(set);
+          }
+          gets.clear();
+        }
       }
-    };
-    auto& gets = loopGetStack.back();
-    for (auto* get : gets) {
-      linkLoopTop(get->index, getSetses[get]);
-    }
-    // and the same for the loop fallthrough: any local that still has the
-    // entry sets should also have the loop-back sets as well
-    for (Index i = 0; i < numLocals; i++) {
-      linkLoopTop(i, currMapping[i]);
-    }
-    // finally, breaks still in flight must be updated too
-    for (auto& iter : breakMappings) {
-      auto name = iter.first;
-      if (name == curr->name) continue; // skip our own (which is still in use)
-      auto& mappings = iter.second;
-      for (auto& mapping : mappings) {
-        for (Index i = 0; i < numLocals; i++) {
-          linkLoopTop(i, mapping[i]);
+      // if anything is left, we must flow it back through other blocks. we
+      // can do that for all gets as a whole, they will get the same results
+      for (Index index = 0; index < numLocals; index++) {
+        auto& gets = allGets[index];
+        if (gets.empty()) continue;
+        work.push_back(block.get());
+        seen.clear();
+        // note that we may need to revisit the later parts of this initial
+        // block, if we are in a loop, so don't mark it as seen
+        while (!work.empty()) {
+          auto* curr = work.back();
+          work.pop_back();
+          // we have gone through this block; now we must handle flowing to
+          // the inputs
+          if (curr->in.empty()) {
+            if (curr == entry) {
+              // these receive a param or zero init value
+              for (auto* get : gets) {
+                getSetses[get].insert(nullptr);
+              }
+            }
+          } else {
+            for (auto* pred : curr->in) {
+              if (seen.count(pred)) continue;
+              seen.insert(pred);
+              auto* lastSet = pred->contents.lastSets[index];
+              if (lastSet) {
+                // there is a set here, apply it, and stop the flow
+                for (auto* get : gets) {
+                  getSetses[get].insert(lastSet);
+                }
+              } else {
+                // keep on flowing
+                work.push_back(pred);
+              }
+            }
+          }
         }
+        gets.clear();
       }
     }
-    // now that we are done with using the mappings, erase our own
-    breakMappings.erase(curr->name);
-  }
-  mappingStack.pop_back();
-  loopGetStack.pop_back();
-}
-void LocalGraph::visitBreak(Break* curr) {
-  if (curr->condition) {
-    breakMappings[curr->name].emplace_back(currMapping);
-  } else {
-    breakMappings[curr->name].emplace_back(std::move(currMapping));
-    setUnreachable(currMapping);
-  }
-}
-void LocalGraph::visitSwitch(Switch* curr) {
-  std::set<Name> all;
-  for (auto target : curr->targets) {
-    all.insert(target);
   }
-  all.insert(curr->default_);
-  for (auto target : all) {
-    breakMappings[target].emplace_back(currMapping);
-  }
-  setUnreachable(currMapping);
-}
-void LocalGraph::visitReturn(Return *curr) {
-  setUnreachable(currMapping);
-}
-void LocalGraph::visitUnreachable(Unreachable *curr) {
-  setUnreachable(currMapping);
-}
+};
 
-// local usage
+} // namespace LocalGraphInternal
 
-void LocalGraph::visitGetLocal(GetLocal* curr) {
-  assert(currMapping.size() == numLocals);
-  assert(curr->index < numLocals);
-  for (auto& loopGets : loopGetStack) {
-    loopGets.push_back(curr);
-  }
-  // current sets are our sets
-  getSetses[curr] = currMapping[curr->index];
-  locations[curr] = getCurrentPointer();
-}
-void LocalGraph::visitSetLocal(SetLocal* curr) {
-  assert(currMapping.size() == numLocals);
-  assert(curr->index < numLocals);
-  // current sets are just this set
-  currMapping[curr->index] = { curr }; // TODO optimize?
-  locations[curr] = getCurrentPointer();
-}
+// LocalGraph implementation
 
-// traversal
+LocalGraph::LocalGraph(Function* func) {
+  LocalGraphInternal::Flower flower(getSetses, locations, func);
 
-void LocalGraph::scan(LocalGraph* self, Expression** currp) {
-  if (auto* iff = (*currp)->dynCast<If>()) {
-    // if needs special handling
-    if (iff->ifFalse) {
-      self->pushTask(LocalGraph::afterIfFalse, currp);
-      self->pushTask(LocalGraph::scan, &iff->ifFalse);
+#ifdef LOCAL_GRAPH_DEBUG
+  std::cout << "LocalGraph::dump\n";
+  for (auto& pair : getSetses) {
+    auto* get = pair.first;
+    auto& sets = pair.second;
+    std::cout << "GET\n" << get << " is influenced by\n";
+    for (auto* set : sets) {
+      std::cout << set << '\n';
     }
-    self->pushTask(LocalGraph::afterIfTrue, currp);
-    self->pushTask(LocalGraph::scan, &iff->ifTrue);
-    self->pushTask(LocalGraph::afterIfCondition, currp);
-    self->pushTask(LocalGraph::scan, &iff->condition);
-  } else {
-    PostWalker<LocalGraph>::scan(self, currp);
   }
-
-  // loops need pre-order visiting too
-  if ((*currp)->is<Loop>()) {
-    self->pushTask(LocalGraph::beforeLoop, currp);
-  }
-}
-
-// helpers
-
-void LocalGraph::setUnreachable(Mapping& mapping) {
-  mapping.resize(numLocals); // may have been emptied by a move
-  mapping[0].clear();
-}
-
-bool LocalGraph::isUnreachable(Mapping& mapping) {
-  // we must have some set for each index, if only the zero init, so empty means we emptied it for unreachable code
-  return mapping[0].empty();
+  std::cout << "total locations: " << locations.size() << '\n';
+#endif
 }
 
-// merges a bunch of infos into one.
-// if we need phis, writes them into the provided vector. the caller should
-// ensure those are placed in the right location
-LocalGraph::Mapping& LocalGraph::merge(std::vector<Mapping>& mappings) {
-  assert(mappings.size() > 0);
-  auto& out = mappings[0];
-  if (mappings.size() == 1) {
-    return out;
-  }
-  // merge into the first
-  for (Index j = 1; j < mappings.size(); j++) {
-    auto& other = mappings[j];
-    for (Index i = 0; i < numLocals; i++) {
-      auto& outSets = out[i];
-      for (auto* set : other[i]) {
-        outSets.insert(set);
+void LocalGraph::computeInfluences() {
+  for (auto& pair : locations) {
+    auto* curr = pair.first;
+    if (auto* set = curr->dynCast<SetLocal>()) {
+      FindAll<GetLocal> findAll(set->value);
+      for (auto* get : findAll.list) {
+        getInfluences[get].insert(set);
+      }
+    } else {
+      auto* get = curr->cast<GetLocal>();
+      for (auto* set : getSetses[get]) {
+        setInfluences[set].insert(get);
       }
     }
   }
-  return out;
 }
 
 } // namespace wasm