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;
   }