BranchUtils additions (#3860)

Add operateOnScopeNameUsesAndSentValues which is like the
existing utilities, but provides the child expression that is sent as a
value on the branch.

Also provide a BranchTargets utility that maps target names to the
control flow structure for them.

Both will be used and tested in escape analysis.

1 files changed, 40 insertions, 0 deletions

diff --git a/src/ir/branch-utils.h b/src/ir/branch-utils.h
index 414ee0e16..cf21facbc 100644
--- a/src/ir/branch-utils.h
+++ b/src/ir/branch-utils.h
@@ -88,6 +88,25 @@ void operateOnScopeNameUsesAndSentTypes(Expression* expr, T func) {
   });
 }
 
+// Similar to operateOnScopeNameUses, but also passes in the expression that is
+// sent if the branch is taken. nullptr is given if there is no value.
+template<typename T>
+void operateOnScopeNameUsesAndSentValues(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);
+    } else if (auto* sw = expr->dynCast<Switch>()) {
+      func(name, sw->value);
+    } else if (auto* br = expr->dynCast<BrOn>()) {
+      func(name, br->ref);
+    } else {
+      assert(expr->is<Try>() || expr->is<Rethrow>()); // delegate or rethrow
+    }
+  });
+}
+
 // 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) {
@@ -333,6 +352,27 @@ public:
   }
 };
 
+// Provides a map of branch target names to the expressions that branch to
+// them.
+struct BranchTargets {
+  BranchTargets(Expression* expr) { inner.walk(expr); }
+
+  Expression* getTarget(Name name) { return inner.map[name]; }
+
+private:
+  struct Inner : public PostWalker<Inner, UnifiedExpressionVisitor<Inner>> {
+    void visitExpression(Expression* curr) {
+      operateOnScopeNameDefs(curr, [&](Name name) {
+        if (name.is()) {
+          map[name] = curr;
+        }
+      });
+    }
+
+    std::map<Name, Expression*> map;
+  } inner;
+};
+
 } // namespace BranchUtils
 
 } // namespace wasm