Move if copy logic from coalesce-locals to remove-unused-brs.

If copies is the case where an if arm is a get that feeds into
a set of the same local:

  (set_local $x
    (if (result i32)
      (..condition..)
      (..result)
      (get_local $x)
    )
  )

We can rework this so that the if-else is only an if, which
executes the code path not going to the get.

This was done in coalesce-locals only because it is likely to
work there as after coalescing there are more copies. However,
the logic is of removing a branch, and so belongs in
remove-unused-brs, and fits alongside existing logic there
for handling ifs with an arm that is a br. Also refactor that
code so that the two optimizations can feed into each other.

2 files changed, 128 insertions, 64 deletions

diff --git a/src/passes/CoalesceLocals.cpp b/src/passes/CoalesceLocals.cpp
index 63912db19..1bcc74f79 100644
--- a/src/passes/CoalesceLocals.cpp
+++ b/src/passes/CoalesceLocals.cpp
@@ -315,25 +315,6 @@ void CoalesceLocals::pickIndices(std::vector<Index>& indices) {
   }
 }
 
-// Remove a copy from a set of an if, where one if arm is a get of the same set
-static void removeIfCopy(Expression** origin, SetLocal* set, If* iff, Expression*& copy, Expression*& other, Module* module) {
-  // replace the origin with the if, and sink the set into the other non-copying arm
-  bool tee = set->isTee();
-  *origin = iff;
-  set->value = other;
-  set->finalize();
-  other = set;
-  // if this is not a tee, then we can get rid of the copy in that arm
-  if (!tee) {
-    // we don't need the copy at all
-    copy = nullptr;
-    if (!iff->ifTrue) {
-      Builder(*module).flip(iff);
-    }
-    iff->finalize();
-  }
-}
-
 void CoalesceLocals::applyIndices(std::vector<Index>& indices, Expression* root) {
   assert(indices.size() == numLocals);
   for (auto& curr : basicBlocks) {
@@ -362,20 +343,6 @@ void CoalesceLocals::applyIndices(std::vector<Index>& indices, Expression* root)
           }
           continue;
         }
-        if (auto* iff = set->value->dynCast<If>()) {
-          if (auto* get = iff->ifTrue->dynCast<GetLocal>()) {
-            if (get->index == set->index) {
-              removeIfCopy(action.origin, set, iff, iff->ifTrue, iff->ifFalse, getModule());
-              continue;
-            }
-          }
-          if (auto* get = iff->ifFalse->dynCast<GetLocal>()) {
-            if (get->index == set->index) {
-              removeIfCopy(action.origin, set, iff, iff->ifFalse, iff->ifTrue, getModule());
-              continue;
-            }
-          }
-        }
       }
     }
   }
diff --git a/src/passes/RemoveUnusedBrs.cpp b/src/passes/RemoveUnusedBrs.cpp
index 62ecb3ed3..7fa94baf0 100644
--- a/src/passes/RemoveUnusedBrs.cpp
+++ b/src/passes/RemoveUnusedBrs.cpp
@@ -804,53 +804,150 @@ 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) || !isConcreteType(iff->condition->type)) return;
+        // Sets of an if can be optimized in various ways that remove part of
+        // the if branching, or all of it.
+        // The optimizations we can do here can recurse and call each
+        // other, so pass around a pointer to the output.
+        optimizeSetIf(getCurrentPointer());
+      }
+
+      void optimizeSetIf(Expression** currp) {
+        if (optimizeSetIfWithBrArm(currp)) return;
+        if (optimizeSetIfWithCopyArm(currp)) return;
+      }
+
+      // If one arm is a br, we prefer a br_if and the set later:
+      //  (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
+      bool optimizeSetIfWithBrArm(Expression** currp) {
+        auto* set = (*currp)->cast<SetLocal>();
+        auto* iff = set->value->dynCast<If>();
+        if (!iff ||
+            !isConcreteType(iff->type) ||
+            !isConcreteType(iff->condition->type)) {
+          return false;
+        }
         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);
+                  builder.flip(iff);
                 }
+                br->condition = iff->condition;
                 br->finalize();
-                curr->value = two;
-                replaceCurrent(
-                  builder.makeSequence(
-                    br,
-                    curr
-                  )
-                );
+                set->value = two;
+                auto* block = builder.makeSequence(br, set);
+                *currp = block;
+                // Recurse on the set, which now has a new value.
+                optimizeSetIf(&block->list[1]);
                 return true;
               }
             }
           }
           return false;
         };
-        if (!tryToOptimize(iff->ifTrue, iff->ifFalse, false)) {
-          tryToOptimize(iff->ifFalse, iff->ifTrue, true);
+        return tryToOptimize(iff->ifTrue, iff->ifFalse, false) ||
+               tryToOptimize(iff->ifFalse, iff->ifTrue, true);
+      }
+
+      // If one arm is a get of the same outer set, it is a copy which
+      // we can remove. If this is not a tee, then we remove the get
+      // as well as the if-else opcode in the binary format, which is
+      // great:
+      //  (set_local $x
+      //    (if (result i32)
+      //      (..condition..)
+      //      (..result)
+      //      (get_local $x)
+      //    )
+      //  )
+      // =>
+      //  (if
+      //    (..condition..)
+      //    (set_local $x
+      //      (..result)
+      //    )
+      //  )
+      // If this is a tee, then we can do the same operation but
+      // inside a block, and keep the get:
+      //  (tee_local $x
+      //    (if (result i32)
+      //      (..condition..)
+      //      (..result)
+      //      (get_local $x)
+      //    )
+      //  )
+      // =>
+      //  (block (result i32)
+      //    (if
+      //      (..condition..)
+      //      (set_local $x
+      //        (..result)
+      //      )
+      //    )
+      //    (get_local $x)
+      //  )
+      // We save the if-else opcode, and add the block's opcodes.
+      // This may be detrimental, however, often the block can be
+      // merged or eliminated given the outside scope, and we
+      // removed one of the if branches.
+      bool optimizeSetIfWithCopyArm(Expression** currp) {
+        auto* set = (*currp)->cast<SetLocal>();
+        auto* iff = set->value->dynCast<If>();
+        if (!iff ||
+            !isConcreteType(iff->type) ||
+            !isConcreteType(iff->condition->type)) {
+          return false;
+        }
+        Builder builder(*getModule());
+        GetLocal* get = iff->ifTrue->dynCast<GetLocal>();
+        if (get && get->index == set->index) {
+          builder.flip(iff);
+        } else {
+          get = iff->ifFalse->dynCast<GetLocal>();
+          if (get && get->index != set->index) {
+            get = nullptr;
+          }
+        }
+        if (!get) return false;
+        // We can do it!
+        bool tee = set->isTee();
+        assert(set->index == get->index);
+        assert(iff->ifFalse == get);
+        set->value = iff->ifTrue;
+        set->finalize();
+        iff->ifTrue = set;
+        iff->ifFalse = nullptr;
+        iff->finalize();
+        Expression* replacement = iff;
+        if (tee) {
+          set->setTee(false);
+          // We need a block too.
+          replacement = builder.makeSequence(
+            iff,
+            get // reuse the get
+          );
         }
+        *currp = replacement;
+        // Recurse on the set, which now has a new value.
+        optimizeSetIf(&iff->ifTrue);
+        return true;
       }
 
       // (br_if)+ => br_table