Merge pull request #386 from WebAssembly/improve-remove-brs

Improve RemoveBrs

3 files changed, 139 insertions, 62 deletions

diff --git a/src/pass.cpp b/src/pass.cpp
index 7095da71d..5c4f510f8 100644
--- a/src/pass.cpp
+++ b/src/pass.cpp
@@ -65,6 +65,7 @@ void PassRunner::addDefaultOptimizationPasses() {
   add("merge-blocks");
   add("reorder-locals");
   add("vacuum");
+  add("optimize-instructions");
 }
 
 void PassRunner::run(Module* module) {
diff --git a/src/passes/RemoveUnusedBrs.cpp b/src/passes/RemoveUnusedBrs.cpp
index cfd9ea7ce..3d6d6973f 100644
--- a/src/passes/RemoveUnusedBrs.cpp
+++ b/src/passes/RemoveUnusedBrs.cpp
@@ -15,89 +15,133 @@
  */
 
 //
-// Removes branches that go to where they go anyhow
+// Removes branches for which we go to where they go anyhow
 //
 
 #include <wasm.h>
 #include <pass.h>
+#include <ast_utils.h>
 
 namespace wasm {
 
 struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs, Visitor<RemoveUnusedBrs>>> {
   bool isFunctionParallel() { return true; }
 
-  // preparation: try to unify branches, as the fewer there are, the higher a chance we can remove them
-  // specifically for if-else, turn an if-else with branches to the same target at the end of each
-  // child, and with a value, to a branch to that target containing the if-else
+  bool anotherCycle;
+
+  typedef std::vector<Break*> Flows;
+
+  // list of breaks that are currently flowing. if they reach their target without
+  // interference, they can be removed (or their value forwarded TODO)
+  Flows flows;
+
+  // a stack for if-else contents, we merge their outputs
+  std::vector<Flows> ifStack;
+
+  static void visitAny(RemoveUnusedBrs* self, Expression** currp) {
+    auto* curr = *currp;
+    auto& flows = self->flows;
+
+    if (curr->is<Break>()) {
+      flows.clear();
+      auto* br = curr->cast<Break>();
+      if (!br->condition && !br->value) { // TODO: optimize in those cases?
+        // a break, let's see where it flows to
+        flows.push_back(br);
+      }
+    } else if (curr->is<Switch>()) {
+      flows.clear();
+    } else if (curr->is<If>()) {
+      auto* iff = curr->cast<If>();
+      if (iff->ifFalse) {
+        assert(self->ifStack.size() > 0);
+        for (auto* flow : self->ifStack.back()) {
+          flows.push_back(flow);
+        }
+        self->ifStack.pop_back();
+      }
+    } else if (curr->is<Block>()) {
+      // any breaks flowing to here are unnecessary, as we get here anyhow
+      auto name = curr->cast<Block>()->name;
+      if (name.is()) {
+        size_t size = flows.size();
+        size_t skip = 0;
+        for (size_t i = 0; i < size; i++) {
+          if (flows[i]->name == name) {
+            ExpressionManipulator::nop(flows[i]);
+            skip++;
+            self->anotherCycle = true;
+          } else if (skip > 0) {
+            flows[i - skip] = flows[i];
+          }
+        }
+        if (skip > 0) {
+          flows.resize(size - skip);
+        }
+      }
+    } else if (curr->is<Loop>()) {
+      // TODO we might optimize branches out of here
+      flows.clear();
+    } else if (curr->is<Nop>()) {
+      // ignore (could be result of a previous cycle)
+    } else {
+      // anything else stops the flow
+      flows.clear();
+    }
+  }
+
+  static void clear(RemoveUnusedBrs* self, Expression** currp) {
+    self->flows.clear();
+  }
+
+  static void saveIfTrue(RemoveUnusedBrs* self, Expression** currp) {
+    self->ifStack.push_back(std::move(self->flows));
+  }
+
   void visitIf(If* curr) {
     if (!curr->ifFalse) {
-      // try to reduce an   if (condition) br  =>  br_if (condition) , which might open up other optimization opportunities
+      // if without an else. try to reduce   if (condition) br  =>  br_if (condition)
       Break* br = curr->ifTrue->dynCast<Break>();
       if (br && !br->condition) { // TODO: if there is a condition, join them
         br->condition = curr->condition;
         replaceCurrent(br);
+        anotherCycle = true;
       }
-      return;
-    }
-    if (isConcreteWasmType(curr->type)) return; // already has a returned value
-    // an if_else that indirectly returns a value by breaking to the same target can potentially remove both breaks, and break outside once
-    auto getLast = [](Expression *side) -> Expression* {
-      Block* b = side->dynCast<Block>();
-      if (!b) return nullptr;
-      if (b->list.size() == 0) return nullptr;
-      return b->list.back();
-    };
-    auto process = [&](Expression *side, bool doIt) {
-      Expression* last = getLast(side);
-      if (!last) return Name();
-      Block* b = side->cast<Block>();
-      Break* br = last->dynCast<Break>();
-      if (!br) return Name();
-      if (br->condition) return Name();
-      if (!br->value) return Name();
-      if (doIt) {
-        b->list[b->list.size()-1] = br->value;
-      }
-      return br->name;
-    };
-    // do both, or none
-    if (process(curr->ifTrue, false).is() && process(curr->ifTrue, false) == process(curr->ifFalse, false)) {
-      auto br = getLast(curr->ifTrue)->cast<Break>(); // we are about to discard this, so why not reuse it!
-      process(curr->ifTrue, true);
-      process(curr->ifFalse, true);
-      curr->type = br->value->type; // if_else now returns a value
-      br->value = curr;
-      // no need to change anything else in the br - target is correct already
-      replaceCurrent(br);
     }
   }
 
-  // main portion
-  void visitBlock(Block *curr) {
-    if (curr->name.isNull()) return;
-    if (curr->list.size() == 0) return;
-    // preparation - remove all code after an unconditional break, since it can't execute, and it might confuse us (we look at the last)
-    for (size_t i = 0; i < curr->list.size()-1; i++) {
-      Break* br = curr->list[i]->dynCast<Break>();
-      if (br && !br->condition) {
-        curr->list.resize(i+1);
-        break;
-      }
-    }
-    Expression* last = curr->list.back();
-    if (Break* br = last->dynCast<Break>()) {
-      if (br->condition) return;
-      if (br->name == curr->name) {
-        if (!br->value) {
-          curr->list.pop_back();
-        } else {
-          curr->list[curr->list.size()-1] = br->value; // can replace with the value
-        }
+  // override scan to add a pre and a post check task to all nodes
+  static void scan(RemoveUnusedBrs* self, Expression** currp) {
+    self->pushTask(visitAny, currp);
+
+    auto* iff = (*currp)->dynCast<If>();
+
+    if (iff) {
+      self->pushTask(doVisitIf, currp);
+      if (iff->ifFalse) {
+      // we need to join up if-else control flow, and clear after the condition
+        self->pushTask(scan, &iff->ifFalse);
+        self->pushTask(saveIfTrue, currp); // safe the ifTrue flow, we'll join it later
       }
+      self->pushTask(scan, &iff->ifTrue);
+      self->pushTask(clear, currp); // clear all flow after the condition
+      self->pushTask(scan, &iff->condition);
+    } else {
+      WalkerPass<PostWalker<RemoveUnusedBrs, Visitor<RemoveUnusedBrs>>>::scan(self, currp);
     }
   }
+
+  void walk(Expression*& root) {
+    // multiple cycles may be needed
+    do {
+      anotherCycle = false;
+      WalkerPass<PostWalker<RemoveUnusedBrs, Visitor<RemoveUnusedBrs>>>::walk(root);
+      assert(ifStack.empty());
+      assert(flows.empty());
+    } while (anotherCycle);
+  }
 };
 
-static RegisterPass<RemoveUnusedBrs> registerPass("remove-unused-brs", "removes breaks from locations that are never branched to");
+static RegisterPass<RemoveUnusedBrs> registerPass("remove-unused-brs", "removes breaks from locations that are not needed");
 
 } // namespace wasm
diff --git a/src/passes/Vacuum.cpp b/src/passes/Vacuum.cpp
index 863d78f14..dd88fa0eb 100644
--- a/src/passes/Vacuum.cpp
+++ b/src/passes/Vacuum.cpp
@@ -21,6 +21,7 @@
 #include <wasm.h>
 #include <pass.h>
 #include <ast_utils.h>
+#include <wasm-builder.h>
 
 namespace wasm {
 
@@ -32,14 +33,26 @@ struct Vacuum : public WalkerPass<PostWalker<Vacuum, Visitor<Vacuum>>> {
     int skip = 0;
     auto& list = curr->list;
     size_t size = list.size();
+    bool needResize = false;
     for (size_t z = 0; z < size; z++) {
       if (list[z]->is<Nop>()) {
         skip++;
-      } else if (skip > 0) {
-        list[z - skip] = list[z];
+        needResize = true;
+      } else {
+        if (skip > 0) {
+          list[z - skip] = list[z];
+        }
+        // if this is an unconditional br, the rest is dead code
+        Break* br = list[z - skip]->dynCast<Break>();
+        Switch* sw = list[z - skip]->dynCast<Switch>();
+        if ((br && !br->condition) || sw) {
+          list.resize(z - skip + 1);
+          needResize = false;
+          break;
+        }
       }
     }
-    if (skip > 0) {
+    if (needResize) {
       list.resize(size - skip);
     }
     if (!curr->name.is()) {
@@ -50,6 +63,25 @@ struct Vacuum : public WalkerPass<PostWalker<Vacuum, Visitor<Vacuum>>> {
       }
     }
   }
+
+  void visitIf(If* curr) {
+    if (curr->ifFalse) {
+      if (curr->ifFalse->is<Nop>()) {
+        curr->ifFalse = nullptr;
+      } else if (curr->ifTrue->is<Nop>()) {
+        curr->ifTrue = curr->ifFalse;
+        curr->ifFalse = nullptr;
+        curr->condition = Builder(*getModule()).makeUnary(EqZ, curr->condition, curr->condition->type);
+      }
+    }
+    if (!curr->ifFalse) {
+      // no else
+      if (curr->ifTrue->is<Nop>()) {
+        // no nothing
+        replaceCurrent(curr->condition);
+      }
+    }
+  }
 };
 
 static RegisterPass<Vacuum> registerPass("vacuum", "removes obviously unneeded code");