2 files changed, 115 insertions, 22 deletions

diff --git a/src/passes/RemoveUnusedBrs.cpp b/src/passes/RemoveUnusedBrs.cpp
index 9f753024d..c8c2277dc 100644
--- a/src/passes/RemoveUnusedBrs.cpp
+++ b/src/passes/RemoveUnusedBrs.cpp
@@ -276,6 +276,8 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
       }
     }
     // TODO: if-else can be turned into a br_if as well, if one of the sides is a dead end
+    //       we handle the case of a returned value to a set_local later down, see
+    //       visitSetLocal.
   }
 
   // override scan to add a pre and a post check task to all nodes
@@ -661,6 +663,56 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
         }
       }
 
+      void visitSetLocal(SetLocal* curr) {
+        // Undo an if return value into a local, if one arm is a br, as we
+        // prefer a br_if:
+        //  (set_local $x
+        //    (if (result i32)
+        //      (..condition..)
+        //      (br $somewhere)
+        //      (..result)
+        //    )
+        //  )
+        // =>
+        //  (br_if $somewhere
+        //    (..condition..)
+        //  )
+        //  (set_local $x
+        //    (..result)
+        //  )
+        // TODO: handle a condition in the br? need to watch for side effects
+        auto* iff = curr->value->dynCast<If>();
+        if (!iff) return;
+        if (!isConcreteType(iff->type)) return;
+        auto tryToOptimize = [&](Expression* one, Expression* two, bool flipCondition) {
+          if (one->type == unreachable && two->type != unreachable) {
+            if (auto* br = one->dynCast<Break>()) {
+              if (ExpressionAnalyzer::isSimple(br)) {
+                // Wonderful, do it!
+                Builder builder(*getModule());
+                br->condition = iff->condition;
+                if (flipCondition) {
+                  br->condition = builder.makeUnary(EqZInt32, br->condition);
+                }
+                br->finalize();
+                curr->value = two;
+                replaceCurrent(
+                  builder.makeSequence(
+                    br,
+                    curr
+                  )
+                );
+                return true;
+              }
+            }
+          }
+          return false;
+        };
+        if (!tryToOptimize(iff->ifTrue, iff->ifFalse, false)) {
+          tryToOptimize(iff->ifFalse, iff->ifTrue, true);
+        }
+      }
+
       // (br_if)+ => br_table
       // we look for the specific pattern of
       //  (br_if ..target1..
diff --git a/src/passes/SimplifyLocals.cpp b/src/passes/SimplifyLocals.cpp
index 4db85eb3e..6c1ea3fcb 100644
--- a/src/passes/SimplifyLocals.cpp
+++ b/src/passes/SimplifyLocals.cpp
@@ -456,16 +456,46 @@ struct SimplifyLocals : public WalkerPass<LinearExecutionWalker<SimplifyLocals<a
     // if this if already has a result, or is unreachable code, we have
     // nothing to do
     if (iff->type != none) return;
-    // We now have the sinkables from both sides of the if.
+    // We now have the sinkables from both sides of the if, and can look
+    // for something to sink. That is either a shared index on both sides,
+    // *or* if one side is unreachable, we can sink anything from the other,
+    //   (if
+    //     (..)
+    //     (br $x)
+    //     (set_local $y (..))
+    //   )
+    //    =>
+    //   (set_local $y
+    //     (if (result i32)
+    //       (..)
+    //       (br $x)
+    //       (..)
+    //     )
+    //   )
     Sinkables& ifFalse = sinkables;
-    Index sharedIndex = -1;
+    Index goodIndex = -1;
     bool found = false;
