Heap2Local: Track interactions in detail (#6834)

Previously we tracked only whether an expression was relevant to analysis, that is,
whether it interacted with the allocation we were tracing the behavior of. That is
not enough for all cases, though, so also track the form of the interaction, namely
whether the allocation flows through or is fully consumed. An example where that
matters:

(ref.eq
  (struct.get $A 0
    (local.tee $x
      (struct.new_default $A)
    )
  )
  (local.get $x)
)

Here the local.get flows out the allocation, but the struct.get only fully consumes
it. Before this PR we thought the struct.get flowed the allocation, and we misoptimized
this to 1.

To make this possible, do a bunch of minor refactoring:

 * Move ParentChildInteraction out of the class.
 * Add a "None" interaction there.
 * Replace the set of reached expressions with a map of them to their interactions.
 * Add helper functions to get an expression's interaction or to update it when replacing.

The new testcase here shows the main fix. The new assertions are covered by existing
testcases.

src/passes/Heap2Local.cpp

@@ -167,6 +167,30 @@ namespace wasm {
 
 namespace {
 
+// Interactions between a child and a parent, with regard to the behavior of the
+// allocation.
+enum class ParentChildInteraction : int8_t {
+  // The parent lets the child escape. E.g. the parent is a call.
+  Escapes,
+  // The parent fully consumes the child in a safe, non-escaping way, and
+  // after consuming it nothing remains to flow further through the parent.
+  // E.g. the parent is a struct.get, which reads from the allocated heap
+  // value and does nothing more with the reference.
+  FullyConsumes,
+  // The parent flows the child out, that is, the child is the single value
+  // that can flow out from the parent. E.g. the parent is a block with no
+  // branches and the child is the final value that is returned.
+  Flows,
+  // The parent does not consume the child completely, so the child's value
+  // can be used through it. However the child does not flow cleanly through.
+  // E.g. the parent is a block with branches, and the value on them may be
+  // returned from the block and not only the child. This means the allocation
+  // is not used in an exclusive way, and we cannot optimize it.
+  Mixes,
+  // No interaction (not relevant to the analysis).
+  None,
+};
+
 // Core analysis that provides an escapes() method to check if an allocation
 // escapes in a way that prevents optimizing it away as described above. It also
 // stashes information about the relevant expressions as it goes, which helps
@@ -194,30 +218,16 @@ struct EscapeAnalyzer {
   // exclusivity.
   std::unordered_set<LocalSet*> sets;
 
-  // All the expressions we reached during the flow analysis. That is exactly
-  // all the places where our allocation is used. We track these so that we
-  // can fix them up at the end, if the optimization ends up possible.
-  std::unordered_set<Expression*> reached;
-
-  enum class ParentChildInteraction {
-    // The parent lets the child escape. E.g. the parent is a call.
-    Escapes,
-    // The parent fully consumes the child in a safe, non-escaping way, and
-    // after consuming it nothing remains to flow further through the parent.
-    // E.g. the parent is a struct.get, which reads from the allocated heap
-    // value and does nothing more with the reference.
-    FullyConsumes,
-    // The parent flows the child out, that is, the child is the single value
-    // that can flow out from the parent. E.g. the parent is a block with no
-    // branches and the child is the final value that is returned.
-    Flows,
-    // The parent does not consume the child completely, so the child's value
-    // can be used through it. However the child does not flow cleanly through.
-    // E.g. the parent is a block with branches, and the value on them may be
-    // returned from the block and not only the child. This means the allocation
-    // is not used in an exclusive way, and we cannot optimize it.
-    Mixes,
-  };
+  // A map of every expression we reached during the flow analysis (which is
+  // exactly all the places where our allocation is used) to the interaction of
+  // the allocation there. If we optimize, anything in this map will be fixed up
+  // at the end, and how we fix it up may depend on the interaction,
+  // specifically, it can matter if the allocations flows out of here (Flows,
+  // which is the case for e.g. a Block that we flow through) or if it is fully
+  // consumed (FullyConsumes, e.g. for a struct.get). We do not store irrelevant
+  // things here (that is, anything not in the map has the interaction |None|,
+  // implicitly).
+  std::unordered_map<Expression*, ParentChildInteraction> reachedInteractions;
 
   // Analyze an allocation to see if it escapes or not.
   bool escapes(Expression* allocation) {
@@ -284,9 +294,10 @@ struct EscapeAnalyzer {
 
       // If we got to here, then we can continue to hope that we can optimize
       // this allocation. Mark the parent and child as reached by it, and
-      // continue.
-      reached.insert(parent);
-      reached.insert(child);
+      // continue. The child flows the value to the parent, and the parent's
+      // behavior was computed before.
+      reachedInteractions[child] = ParentChildInteraction::Flows;
+      reachedInteractions[parent] = interaction;
     }
 
     // We finished the loop over the flows. Do the final checks.
@@ -300,7 +311,7 @@ struct EscapeAnalyzer {
 
   ParentChildInteraction getParentChildInteraction(Expression* allocation,
                                                    Expression* parent,
-                                                   Expression* child) {
+                                                   Expression* child) const {
     // If there is no parent then we are the body of the function, and that
     // means we escape by flowing to the caller.
     if (!parent) {
@@ -517,6 +528,36 @@ struct EscapeAnalyzer {
 
     return true;
   }
+
+  // Helper function for Struct2Local and Array2Struct. Given an old expression
+  // that is being replaced by a new one, add the proper interaction for the
+  // replacement.
+  void applyOldInteractionToReplacement(Expression* old, Expression* rep) {
+    // We can only replace something relevant that we found in the analysis.
+    // (Not only would anything else be invalid to process, but also we wouldn't
+    // know what interaction to give the replacement.)
+    assert(reachedInteractions.count(old));
+
+    // The replacement should have the same interaction as the thing it
+    // replaces, since it is a drop-in replacement for it. The one exception is
+    // when we replace with something unreachable, which is the result of us
+    // figuring out that some code will trap at runtime. In that case, we've
+    // made the code unreachable and the allocation does not interact with that
+    // code at all.
+    if (rep->type != Type::unreachable) {
+      reachedInteractions[rep] = reachedInteractions[old];
+    }
+  }
+
+  // Get the interaction of an expression.
+  ParentChildInteraction getInteraction(Expression* curr) {
+    auto iter = reachedInteractions.find(curr);
+    if (iter == reachedInteractions.end()) {
+      // This is not interacted with.
+      return ParentChildInteraction::None;
+    }
+    return iter->second;
+  }
 };
 
 // An optimizer that handles the rewriting to turn a struct allocation into
@@ -527,8 +568,8 @@ struct EscapeAnalyzer {
 struct Struct2Local : PostWalker<Struct2Local> {
   StructNew* allocation;
 
-  // The analyzer is not |const| because we update |analyzer.reached| as we go
-  // (see replaceCurrent, below).
+  // The analyzer is not |const| because we update
+  // |analyzer.reachedInteractions| as we go (see replaceCurrent, below).
   EscapeAnalyzer& analyzer;
 
   Function* func;
@@ -563,11 +604,8 @@ struct Struct2Local : PostWalker<Struct2Local> {
   bool refinalize = false;
 
   Expression* replaceCurrent(Expression* expression) {
+    analyzer.applyOldInteractionToReplacement(getCurrent(), expression);
     PostWalker<Struct2Local>::replaceCurrent(expression);
-    // Also update |reached|: we are replacing something that was reached, so
-    // logically the replacement is also reached. This update is necessary if
-    // the parent of an expression cares about whether a child was reached.
-    analyzer.reached.insert(expression);
     return expression;
   }
 
@@ -583,7 +621,7 @@ struct Struct2Local : PostWalker<Struct2Local> {
   // Adjust the type that flows through an expression, updating that type as
   // necessary.
   void adjustTypeFlowingThrough(Expression* curr) {
-    if (!analyzer.reached.count(curr)) {
+    if (analyzer.getInteraction(curr) != ParentChildInteraction::Flows) {
       return;
     }
 
@@ -603,7 +641,7 @@ struct Struct2Local : PostWalker<Struct2Local> {
   void visitLoop(Loop* curr) { adjustTypeFlowingThrough(curr); }
 
   void visitLocalSet(LocalSet* curr) {
-    if (!analyzer.reached.count(curr)) {
+    if (analyzer.getInteraction(curr) == ParentChildInteraction::None) {
       return;
     }
 
@@ -617,7 +655,7 @@ struct Struct2Local : PostWalker<Struct2Local> {
   }
 
   void visitLocalGet(LocalGet* curr) {
-    if (!analyzer.reached.count(curr)) {
+    if (analyzer.getInteraction(curr) == ParentChildInteraction::None) {
       return;
     }
 
@@ -639,7 +677,7 @@ struct Struct2Local : PostWalker<Struct2Local> {
   }
 
   void visitBreak(Break* curr) {
-    if (!analyzer.reached.count(curr)) {
+    if (analyzer.getInteraction(curr) == ParentChildInteraction::None) {
       return;
     }
 
@@ -721,7 +759,7 @@ struct Struct2Local : PostWalker<Struct2Local> {
   }
 
   void visitRefIsNull(RefIsNull* curr) {
-    if (!analyzer.reached.count(curr)) {
+    if (analyzer.getInteraction(curr) == ParentChildInteraction::None) {
       return;
     }
 
@@ -732,7 +770,7 @@ struct Struct2Local : PostWalker<Struct2Local> {
   }
 
   void visitRefEq(RefEq* curr) {
-    if (!analyzer.reached.count(curr)) {
+    if (analyzer.getInteraction(curr) == ParentChildInteraction::None) {
       return;
     }
 
@@ -745,15 +783,16 @@ struct Struct2Local : PostWalker<Struct2Local> {
     // If our reference is compared to itself, the result is 1. If it is
     // compared to something else, the result must be 0, as our reference does
     // not escape to any other place.
-    int32_t result = analyzer.reached.count(curr->left) > 0 &&
-                     analyzer.reached.count(curr->right) > 0;
+    int32_t result =
+      analyzer.getInteraction(curr->left) == ParentChildInteraction::Flows &&
+      analyzer.getInteraction(curr->right) == ParentChildInteraction::Flows;
     replaceCurrent(builder.makeBlock({builder.makeDrop(curr->left),
                                       builder.makeDrop(curr->right),
                                       builder.makeConst(Literal(result))}));
   }
 
   void visitRefAs(RefAs* curr) {
-    if (!analyzer.reached.count(curr)) {
+    if (analyzer.getInteraction(curr) == ParentChildInteraction::None) {
       return;
     }
 
@@ -764,7 +803,7 @@ struct Struct2Local : PostWalker<Struct2Local> {
   }
 
   void visitRefTest(RefTest* curr) {
-    if (!analyzer.reached.count(curr)) {
+    if (analyzer.getInteraction(curr) == ParentChildInteraction::None) {
       return;
     }
 
@@ -781,7 +820,7 @@ struct Struct2Local : PostWalker<Struct2Local> {
   }
 
   void visitRefCast(RefCast* curr) {
-    if (!analyzer.reached.count(curr)) {
+    if (analyzer.getInteraction(curr) == ParentChildInteraction::None) {
       return;
     }
 
@@ -805,7 +844,7 @@ struct Struct2Local : PostWalker<Struct2Local> {
   }
 
   void visitStructSet(StructSet* curr) {
-    if (!analyzer.reached.count(curr)) {
+    if (analyzer.getInteraction(curr) == ParentChildInteraction::None) {
       return;
     }
 
@@ -817,7 +856,7 @@ struct Struct2Local : PostWalker<Struct2Local> {
   }
 
   void visitStructGet(StructGet* curr) {
-    if (!analyzer.reached.count(curr)) {
+    if (analyzer.getInteraction(curr) == ParentChildInteraction::None) {
       return;
     }
 
@@ -913,13 +952,14 @@ struct Array2Struct : PostWalker<Array2Struct> {
       WASM_UNREACHABLE("bad allocation");
     }
 
-    // Mark the new expressions we created as reached by the analysis. We need
-    // to inform the analysis of this because Struct2Local will only process
-    // such code (it depends on the analysis to tell it what the allocation is
-    // and where it flowed). Note that the two values here may be identical but
-    // there is no harm to calling it twice in that case.
-    noteIsReached(structNew);
-    noteIsReached(arrayNewReplacement);
+    // Mark new expressions we created as flowing out the allocation. We need to
+    // inform the analysis of this because Struct2Local will only process such
+    // code (it depends on the analysis to tell it what the allocation is and
+    // where it flowed). Note that the two values here may be identical but
+    // there is no harm to doing this twice in that case.
+    analyzer.reachedInteractions[structNew] =
+      analyzer.reachedInteractions[arrayNewReplacement] =
+        ParentChildInteraction::Flows;
 
     // Update types along the path reached by the allocation: whenever we see
     // the array type, it should be the struct type. Note that we do this before
@@ -935,7 +975,7 @@ struct Array2Struct : PostWalker<Array2Struct> {
     auto nonNullArray = Type(arrayType, NonNullable);
     nullStruct = Type(structType, Nullable);
     nonNullStruct = Type(structType, NonNullable);
-    for (auto* reached : analyzer.reached) {
+    for (auto& [reached, _] : analyzer.reachedInteractions) {
       if (reached->is<RefCast>()) {
         // Casts must be handled later: We need to see the old type, and to
         // potentially replace the cast based on that, see below.
@@ -971,6 +1011,12 @@ struct Array2Struct : PostWalker<Array2Struct> {
     }
   }
 
+  Expression* replaceCurrent(Expression* expression) {
+    analyzer.applyOldInteractionToReplacement(getCurrent(), expression);
+    PostWalker<Array2Struct>::replaceCurrent(expression);
+    return expression;
+  }
+
   // In rare cases we may need to refinalize, as with Struct2Local.
   bool refinalize = false;
 
@@ -989,19 +1035,17 @@ struct Array2Struct : PostWalker<Array2Struct> {
   void visitArrayNew(ArrayNew* curr) {
     if (curr == allocation) {
       replaceCurrent(arrayNewReplacement);
-      noteCurrentIsReached();
     }
   }
 
   void visitArrayNewFixed(ArrayNewFixed* curr) {
     if (curr == allocation) {
       replaceCurrent(arrayNewReplacement);
-      noteCurrentIsReached();
     }
   }
 
   void visitArraySet(ArraySet* curr) {
-    if (!analyzer.reached.count(curr)) {
+    if (analyzer.getInteraction(curr) == ParentChildInteraction::None) {
       return;
     }
 
@@ -1018,11 +1062,10 @@ struct Array2Struct : PostWalker<Array2Struct> {
 
     // Convert the ArraySet into a StructSet.
     replaceCurrent(builder.makeStructSet(index, curr->ref, curr->value));
-    noteCurrentIsReached();
   }
 
   void visitArrayGet(ArrayGet* curr) {
-    if (!analyzer.reached.count(curr)) {
+    if (analyzer.getInteraction(curr) == ParentChildInteraction::None) {
       return;
     }
 
@@ -1039,13 +1082,12 @@ struct Array2Struct : PostWalker<Array2Struct> {
     // Convert the ArrayGet into a StructGet.
     replaceCurrent(
       builder.makeStructGet(index, curr->ref, curr->type, curr->signed_));
-    noteCurrentIsReached();
   }
 
   // Some additional operations need special handling
 
   void visitRefTest(RefTest* curr) {
-    if (!analyzer.reached.count(curr)) {
+    if (analyzer.getInteraction(curr) == ParentChildInteraction::None) {
       return;
     }
 
@@ -1061,7 +1103,7 @@ struct Array2Struct : PostWalker<Array2Struct> {
   }
 
   void visitRefCast(RefCast* curr) {
-    if (!analyzer.reached.count(curr)) {
+    if (analyzer.getInteraction(curr) == ParentChildInteraction::None) {
       return;
     }
 
@@ -1089,13 +1131,6 @@ struct Array2Struct : PostWalker<Array2Struct> {
     return curr->cast<Const>()->value.getUnsigned();
   }
 
-  // Inform the analyzer that the current expression (which we just replaced)
-  // has been reached in its analysis. We are replacing something it reached,
-  // and want it to consider it as its equivalent.
-  void noteCurrentIsReached() { noteIsReached(getCurrent()); }
-
-  void noteIsReached(Expression* curr) { analyzer.reached.insert(curr); }
-
   // Given an ArrayNew or ArrayNewFixed, return the size of the array that is
   // being allocated.
   Index getArrayNewSize(Expression* allocation) {