2 files changed, 58 insertions, 4 deletions

diff --git a/src/passes/CoalesceLocals.cpp b/src/passes/CoalesceLocals.cpp
index 9a2f8d220..ce9b2d291 100644
--- a/src/passes/CoalesceLocals.cpp
+++ b/src/passes/CoalesceLocals.cpp
@@ -30,6 +30,7 @@
 #include "pass.h"
 #include "ast_utils.h"
 #include "cfg/cfg-traversal.h"
+#include "wasm-builder.h"
 #include "support/learning.h"
 #ifdef CFG_PROFILE
 #include "support/timing.h"
@@ -177,14 +178,28 @@ struct CoalesceLocals : public WalkerPass<CFGWalker<CoalesceLocals, Visitor<Coal
     }
     self->currBasicBlock->contents.actions.emplace_back(Action::Set, curr->index, currp);
     // if this is a copy, note it
-    auto* get = curr->value->dynCast<GetLocal>();
-    if (get) {
+    if (auto* get = self->getCopy(curr)) {
       // add 2 units, so that backedge prioritization can decide ties, but not much more
       self->addCopy(curr->index, get->index);
       self->addCopy(curr->index, get->index);
     }
   }
 
+  // A simple copy is a set of a get. A more interesting copy
+  // is a set of an if with a value, where one side a get.
+  // That can happen when we create an if value in simplify-locals. TODO: recurse into
+  // nested ifs, and block return values? Those cases are trickier, need to
+  // count to see if worth it.
+  // TODO: an if can have two copies
+  GetLocal* getCopy(SetLocal* set) {
+    if (auto* get = set->value->dynCast<GetLocal>()) return get;
+    if (auto* iff = set->value->dynCast<If>()) {
+      if (auto* get = iff->ifTrue->dynCast<GetLocal>()) return get;
+      if (auto* get = iff->ifFalse->dynCast<GetLocal>()) return get;
+    }
+    return nullptr;
+  }
+
   // main entry point
 
   void doWalkFunction(Function* func);
@@ -293,7 +308,7 @@ void CoalesceLocals::increaseBackEdgePriorities() {
       for (auto& action : arrivingBlock->contents.actions) {
         if (action.what == Action::Set) {
           auto* set = (*action.origin)->cast<SetLocal>();
-          if (auto* get = set->value->dynCast<GetLocal>()) {
+          if (auto* get = getCopy(set)) {
             // this is indeed a copy, add to the cost (default cost is 2, so this adds 50%, and can mostly break ties)
             addCopy(set->index, get->index);
           }
@@ -605,6 +620,22 @@ 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
+  *origin = iff;
+  set->value = other;
+  other = set;
+  if (!set->isTee()) {
+    // 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) {
@@ -624,7 +655,9 @@ void CoalesceLocals::applyIndices(std::vector<Index>& indices, Expression* root)
           } else {
             ExpressionManipulator::nop(set);
           }
-        } else if (!action.effective) {
+          continue;
+        }
+        if (!action.effective) {
           *action.origin = set->value; // value may have no side effects, further optimizations can eliminate it
           if (!set->isTee()) {
             // we need to drop it
@@ -632,6 +665,21 @@ void CoalesceLocals::applyIndices(std::vector<Index>& indices, Expression* root)
             drop->value = *action.origin;
             *action.origin = drop;
           }
+          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/wasm-builder.h b/src/wasm-builder.h
index 4f3a80091..0f2c333f9 100644
--- a/src/wasm-builder.h
+++ b/src/wasm-builder.h
@@ -245,6 +245,7 @@ public:
   }
 
   // Additional utility functions for building on top of nodes
+  // Convenient to have these on Builder, as it has allocation built in
 
   static Index addParam(Function* func, Name name, WasmType type) {
     // only ok to add a param if no vars, otherwise indices are invalidated
@@ -368,6 +369,11 @@ public:
     if (!isConcreteWasmType(curr->type)) return curr;
     return makeDrop(curr);
   }
+
+  void flip(If* iff) {
+    std::swap(iff->ifTrue, iff->ifFalse);
+    iff->condition = makeUnary(EqZInt32, iff->condition);
+  }
 };
 
 } // namespace wasm