-    for (auto& sinkable : ifTrue) {
-      Index index = sinkable.first;
-      if (ifFalse.count(index) > 0) {
-        sharedIndex = index;
+    if (iff->ifTrue->type == unreachable) {
+      assert(iff->ifFalse->type != unreachable); // since the if type is none
+      if (!ifFalse.empty()) {
+        goodIndex = ifFalse.begin()->first;
         found = true;
-        break;
+      }
+    } else if (iff->ifFalse->type == unreachable) {
+      assert(iff->ifTrue->type != unreachable); // since the if type is none
+      if (!ifTrue.empty()) {
+        goodIndex = ifTrue.begin()->first;
+        found = true;
+      }
+    } else {
+      // Look for a shared index.
+      for (auto& sinkable : ifTrue) {
+        Index index = sinkable.first;
+        if (ifFalse.count(index) > 0) {
+          goodIndex = index;
+          found = true;
+          break;
+        }
       }
     }
     if (!found) return;
@@ -473,27 +503,38 @@ struct SimplifyLocals : public WalkerPass<LinearExecutionWalker<SimplifyLocals<a
     // ensure we have a place to write the return values for, if not, we
     // need another cycle
     auto* ifTrueBlock  = iff->ifTrue->dynCast<Block>();
+    if (iff->ifTrue->type != unreachable) {
+      if (!ifTrueBlock  || ifTrueBlock->list.size() == 0  || !ifTrueBlock->list.back()->is<Nop>()) {
+        ifsToEnlarge.push_back(iff);
+        return;
+      }
+    }
     auto* ifFalseBlock = iff->ifFalse->dynCast<Block>();
-    if (!ifTrueBlock  || ifTrueBlock->list.size() == 0  || !ifTrueBlock->list.back()->is<Nop>() ||
-        !ifFalseBlock || ifFalseBlock->list.size() == 0 || !ifFalseBlock->list.back()->is<Nop>()) {
-      ifsToEnlarge.push_back(iff);
-      return;
+    if (iff->ifFalse->type != unreachable) {
+      if (!ifFalseBlock  || ifFalseBlock->list.size() == 0  || !ifFalseBlock->list.back()->is<Nop>()) {
+        ifsToEnlarge.push_back(iff);
+        return;
+      }
     }
     // all set, go
-    auto *ifTrueItem = ifTrue.at(sharedIndex).item;
-    ifTrueBlock->list[ifTrueBlock->list.size() - 1] = (*ifTrueItem)->template cast<SetLocal>()->value;
-    ExpressionManipulator::nop(*ifTrueItem);
-    ifTrueBlock->finalize();
-    assert(ifTrueBlock->type != none);
-    auto *ifFalseItem = ifFalse.at(sharedIndex).item;
-    ifFalseBlock->list[ifFalseBlock->list.size() - 1] = (*ifFalseItem)->template cast<SetLocal>()->value;
-    ExpressionManipulator::nop(*ifFalseItem);
-    ifFalseBlock->finalize();
-    assert(ifTrueBlock->type != none);
+    if (iff->ifTrue->type != unreachable) {
+      auto *ifTrueItem = ifTrue.at(goodIndex).item;
+      ifTrueBlock->list[ifTrueBlock->list.size() - 1] = (*ifTrueItem)->template cast<SetLocal>()->value;
+      ExpressionManipulator::nop(*ifTrueItem);
+      ifTrueBlock->finalize();
+      assert(ifTrueBlock->type != none);
+    }
+    if (iff->ifFalse->type != unreachable) {
+      auto *ifFalseItem = ifFalse.at(goodIndex).item;
+      ifFalseBlock->list[ifFalseBlock->list.size() - 1] = (*ifFalseItem)->template cast<SetLocal>()->value;
+      ExpressionManipulator::nop(*ifFalseItem);
+      ifFalseBlock->finalize();
+      assert(ifFalseBlock->type != none);
+    }
     iff->finalize(); // update type
     assert(iff->type != none);
     // finally, create a set_local on the iff itself
-    auto* newSetLocal = Builder(*this->getModule()).makeSetLocal(sharedIndex, iff);
+    auto* newSetLocal = Builder(*this->getModule()).makeSetLocal(goodIndex, iff);
     *currp = newSetLocal;
     anotherCycle = true;
   }