2 files changed, 42 insertions, 56 deletions

diff --git a/src/ir/branch-utils.h b/src/ir/branch-utils.h
index e6a3486ca..6dfe216a2 100644
--- a/src/ir/branch-utils.h
+++ b/src/ir/branch-utils.h
@@ -68,6 +68,25 @@ template<typename T> void operateOnScopeNameUses(Expression* expr, T func) {
 #include "wasm-delegations-fields.h"
 }
 
+// Similar to operateOnScopeNameUses, but also passes in the type that is sent
+// if the branch is taken. The type is none if there is no value.
+template<typename T>
+void operateOnScopeNameUsesAndSentTypes(Expression* expr, T func) {
+  operateOnScopeNameUses(expr, [&](Name& name) {
+    // There isn't a delegate mechanism for getting a sent value, so do a direct
+    // if-else chain. This will need to be updated with new br variants.
+    if (auto* br = expr->dynCast<Break>()) {
+      func(name, br->value ? br->value->type : Type::none);
+    } else if (auto* sw = expr->dynCast<Switch>()) {
+      func(name, sw->value ? sw->value->type : Type::none);
+    } else if (auto* br = expr->dynCast<BrOnExn>()) {
+      func(name, br->sent);
+    } else {
+      WASM_UNREACHABLE("bad br type");
+    }
+  });
+}
+
 // Perform a generic operation on definitions of scope names in an expression.
 // The provided function receives a Name& which it can modify if it needs to.
 template<typename T> void operateOnScopeNameDefs(Expression* expr, T func) {
@@ -164,36 +183,26 @@ struct BranchSeeker
   Name target;
 
   Index found = 0;
-  Type valueType;
+  // None indicates no value is sent.
+  Type valueType = Type::none;
 
   BranchSeeker(Name target) : target(target) {}
 
-  void noteFound(Expression* value) {
-    noteFound(value ? value->type : Type::none);
-  }
-
-  void noteFound(Type type) {
+  void noteFound(Type newType) {
     found++;
-    if (found == 1) {
-      valueType = Type::unreachable;
-    }
-    if (type != Type::unreachable) {
-      valueType = type;
+    if (newType != Type::none) {
+      if (found == 1) {
+        valueType = newType;
+      } else {
+        valueType = Type::getLeastUpperBound(valueType, newType);
+      }
     }
   }
 
   void visitExpression(Expression* curr) {
-    operateOnScopeNameUses(curr, [&](Name& name) {
+    operateOnScopeNameUsesAndSentTypes(curr, [&](Name& name, Type type) {
       if (name == target) {
-        if (auto* br = curr->dynCast<Break>()) {
-          noteFound(br->value);
-        } else if (auto* sw = curr->dynCast<Switch>()) {
-          noteFound(sw->value);
-        } else if (auto* br = curr->dynCast<BrOnExn>()) {
-          noteFound(br->sent);
-        } else {
-          WASM_UNREACHABLE("bad br type");
-        }
+        noteFound(type);
       }
     });
   }
diff --git a/src/ir/type-updating.h b/src/ir/type-updating.h
index 1a1086a0c..1a117c239 100644
--- a/src/ir/type-updating.h
+++ b/src/ir/type-updating.h
@@ -17,6 +17,7 @@
 #ifndef wasm_ir_type_updating_h
 #define wasm_ir_type_updating_h
 
+#include "ir/branch-utils.h"
 #include "wasm-traversal.h"
 
 namespace wasm {
@@ -56,17 +57,11 @@ struct TypeUpdater
       if (block->name.is()) {
         blockInfos[block->name].block = block;
       }
-    } else if (auto* br = curr->dynCast<Break>()) {
-      // ensure info exists, discoverBreaks can then fill it
-      blockInfos[br->name];
-    } else if (auto* sw = curr->dynCast<Switch>()) {
-      // ensure info exists, discoverBreaks can then fill it
-      for (auto target : sw->targets) {
-        blockInfos[target];
-      }
-      blockInfos[sw->default_];
-    } else if (auto* br = curr->dynCast<BrOnExn>()) {
-      blockInfos[br->name];
+    } else {
+      BranchUtils::operateOnScopeNameUses(curr, [&](Name& name) {
+        // ensure info exists, discoverBreaks can then fill it
+        blockInfos[name];
+      });
     }
     // add a break to the info, for break and switch
     discoverBreaks(curr, +1);
@@ -157,30 +152,12 @@ struct TypeUpdater
 
   // adds (or removes) breaks depending on break/switch contents
   void discoverBreaks(Expression* curr, int change) {
-    if (auto* br = curr->dynCast<Break>()) {
-      noteBreakChange(br->name, change, br->value);
-    } else if (auto* sw = curr->dynCast<Switch>()) {
-      applySwitchChanges(sw, change);
-    } else if (auto* br = curr->dynCast<BrOnExn>()) {
-      noteBreakChange(br->name, change, br->sent);
-    }
-  }
-
-  void applySwitchChanges(Switch* sw, int change) {
-    std::set<Name> seen;
-    for (auto target : sw->targets) {
-      if (seen.insert(target).second) {
-        noteBreakChange(target, change, sw->value);
-      }
-    }
-    if (seen.insert(sw->default_).second) {
-      noteBreakChange(sw->default_, change, sw->value);
-    }
-  }
-
-  // note the addition of a node
-  void noteBreakChange(Name name, int change, Expression* value) {
-    noteBreakChange(name, change, value ? value->type : Type::none);
+    BranchUtils::operateOnScopeNameUsesAndSentTypes(
+      curr,
+      [&](Name& name, Type type) { noteBreakChange(name, change, type); });
+    // TODO: it may be faster to accumulate all changes to a set first, then
+    // call noteBreakChange on the unique values, as a switch can be quite
+    // large and have lots of repeated targets.
   }
 
   void noteBreakChange(Name name, int change, Type type) {