More refactoring of branch utility code to remove boilerplate. (#3448)

This is almost NFC, but it may emit slightly different IR in cases that
don't matter much. Specifically,

(block (result i32) ;; can also be unreachable
  (unreachable)
  (i32.const 1)
)

That can be finalized to have type unreachable or i32, as both are
valid. After this PR we should consistently do the same thing in all
places. (Either option would be ok - we prefer to keep the type if
there is one.)

In practice, DCE will remove all the dead code anyhow, leaving no
difference to matter. However, the IR is different without DCE, and
that may be noticeable in an unoptimized build - but it should have
no effect on behavior, just on the binary.

5 files changed, 83 insertions, 158 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) {
diff --git a/src/wasm.h b/src/wasm.h
index bc39eae7b..37c12ffc6 100644
--- a/src/wasm.h
+++ b/src/wasm.h
@@ -703,11 +703,13 @@ public:
   // needed (which may require scanning the block)
   void finalize(Type type_);
 
+  enum Breakability { Unknown, HasBreak, NoBreak };
+
   // set the type given you know its type, and you know if there is a break to
   // this block. this avoids the need to scan the contents of the block in the
   // case that it might be unreachable, so it is recommended if you already know
   // the type and breakability anyhow.
-  void finalize(Type type_, bool hasBreak);
+  void finalize(Type type_, Breakability breakability);
 };
 
 class If : public SpecificExpression<Expression::IfId> {
diff --git a/src/wasm/wasm-binary.cpp b/src/wasm/wasm-binary.cpp
index 66715686e..a4d652608 100644
--- a/src/wasm/wasm-binary.cpp
+++ b/src/wasm/wasm-binary.cpp
@@ -3154,8 +3154,9 @@ void WasmBinaryBuilder::visitBlock(Block* curr) {
     }
     pushBlockElements(curr, curr->type, start);
     curr->finalize(curr->type,
-                   breakTargetNames.find(curr->name) !=
-                     breakTargetNames.end() /* hasBreak */);
+                   breakTargetNames.find(curr->name) != breakTargetNames.end()
+                     ? Block::HasBreak
+                     : Block::NoBreak);
     breakStack.pop_back();
     breakTargetNames.erase(curr->name);
   }
diff --git a/src/wasm/wasm.cpp b/src/wasm/wasm.cpp
index 851e6a804..d9ca03185 100644
--- a/src/wasm/wasm.cpp
+++ b/src/wasm/wasm.cpp
@@ -267,72 +267,12 @@ Literals getLiteralsFromConstExpression(Expression* curr) {
   }
 }
 
-// core AST type checking
-
-struct TypeSeeker : public PostWalker<TypeSeeker> {
-  Expression* target; // look for this one
-  Name targetName;
-  std::vector<Type> types;
-
-  TypeSeeker(Expression* target, Name targetName)
-    : target(target), targetName(targetName) {
-    Expression* temp = target;
-    walk(temp);
-  }
-
-  void visitBreak(Break* curr) {
-    if (curr->name == targetName) {
-      types.push_back(curr->value ? curr->value->type : Type::none);
-    }
-  }
-
-  void visitSwitch(Switch* curr) {
-    for (auto name : curr->targets) {
-      if (name == targetName) {
-        types.push_back(curr->value ? curr->value->type : Type::none);
-      }
-    }
-    if (curr->default_ == targetName) {
-      types.push_back(curr->value ? curr->value->type : Type::none);
-    }
-  }
-
-  void visitBrOnExn(BrOnExn* curr) {
-    if (curr->name == targetName) {
-      types.push_back(curr->sent);
-    }
-  }
-
-  void visitBlock(Block* curr) {
-    if (curr == target) {
-      if (curr->list.size() > 0) {
-        types.push_back(curr->list.back()->type);
-      } else {
-        types.push_back(Type::none);
-      }
-    } else if (curr->name == targetName) {
-      // ignore all breaks til now, they were captured by someone with the same
-      // name
-      types.clear();
-    }
-  }
-
-  void visitLoop(Loop* curr) {
-    if (curr == target) {
-      types.push_back(curr->body->type);
-    } else if (curr->name == targetName) {
-      // ignore all breaks til now, they were captured by someone with the same
-      // name
-      types.clear();
-    }
-  }
-};
-
 // a block is unreachable if one of its elements is unreachable,
 // and there are no branches to it
-static void handleUnreachable(Block* block,
-                              bool breakabilityKnown = false,
-                              bool hasBreak = false) {
+
+static void
+handleUnreachable(Block* block,
+                  Block::Breakability breakability = Block::Unknown) {
   if (block->type == Type::unreachable) {
     return; // nothing to do
   }
@@ -350,10 +290,12 @@ static void handleUnreachable(Block* block,
     if (child->type == Type::unreachable) {
       // there is an unreachable child, so we are unreachable, unless we have a
       // break
-      if (!breakabilityKnown) {
-        hasBreak = BranchUtils::BranchSeeker::has(block, block->name);
+      if (breakability == Block::Unknown) {
+        breakability = BranchUtils::BranchSeeker::has(block, block->name)
+                         ? Block::HasBreak
+                         : Block::NoBreak;
       }
-      if (!hasBreak) {
+      if (breakability == Block::NoBreak) {
         block->type = Type::unreachable;
       }
       return;
@@ -362,41 +304,35 @@ static void handleUnreachable(Block* block,
 }
 
 void Block::finalize() {
+  if (list.size() == 0) {
+    type = Type::none;
+    return;
+  }
+  // The default type is what is at the end. Next we need to see if breaks and/
+  // or unreachabitily change that.
+  type = list.back()->type;
   if (!name.is()) {
-    if (list.size() > 0) {
-      // nothing branches here, so this is easy
-      // normally the type is the type of the final child
-      type = list.back()->type;
-      // and even if we have an unreachable child somewhere,
-      // we still mark ourselves as having that type,
-      // (block (result i32)
-      //  (return)
-      //  (i32.const 10)
-      // )
-      if (type.isConcrete()) {
-        return;
-      }
-      // if we are unreachable, we are done
-      if (type == Type::unreachable) {
-        return;
-      }
-      // we may still be unreachable if we have an unreachable
-      // child
-      for (auto* child : list) {
-        if (child->type == Type::unreachable) {
-          type = Type::unreachable;
-          return;
-        }
-      }
+    // Nothing branches here, so this is easy.
+    handleUnreachable(this, NoBreak);
+    return;
+  }
+
+  // The default type is according to the value that flows out.
+  BranchUtils::BranchSeeker seeker(this->name);
+  Expression* temp = this;
+  seeker.walk(temp);
+  if (seeker.found) {
+    // Take the branch values into account.
+    if (seeker.valueType != Type::none) {
+      type = Type::getLeastUpperBound(type, seeker.valueType);
     } else {
+      // No value is sent, but as we have a branch we are not unreachable.
       type = Type::none;
     }
-    return;
+  } else {
+    // There are no branches, so this block may be unreachable.
+    handleUnreachable(this, NoBreak);
   }
-
-  TypeSeeker seeker(this, this->name);
-  type = Type::mergeTypes(seeker.types);
-  handleUnreachable(this);
 }
 
 void Block::finalize(Type type_) {
@@ -406,10 +342,10 @@ void Block::finalize(Type type_) {
   }
 }
 
-void Block::finalize(Type type_, bool hasBreak) {
+void Block::finalize(Type type_, Breakability breakability) {
   type = type_;
   if (type == Type::none && list.size() > 0) {
-    handleUnreachable(this, true, hasBreak);
+    handleUnreachable(this, breakability);
   }
 }