Switch optimizations in remove-unused-brs (#1753)

* Switch optimizations in remove-unused-brs: thread switch jumps, and turn a switch with all identical targets into a br

* refinalize in interm operations in remove-unused-brs, as we can be confused by it

4 files changed, 95 insertions, 47 deletions

diff --git a/src/ir/branch-utils.h b/src/ir/branch-utils.h
index 84be9f897..6f9299bf5 100644
--- a/src/ir/branch-utils.h
+++ b/src/ir/branch-utils.h
@@ -45,6 +45,40 @@ inline bool isBranchReachable(Expression* expr) {
   WASM_UNREACHABLE();
 }
 
+inline std::set<Name> getUniqueTargets(Switch* sw) {
+  std::set<Name> ret;
+  for (auto target : sw->targets) {
+    ret.insert(target);
+  }
+  ret.insert(sw->default_);
+  return ret;
+}
+
+// If we branch to 'from', change that to 'to' instead.
+inline bool replacePossibleTarget(Expression* branch, Name from, Name to) {
+  bool worked = false;
+  if (auto* br = branch->dynCast<Break>()) {
+    if (br->name == from) {
+      br->name = to;
+      worked = true;
+    }
+  } else if (auto* sw = branch->dynCast<Switch>()) {
+    for (auto& target : sw->targets) {
+      if (target == from) {
+        target = to;
+        worked = true;
+      }
+    }
+    if (sw->default_ == from) {
+      sw->default_ = to;
+      worked = true;
+    }
+  } else {
+    WASM_UNREACHABLE();
+  }
+  return worked;
+}
+
 // returns the set of targets to which we branch that are
 // outside of a node
 inline std::set<Name> getExitingBranches(Expression* ast) {
diff --git a/src/passes/Flatten.cpp b/src/passes/Flatten.cpp
index f4b468098..aa5c1a491 100644
--- a/src/passes/Flatten.cpp
+++ b/src/passes/Flatten.cpp
@@ -53,8 +53,9 @@
 #include <wasm.h>
 #include <pass.h>
 #include <wasm-builder.h>
-#include <ir/utils.h>
+#include <ir/branch-utils.h>
 #include <ir/effects.h>
+#include <ir/utils.h>
 
 namespace wasm {
 
@@ -232,11 +233,7 @@ struct Flatten : public WalkerPass<ExpressionStackWalker<Flatten, UnifiedExpress
             Index temp = builder.addVar(getFunction(), type);
             ourPreludes.push_back(builder.makeSetLocal(temp, sw->value));
             // we don't know which break target will be hit - assign to them all
-            std::set<Name> names;
-            for (auto target : sw->targets) {
-              names.insert(target);
-            }
-            names.insert(sw->default_);
+            auto names = BranchUtils::getUniqueTargets(sw);
             for (auto name : names) {
               ourPreludes.push_back(builder.makeSetLocal(
                 getTempForBreakTarget(name, type),
diff --git a/src/passes/RemoveUnusedBrs.cpp b/src/passes/RemoveUnusedBrs.cpp
index eb81b0a5a..62ecb3ed3 100644
--- a/src/passes/RemoveUnusedBrs.cpp
+++ b/src/passes/RemoveUnusedBrs.cpp
@@ -435,7 +435,7 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
     }
   }
 
-  void sinkBlocks(Function* func) {
+  bool sinkBlocks(Function* func) {
     struct Sinker : public PostWalker<Sinker> {
       bool worked = false;
 
@@ -501,13 +501,14 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
 
     sinker.doWalkFunction(func);
     if (sinker.worked) {
-      anotherCycle = true;
+      ReFinalize().walkFunctionInModule(func, getModule());
+      return true;
     }
+    return false;
   }
 
   void doWalkFunction(Function* func) {
     // multiple cycles may be needed
-    bool worked = false;
     do {
       anotherCycle = false;
       super::doWalkFunction(func);
@@ -532,32 +533,39 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
         anotherCycle |= optimizeLoop(loop);
       }
       loops.clear();
+      if (anotherCycle) {
+        ReFinalize().walkFunctionInModule(func, getModule());
+      }
       // sink blocks
-      sinkBlocks(func);
-      if (anotherCycle) worked = true;
+      if (sinkBlocks(func)) {
+        anotherCycle = true;
+      }
     } while (anotherCycle);
 
-    if (worked) {
-      // Our work may alter block and if types, they may now return values that we made flow through them
-      ReFinalize().walkFunctionInModule(func, getModule());
-    }
-
     // thread trivial jumps
     struct JumpThreader : public ControlFlowWalker<JumpThreader> {
-      // map of all value-less breaks going to a block (and not a loop)
-      std::map<Block*, std::vector<Break*>> breaksToBlock;
+      // map of all value-less breaks and switches going to a block (and not a loop)
+      std::map<Block*, std::vector<Expression*>> branchesToBlock;
 
-      // the names to update
-      std::map<Break*, Name> newNames;
+      bool worked = false;
 
       void visitBreak(Break* curr) {
         if (!curr->value) {
           if (auto* target = findBreakTarget(curr->name)->dynCast<Block>()) {
-            breaksToBlock[target].push_back(curr);
+            branchesToBlock[target].push_back(curr);
+          }
+        }
+      }
+      void visitSwitch(Switch* curr) {
+        if (!curr->value) {
+          auto names = BranchUtils::getUniqueTargets(curr);
+          for (auto name : names) {
+            if (auto* target = findBreakTarget(name)->dynCast<Block>()) {
+              branchesToBlock[target].push_back(curr);
+            }
           }
         }
       }
-      // TODO: Switch?
       void visitBlock(Block* curr) {
         auto& list = curr->list;
         if (list.size() == 1 && curr->name.is()) {
@@ -566,12 +574,7 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
             // the two blocks must have the same type for us to update the branch, as otherwise
             // one block may be unreachable and the other concrete, so one might lack a value
             if (child->name.is() && child->name != curr->name && child->type == curr->type) {
-              auto& breaks = breaksToBlock[child];
-              for (auto* br : breaks) {
-                newNames[br] = curr->name;
-                breaksToBlock[curr].push_back(br); // update the list - we may push it even more later
-              }
-              breaksToBlock.erase(child);
+              redirectBranches(child, curr->name);
             }
           }
         } else if (list.size() == 2) {
@@ -579,28 +582,28 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
           auto* child = list[0]->dynCast<Block>();
           auto* jump = list[1]->dynCast<Break>();
           if (child && child->name.is() && jump && ExpressionAnalyzer::isSimple(jump)) {
-            auto& breaks = breaksToBlock[child];
-            for (auto* br : breaks) {
-              newNames[br] = jump->name;
-            }
-            // if the jump is to another block then we can update the list, and maybe push it even more later
-            if (auto* newTarget = findBreakTarget(jump->name)->dynCast<Block>()) {
-              for (auto* br : breaks) {
-                breaksToBlock[newTarget].push_back(br);
-              }
-            }
-            breaksToBlock.erase(child);
+            redirectBranches(child, jump->name);
           }
         }
       }
 
-      void finish(Function* func) {
-        for (auto& iter : newNames) {
-          auto* br = iter.first;
-          auto name = iter.second;
-          br->name = name;
+      void redirectBranches(Block* from, Name to) {
+        auto& branches = branchesToBlock[from];
+        for (auto* branch : branches) {
+          if (BranchUtils::replacePossibleTarget(branch, from->name, to)) {
+            worked = true;
+          }
+        }
+        // if the jump is to another block then we can update the list, and maybe push it even more later
+        if (auto* newTarget = findBreakTarget(to)->dynCast<Block>()) {
+          for (auto* branch : branches) {
+            branchesToBlock[newTarget].push_back(branch);
+          }
         }
-        if (newNames.size() > 0) {
+      }
+
+      void finish(Function* func) {
+        if (worked) {
           // by changing where brs go, we may change block types etc.
           ReFinalize().walkFunctionInModule(func, getModule());
         }
@@ -686,6 +689,19 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
         }
       }
 
+      void visitSwitch(Switch* curr) {
+        if (BranchUtils::getUniqueTargets(curr).size() == 1) {
+          // This switch has just one target no matter what; replace with a br.
+          Builder builder(*getModule());
+          replaceCurrent(
+            builder.makeSequence(
+              builder.makeDrop(curr->condition), // might have side effects
+              builder.makeBreak(curr->default_, curr->value)
+            )
+          );
+        }
+      }
+
       // Restructuring of ifs: if we have
       //   (block $x
       //     (br_if $x (cond))
diff --git a/src/passes/SimplifyLocals.cpp b/src/passes/SimplifyLocals.cpp
index aadf766ac..bf4eed24e 100644
--- a/src/passes/SimplifyLocals.cpp
+++ b/src/passes/SimplifyLocals.cpp
@@ -50,6 +50,7 @@
 #include <wasm-builder.h>
 #include <wasm-traversal.h>
 #include <pass.h>
+#include <ir/branch-utils.h>
 #include <ir/count.h>
 #include <ir/effects.h>
 #include "ir/equivalent_sets.h"
@@ -128,10 +129,10 @@ struct SimplifyLocals : public WalkerPass<LinearExecutionWalker<SimplifyLocals<a
       assert(!curr->cast<If>()->ifFalse); // if-elses are handled by doNoteIfElse* methods
     } else if (curr->is<Switch>()) {
       auto* sw = curr->cast<Switch>();
-      for (auto target : sw->targets) {
+      auto targets = BranchUtils::getUniqueTargets(sw);
+      for (auto target : targets) {
         self->unoptimizableBlocks.insert(target);
       }
-      self->unoptimizableBlocks.insert(sw->default_);
       // TODO: we could use this info to stop gathering data on these blocks
     }
     self->sinkables.clear();