[Wasm GC] Heap2Local: Handle branches (#3881)

If we branch to a block, and there are no other branches or a final
value on the block either, then there is no mixing, and we may be able
to optimize the allocation. Before this PR, all branches stopped us.

To do this, add some helpers in BranchUtils.

The main flow logic in Heap2Local used to stop when we reached a
child for the second time. With branches, however, a child can flow both to
its immediate parent, and to branch targets, and so the proper thing to look
at is when we reach a parent for the second time (which would definitely
indicate mixing).

Tests are added for the new functionality. Note that some existing tests
already covered some things we should not optimize, and so no tests
were needed for them. The existing ones are: $get-through-block,
$branch-to-block.

2 files changed, 62 insertions, 16 deletions

diff --git a/src/ir/branch-utils.h b/src/ir/branch-utils.h
index cf21facbc..662f70725 100644
--- a/src/ir/branch-utils.h
+++ b/src/ir/branch-utils.h
@@ -236,6 +236,22 @@ inline NameSet getBranchTargets(Expression* ast) {
   return scanner.targets;
 }
 
+// Get the name of the branch target that is defined in the expression, or an
+// empty name if there is none.
+inline Name getDefinedName(Expression* curr) {
+  Name ret;
+  operateOnScopeNameDefs(curr, [&](Name& name) { ret = name; });
+  return ret;
+}
+
+// Return the value sent by a branch instruction, or nullptr if there is none.
+inline Expression* getSentValue(Expression* curr) {
+  Expression* ret = nullptr;
+  operateOnScopeNameUsesAndSentValues(
+    curr, [&](Name name, Expression* value) { ret = value; });
+  return ret;
+}
+
 // Finds if there are branches targeting a name. Note that since names are
 // unique in our IR, we just need to look for the name, and do not need
 // to analyze scoping.
@@ -352,24 +368,41 @@ public:
   }
 };
 
-// Provides a map of branch target names to the expressions that branch to
-// them.
+// Stores information about branch targets, specifically, finding them by their
+// name, and finding the branches to them.
 struct BranchTargets {
   BranchTargets(Expression* expr) { inner.walk(expr); }
 
-  Expression* getTarget(Name name) { return inner.map[name]; }
+  // Gets the expression that defines this branch target, i.e., where we branch
+  // to if we branch to that name.
+  Expression* getTarget(Name name) { return inner.targets[name]; }
+
+  // Gets the expressions branching to a target.
+  std::unordered_set<Expression*> getBranches(Name name) {
+    auto iter = inner.branches.find(name);
+    if (iter != inner.branches.end()) {
+      return iter->second;
+    }
+    return {};
+  }
 
 private:
   struct Inner : public PostWalker<Inner, UnifiedExpressionVisitor<Inner>> {
     void visitExpression(Expression* curr) {
       operateOnScopeNameDefs(curr, [&](Name name) {
         if (name.is()) {
-          map[name] = curr;
+          targets[name] = curr;
+        }
+      });
+      operateOnScopeNameUses(curr, [&](Name& name) {
+        if (name.is()) {
+          branches[name].insert(curr);
         }
       });
     }
 
-    std::map<Name, Expression*> map;
+    std::map<Name, Expression*> targets;
+    std::map<Name, std::unordered_set<Expression*>> branches;
   } inner;
 };
 
diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index 7e8f12990..afdcab06f 100644
--- a/src/passes/Heap2Local.cpp
+++ b/src/passes/Heap2Local.cpp
@@ -394,6 +394,9 @@ struct Heap2LocalOptimizer {
     }
   };
 
+  // All the expressions we have already looked at.
+  std::unordered_set<Expression*> seen;
+
   enum class ParentChildInteraction {
     // The parent lets the child escape. E.g. the parent is a call.
     Escapes,
@@ -444,10 +447,10 @@ struct Heap2LocalOptimizer {
       // look at something that another allocation reached, which would be in a
       // different call to this function and use a different queue (any overlap
       // between calls would prove non-exclusivity).
-      if (seen.count(child)) {
+      if (seen.count(parent)) {
         return false;
       }
-      seen.insert(child);
+      seen.insert(parent);
 
       switch (getParentChildInteraction(parent, child)) {
         case ParentChildInteraction::Escapes: {
@@ -490,7 +493,7 @@ struct Heap2LocalOptimizer {
       // If the parent may send us on a branch, we will need to look at the flow
       // to the branch target(s).
       for (auto name : branchesSentByParent(child, parent)) {
-        flows.push({parent, branchTargets.getTarget(name)});
+        flows.push({child, branchTargets.getTarget(name)});
       }
 
       // If we got to here, then we can continue to hope that we can optimize
@@ -509,9 +512,6 @@ struct Heap2LocalOptimizer {
     return true;
   }
 
-  // All the expressions we have already looked at.
-  std::unordered_set<Expression*> seen;
-
   ParentChildInteraction getParentChildInteraction(Expression* parent,
                                                    Expression* child) {
     // If there is no parent then we are the body of the function, and that
@@ -611,16 +611,29 @@ struct Heap2LocalOptimizer {
 
     // Finally, check for mixing. If the child is the immediate fallthrough
     // of the parent then no other values can be mixed in.
-    //
-    // TODO: Also check if we are sent via a branch (and that branch sends the
-    //       single value exiting the parent).
-    // TODO: Also check for safe merges where our allocation is in all places,
-    //       like two if or select arms, or branches.
     if (Properties::getImmediateFallthrough(
           parent, passOptions, module->features) == child) {
       return ParentChildInteraction::Flows;
     }
 
+    // Likewise, if the child branches to the parent, and it is the sole branch,
+    // with no other value exiting the block (in particular, no final value at
+    // the end that flows out), then there is no mixing.
+    auto branches =
+      branchTargets.getBranches(BranchUtils::getDefinedName(parent));
+    if (branches.size() == 1 &&
+        BranchUtils::getSentValue(*branches.begin()) == child) {
+      // TODO: support more types of branch targets.
+      if (auto* parentAsBlock = parent->dynCast<Block>()) {
+        if (parentAsBlock->list.back()->type == Type::unreachable) {
+          return ParentChildInteraction::Flows;
+        }
+      }
+    }
+
+    // TODO: Also check for safe merges where our allocation is in all places,
+    //       like two if or select arms, or branches.
+
     return ParentChildInteraction::Mixes;
   }