1 files changed, 44 insertions, 16 deletions

diff --git a/src/passes/SimplifyLocals.cpp b/src/passes/SimplifyLocals.cpp
index dfbc7551a..c368740f1 100644
--- a/src/passes/SimplifyLocals.cpp
+++ b/src/passes/SimplifyLocals.cpp
@@ -1000,8 +1000,10 @@ struct SimplifyLocals
       std::vector<Index>* numLocalGets;
       bool removeEquivalentSets;
       Module* module;
+      PassOptions passOptions;
 
       bool anotherCycle = false;
+      bool refinalize = false;
 
       // We track locals containing the same value.
       EquivalentSets equivalences;
@@ -1015,12 +1017,9 @@ struct SimplifyLocals
 
       void visitLocalSet(LocalSet* curr) {
         // Remove trivial copies, even through a tee
-        auto* value = curr->value;
-        Function* func = this->getFunction();
-        while (auto* subSet = value->dynCast<LocalSet>()) {
-          value = subSet->value;
-        }
-        if (auto* get = value->dynCast<LocalGet>()) {
+        auto* value =
+          Properties::getFallthrough(curr->value, passOptions, *module);
+        if (auto* get = value->template dynCast<LocalGet>()) {
           if (equivalences.check(curr->index, get->index)) {
             // This is an unnecessary copy!
             if (removeEquivalentSets) {
@@ -1033,13 +1032,9 @@ struct SimplifyLocals
             }
             // Nothing more to do, ignore the copy.
             return;
-          } else if (func->getLocalType(curr->index) ==
-                     func->getLocalType(get->index)) {
+          } else {
             // There is a new equivalence now. Remove all the old ones, and add
             // the new one.
-            // Note that we ignore the case of subtyping here, to keep this
-            // optimization simple by assuming all equivalent indexes also have
-            // the same type. TODO: consider optimizing this.
             equivalences.reset(curr->index);
             equivalences.add(curr->index, get->index);
             return;
@@ -1054,8 +1049,6 @@ struct SimplifyLocals
         // Canonicalize gets: if some are equivalent, then we can pick more
         // then one, and other passes may benefit from having more uniformity.
         if (auto* set = equivalences.getEquivalents(curr->index)) {
-          // Pick the index with the most uses - maximizing the chance to
-          // lower one's uses to zero.
           // Helper method that returns the # of gets *ignoring the current
           // get*, as we want to see what is best overall, treating this one as
           // to be decided upon.
@@ -1067,9 +1060,31 @@ struct SimplifyLocals
             }
             return ret;
           };
+
+          // Pick the index with the most uses - maximizing the chance to
+          // lower one's uses to zero. If types differ though then we prefer to
+          // switch to a more refined type even if there are fewer uses, as that
+          // may have significant benefits to later optimizations (we may be
+          // able to use it to remove casts, etc.).
+          auto* func = this->getFunction();
           Index best = -1;
           for (auto index : *set) {
-            if (best == Index(-1) ||
+            if (best == Index(-1)) {
+              // This is the first possible option we've seen.
+              best = index;
+              continue;
+            }
+
+            auto bestType = func->getLocalType(best);
+            auto indexType = func->getLocalType(index);
+            if (!Type::isSubType(indexType, bestType)) {
+              // This is less refined than the current best; ignore.
+              continue;
+            }
+
+            // This is better if it has a more refined type, or if it has more
+            // uses.
+            if (indexType != bestType ||
                 getNumGetsIgnoringCurr(index) > getNumGetsIgnoringCurr(best)) {
               best = index;
             }
@@ -1077,8 +1092,11 @@ struct SimplifyLocals
           assert(best != Index(-1));
           // Due to ordering, the best index may be different from us but have
           // the same # of locals - make sure we actually improve.
-          if (best != curr->index && getNumGetsIgnoringCurr(best) >
-                                       getNumGetsIgnoringCurr(curr->index)) {
+          auto bestType = func->getLocalType(best);
+          auto oldType = func->getLocalType(curr->index);
+          if (best != curr->index && (getNumGetsIgnoringCurr(best) >
+                                        getNumGetsIgnoringCurr(curr->index) ||
+                                      bestType != oldType)) {
             // Update the get counts.
             (*numLocalGets)[best]++;
             assert((*numLocalGets)[curr->index] >= 1);
@@ -1086,6 +1104,12 @@ struct SimplifyLocals
             // Make the change.
             curr->index = best;
             anotherCycle = true;
+            if (bestType != oldType) {
+              curr->type = func->getLocalType(best);
+              // We are switching to a more refined type, which might require
+              // changes in the user of the local.get.
+              refinalize = true;
+            }
           }
         }
       }
@@ -1093,9 +1117,13 @@ struct SimplifyLocals
 
     EquivalentOptimizer eqOpter;
     eqOpter.module = this->getModule();
+    eqOpter.passOptions = this->getPassOptions();
     eqOpter.numLocalGets = &getCounter.num;
     eqOpter.removeEquivalentSets = allowStructure;
     eqOpter.walkFunction(func);
+    if (eqOpter.refinalize) {
+      ReFinalize().walkFunctionInModule(func, this->getModule());
+    }
 
     // We may have already had a local with no uses, or we may have just
     // gotten there thanks to the EquivalentOptimizer. If there are such