Heap2Local: Fix an ordering issue with children having different interactions with a parent (#6089)

We had a simple rule that if we reach an expression twice then we give up, which makes
sense for say a block: if one allocation flows out of it, then another can't - it would get
mixed in with the other one, which is a case we don't optimize. However, there are
cases where a parent has multiple children and different interactions with them, like
a struct.set: the reference child does not escape, but the value child does. Before this
PR if we reached the value child first, we'd mark the parent as seen, and then the reference
child would see it isn't the first to get here, and not optimize.

To fix this, reorder the code to handle this case. The manner of interaction between the
child and the parent decides whether we mark the parent as seen and to be further
avoided.

Noticed by the determinism fuzzer, since the order of analysis mattered here.

1 files changed, 33 insertions, 21 deletions

diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index f0f691299..0f569374b 100644
--- a/src/passes/Heap2Local.cpp
+++ b/src/passes/Heap2Local.cpp
@@ -541,6 +541,19 @@ struct Heap2LocalOptimizer {
       auto* child = flow.first;
       auto* parent = flow.second;
 
+      auto interaction = getParentChildInteraction(allocation, parent, child);
+      if (interaction == ParentChildInteraction::Escapes ||
+          interaction == ParentChildInteraction::Mixes) {
+        // If the parent may let us escape, or the parent mixes other values
+        // up with us, give up.
+        return;
+      }
+
+      // The parent either fully consumes us, or flows us onwards; either way,
+      // we can proceed here, hopefully.
+      assert(interaction == ParentChildInteraction::FullyConsumes ||
+             interaction == ParentChildInteraction::Flows);
+
       // If we've already seen an expression, stop since we cannot optimize
       // things that overlap in any way (see the notes on exclusivity, above).
       // Note that we use a nonrepeating queue here, so we already do not visit
@@ -548,31 +561,30 @@ 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).
+      //
+      // Note that we do this after the check for Escapes/Mixes above: it is
+      // possible for a parent to receive two children and handle them
+      // differently:
+      //
+      //  (struct.set
+      //    (local.get $ref)
+      //    (local.get $value)
+      //  )
+      //
+      // The value escapes, but the ref does not, and might be optimized. If we
+      // added the parent to |seen| for both children, the reference would get
+      // blocked from being optimized.
       if (!seen.emplace(parent).second) {
         return;
       }
 
-      switch (getParentChildInteraction(allocation, parent, child)) {
-        case ParentChildInteraction::Escapes: {
-          // If the parent may let us escape then we are done.
-          return;
-        }
-        case ParentChildInteraction::FullyConsumes: {
-          // If the parent consumes us without letting us escape then all is
-          // well (and there is nothing flowing from the parent to check).
-          break;
-        }
-        case ParentChildInteraction::Flows: {
-          // The value flows through the parent; we need to look further at the
-          // grandparent.
-          flows.push({parent, parents.getParent(parent)});
-          break;
-        }
-        case ParentChildInteraction::Mixes: {
-          // Our allocation is not used exclusively via the parent, as other
-          // values are mixed with it. Give up.
-          return;
-        }
+      // We can proceed, as the parent interacts with us properly, and we are
+      // the only allocation to get here.
+
+      if (interaction == ParentChildInteraction::Flows) {
+        // The value flows through the parent; we need to look further at the
+        // grandparent.
+        flows.push({parent, parents.getParent(parent)});
       }
 
       if (auto* set = parent->dynCast<LocalSet>()) {