3 files changed, 67 insertions, 18 deletions

diff --git a/src/ir/lubs.h b/src/ir/lubs.h
new file mode 100644
index 000000000..b75932091
--- /dev/null
+++ b/src/ir/lubs.h
@@ -0,0 +1,46 @@
+/*
+ * Copyright 2021 WebAssembly Community Group participants
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef wasm_ir_lubs_h
+#define wasm_ir_lubs_h
+
+#include "ir/module-utils.h"
+#include "wasm.h"
+
+namespace wasm {
+
+// Helper to find a LUB of a series of expressions. This works incrementally so
+// that if we see we are not improving on an existing type then we can stop
+// early.
+struct LUBFinder {
+  // The least upper bound we found so far.
+  Type lub = Type::unreachable;
+
+  // Note another type to take into account in the lub. Returns the new lub.
+  Type note(Type type) { return lub = Type::getLeastUpperBound(lub, type); }
+
+  Type note(Expression* curr) { return note(curr->type); }
+
+  // Returns whether we noted any (reachable) value.
+  bool noted() { return lub != Type::unreachable; }
+
+  // Returns the lub that we found.
+  Type get() { return lub; }
+};
+
+} // namespace wasm
+
+#endif // wasm_ir_lubs_h
diff --git a/src/passes/DeadArgumentElimination.cpp b/src/passes/DeadArgumentElimination.cpp
index afe9f19b7..033e766b6 100644
--- a/src/passes/DeadArgumentElimination.cpp
+++ b/src/passes/DeadArgumentElimination.cpp
@@ -41,6 +41,7 @@
 #include "ir/effects.h"
 #include "ir/element-utils.h"
 #include "ir/find_all.h"
+#include "ir/lubs.h"
 #include "ir/module-utils.h"
 #include "ir/type-updating.h"
 #include "ir/utils.h"
@@ -556,21 +557,20 @@ private:
         newParamTypes.push_back(originalType);
         continue;
       }
-      Type refinedType = Type::unreachable;
+      LUBFinder lub;
       for (auto* call : calls) {
         auto* operand = call->operands[i];
-        refinedType = Type::getLeastUpperBound(refinedType, operand->type);
-        if (refinedType == originalType) {
+        if (lub.note(operand) == originalType) {
           // We failed to refine this parameter to anything more specific.
           break;
         }
       }
 
       // Nothing is sent here at all; leave such optimizations to DCE.
-      if (refinedType == Type::unreachable) {
+      if (!lub.noted()) {
         return;
       }
-      newParamTypes.push_back(refinedType);
+      newParamTypes.push_back(lub.get());
     }
 
     // Check if we are able to optimize here before we do the work to scan the
@@ -654,17 +654,16 @@ private:
     //  )
     ReFinalize().walkFunctionInModule(func, module);
 
-    Type refinedType = func->body->type;
-    if (refinedType == originalType) {
+    LUBFinder lub;
+    if (lub.note(func->body) == originalType) {
       return false;
     }
 
     // Scan the body and look at the returns.
     auto processReturnType = [&](Type type) {
-      refinedType = Type::getLeastUpperBound(refinedType, type);
       // Return whether we still look ok to do the optimization. If this is
       // false then we can stop here.
-      return refinedType != originalType;
+      return lub.note(type) != originalType;
     };
     for (auto* ret : FindAll<Return>(func->body).list) {
       if (!processReturnType(ret->value->type)) {
@@ -694,20 +693,20 @@ private:
         }
       }
     }
-    assert(refinedType != originalType);
+    assert(lub.get() != originalType);
 
     // If the refined type is unreachable then nothing actually returns from
     // this function.
     // TODO: We can propagate that to the outside, and not just for GC.
-    if (refinedType == Type::unreachable) {
+    if (!lub.noted()) {
       return false;
     }
 
     // Success. Update the type, and the calls.
-    func->setResults(refinedType);
+    func->setResults(lub.get());
     for (auto* call : calls) {
       if (call->type != Type::unreachable) {
-        call->type = refinedType;
+        call->type = lub.get();
       }
     }
     return true;
diff --git a/src/passes/LocalSubtyping.cpp b/src/passes/LocalSubtyping.cpp
index 24a2e48e4..57119a253 100644
--- a/src/passes/LocalSubtyping.cpp
+++ b/src/passes/LocalSubtyping.cpp
@@ -25,6 +25,7 @@
 #include <ir/find_all.h>
 #include <ir/linear-execution.h>
 #include <ir/local-graph.h>
+#include <ir/lubs.h>
 #include <ir/utils.h>
 #include <pass.h>
 #include <wasm-builder.h>
@@ -118,18 +119,21 @@ struct LocalSubtyping : public WalkerPass<PostWalker<LocalSubtyping>> {
       // type.
 
       for (Index i = varBase; i < numLocals; i++) {
+        auto oldType = func->getLocalType(i);
+
         // Find all the types assigned to the var, and compute the optimal LUB.
-        std::unordered_set<Type> types;
+        LUBFinder lub;
         for (auto* set : setsForLocal[i]) {
-          types.insert(set->value->type);
+          if (lub.note(set->value) == oldType) {
+            break;
+          }
         }
-        if (types.empty()) {
+        if (!lub.noted()) {
           // Nothing is assigned to this local (other opts will remove it).
           continue;
         }
 
-        auto oldType = func->getLocalType(i);
-        auto newType = Type::getLeastUpperBound(types);
+        auto newType = lub.get();
         assert(newType != Type::none); // in valid wasm there must be a LUB
 
         // Remove non-nullability if we disallow that in locals.