From a4e177899d1310e91bf906add718853734298dd1 Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Fri, 21 Jun 2024 16:32:16 -0700
Subject: [PATCH 01/21] almost

---
 src/passes/Heap2Local.cpp       |  24 +++++++
 test/lit/passes/heap2local.wast | 122 +++++++++++++++++++++++++++++---
 2 files changed, 138 insertions(+), 8 deletions(-)

diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index 9d1763ddc1b..5494abf8a22 100644
--- a/src/passes/Heap2Local.cpp
+++ b/src/passes/Heap2Local.cpp
@@ -382,6 +382,12 @@ struct EscapeAnalyzer {
       void visitLocalSet(LocalSet* curr) { escapes = false; }
 
       // Reference operations. TODO add more
+      void visitRefEq(RefEq* curr) {
+        // The reference is compared for identity, but nothing more.
+        escapes = false;
+        fullyConsumes = true;
+      }
+
       void visitRefAs(RefAs* curr) {
         // TODO General OptimizeInstructions integration, that is, since we know
         //      that our allocation is what flows into this RefAs, we can
@@ -721,6 +727,24 @@ struct Struct2Local : PostWalker<Struct2Local> {
     replaceCurrent(builder.makeBlock(contents));
   }
 
+  void visitRefEq(RefEq* curr) {
+    if (!analyzer.reached.count(curr)) {
+      return;
+    }
+
+    // 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;
+    auto* block = builder.makeBlock({
+      builder.makeDrop(curr->left),
+      builder.makeDrop(curr->right),
+      builder.makeConst(Literal(result))
+    });
+    replaceCurrent(block);
+  }
+
   void visitRefAs(RefAs* curr) {
     if (!analyzer.reached.count(curr)) {
       return;
diff --git a/test/lit/passes/heap2local.wast b/test/lit/passes/heap2local.wast
index 80eb9bd2614..e6d5d4d88df 100644
--- a/test/lit/passes/heap2local.wast
+++ b/test/lit/passes/heap2local.wast
@@ -13,9 +13,9 @@
 
   ;; CHECK:      (type $struct.recursive (struct (field (mut (ref null $struct.recursive)))))
 
-  ;; CHECK:      (type $4 (func (param (ref null $struct.A))))
+  ;; CHECK:      (type $4 (func (result i32)))
 
-  ;; CHECK:      (type $5 (func (result i32)))
+  ;; CHECK:      (type $5 (func (param (ref null $struct.A))))
 
   ;; CHECK:      (type $6 (func (result anyref)))
 
@@ -36,6 +36,8 @@
 
   ;; CHECK:      (type $11 (func (param i32)))
 
+  ;; CHECK:      (type $12 (func (param eqref) (result i32)))
+
   ;; CHECK:      (func $simple (type $1)
   ;; CHECK-NEXT:  (local $0 i32)
   ;; CHECK-NEXT:  (local $1 f64)
@@ -474,7 +476,7 @@
     )
   )
 
-  ;; CHECK:      (func $send-ref (type $4) (param $0 (ref null $struct.A))
+  ;; CHECK:      (func $send-ref (type $5) (param $0 (ref null $struct.A))
   ;; CHECK-NEXT:  (nop)
   ;; CHECK-NEXT: )
   (func $send-ref (param (ref null $struct.A))
@@ -887,7 +889,7 @@
     )
   )
 
-  ;; CHECK:      (func $tee (type $5) (result i32)
+  ;; CHECK:      (func $tee (type $4) (result i32)
   ;; CHECK-NEXT:  (local $ref (ref null $struct.A))
   ;; CHECK-NEXT:  (local $1 i32)
   ;; CHECK-NEXT:  (local $2 f64)
@@ -1786,7 +1788,7 @@
     )
   )
 
-  ;; CHECK:      (func $ref-as-non-null-through-local (type $5) (result i32)
+  ;; CHECK:      (func $ref-as-non-null-through-local (type $4) (result i32)
   ;; CHECK-NEXT:  (local $ref (ref null $struct.A))
   ;; CHECK-NEXT:  (local $1 i32)
   ;; CHECK-NEXT:  (local $2 f64)
@@ -1959,7 +1961,7 @@
    (local.get $0)
   )
 
-  ;; CHECK:      (func $to-param (type $4) (param $ref (ref null $struct.A))
+  ;; CHECK:      (func $to-param (type $5) (param $ref (ref null $struct.A))
   ;; CHECK-NEXT:  (local $1 i32)
   ;; CHECK-NEXT:  (local $2 f64)
   ;; CHECK-NEXT:  (drop
@@ -2006,7 +2008,7 @@
     )
   )
 
-  ;; CHECK:      (func $to-param-loop (type $4) (param $ref (ref null $struct.A))
+  ;; CHECK:      (func $to-param-loop (type $5) (param $ref (ref null $struct.A))
   ;; CHECK-NEXT:  (loop $loop
   ;; CHECK-NEXT:   (drop
   ;; CHECK-NEXT:    (struct.get $struct.A 0
@@ -2046,7 +2048,7 @@
     )
   )
 
-  ;; CHECK:      (func $ref-cast (type $5) (result i32)
+  ;; CHECK:      (func $ref-cast (type $4) (result i32)
   ;; CHECK-NEXT:  (local $0 i32)
   ;; CHECK-NEXT:  (local $1 f64)
   ;; CHECK-NEXT:  (local $2 i32)
@@ -2080,6 +2082,110 @@
       )
     )
   )
+
+  ;; CHECK:      (func $ref-eq (type $4) (result i32)
+  ;; CHECK-NEXT:  (local $0 i32)
+  ;; CHECK-NEXT:  (local $1 f64)
+  ;; CHECK-NEXT:  (local $2 i32)
+  ;; CHECK-NEXT:  (local $3 f64)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result nullref)
+  ;; CHECK-NEXT:    (local.set $2
+  ;; CHECK-NEXT:     (i32.const 0)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $3
+  ;; CHECK-NEXT:     (f64.const 0)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $0
+  ;; CHECK-NEXT:     (local.get $2)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $1
+  ;; CHECK-NEXT:     (local.get $3)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (ref.null none)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (ref.null none)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (i32.const 0)
+  ;; CHECK-NEXT: )
+  (func $ref-eq (result i32)
+    ;; Comparing an allocation to something else results in 0, and we can
+    ;; optimize away the allocation.
+    (ref.eq
+      (struct.new $struct.A
+        (i32.const 0)
+        (f64.const 0)
+      )
+      (ref.null eq)
+    )
+  )
+
+  ;; CHECK:      (func $ref-eq-flip (type $12) (param $other eqref) (result i32)
+  ;; CHECK-NEXT:  (local $1 i32)
+  ;; CHECK-NEXT:  (local $2 f64)
+  ;; CHECK-NEXT:  (local $3 i32)
+  ;; CHECK-NEXT:  (local $4 f64)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (ref.null none)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result nullref)
+  ;; CHECK-NEXT:    (local.set $3
+  ;; CHECK-NEXT:     (i32.const 0)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $4
+  ;; CHECK-NEXT:     (f64.const 0)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $1
+  ;; CHECK-NEXT:     (local.get $3)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $2
+  ;; CHECK-NEXT:     (local.get $4)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (ref.null none)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (i32.const 0)
+  ;; CHECK-NEXT: )
+  (func $ref-eq-flip (param $other eqref) (result i32)
+    ;; As above, but flipped, and compared to a local.
+    (ref.eq
+      (ref.null eq)
+      (struct.new $struct.A
+        (i32.const 0)
+        (f64.const 0)
+      )
+    )
+  )
+
+  ;; CHECK:      (func $ref-eq-self (type $4) (result i32)
+  ;; CHECK-NEXT:  (local $eq eqref)
+  ;; CHECK-NEXT:  (ref.eq
+  ;; CHECK-NEXT:   (local.tee $eq
+  ;; CHECK-NEXT:    (struct.new $struct.A
+  ;; CHECK-NEXT:     (i32.const 0)
+  ;; CHECK-NEXT:     (f64.const 0)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (local.get $eq)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $ref-eq-self (result i32)
+    (local $eq eqref)
+    ;; Comparing to oneself results in 1, and we can optimize away the
+    ;; allocation.
+    (ref.eq
+      (local.tee $eq
+        (struct.new $struct.A
+          (i32.const 0)
+          (f64.const 0)
+        )
+      )
+      (local.get $eq)
+    )
+  )
 )
 
 (module

From 38f22977cb571c6ea193e06f0defda8eacf8d0b4 Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Mon, 24 Jun 2024 12:42:05 -0700
Subject: [PATCH 02/21] work

---
 src/passes/Heap2Local.cpp | 60 ++++++++++++++++++++++-----------------
 1 file changed, 34 insertions(+), 26 deletions(-)

diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index 5494abf8a22..9b46413a0d4 100644
--- a/src/passes/Heap2Local.cpp
+++ b/src/passes/Heap2Local.cpp
@@ -248,6 +248,25 @@ struct EscapeAnalyzer {
       auto* child = flow.first;
       auto* parent = flow.second;
 
+      // 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
+      // the same thing more than once; what this check does is verify we don't
+      // 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).
+      //
+      // It is ok, however, to see a parent more than once, if the allocation
+      // flows to it from two children, as is the case for ref.eq. Likewise, for
+      // struct.set, where we also have different considerations for the two
+      // children (the reference does not escape, but the value does). We only
+      // stop here if we've seen this parent in a *different* allocation.
+      auto seenBefore = !seen.emplace(parent).second;
+      auto reachedInThisAllocation = reached.count(parent) > 0;
+      if (seenBefore && !reachedInThisAllocation) {
+        return true;
+      }
+
       auto interaction = getParentChildInteraction(allocation, parent, child);
       if (interaction == ParentChildInteraction::Escapes ||
           interaction == ParentChildInteraction::Mixes) {
@@ -261,30 +280,6 @@ struct EscapeAnalyzer {
       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
-      // the same thing more than once; what this check does is verify we don't
-      // 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 true;
-      }
-
       // We can proceed, as the parent interacts with us properly, and we are
       // the only allocation to get here.
 
@@ -546,14 +541,18 @@ struct EscapeAnalyzer {
 //       efficient, but it would need to be more complex.
 struct Struct2Local : PostWalker<Struct2Local> {
   StructNew* allocation;
-  const EscapeAnalyzer& analyzer;
+
+  // The analyzer is not |const| because we update |analyzer.reached| as we go
+  // (see replaceCurrent, below).
+  EscapeAnalyzer& analyzer;
+
   Function* func;
   Module& wasm;
   Builder builder;
   const FieldList& fields;
 
   Struct2Local(StructNew* allocation,
-               const EscapeAnalyzer& analyzer,
+               EscapeAnalyzer& analyzer,
                Function* func,
                Module& wasm)
     : allocation(allocation), analyzer(analyzer), func(func), wasm(wasm),
@@ -578,6 +577,15 @@ struct Struct2Local : PostWalker<Struct2Local> {
   // In rare cases we may need to refinalize, see below.
   bool refinalize = false;
 
+  Expression* replaceCurrent(Expression* 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;
+  }
+
   // Rewrite the code in visit* methods. The general approach taken is to
   // replace the allocation with a null reference (which may require changing
   // types in some places, like making a block return value nullable), and to

From 3009f94c0c203f3f7ab06ae5a77da6a8f8bae06e Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Mon, 24 Jun 2024 12:46:20 -0700
Subject: [PATCH 03/21] work

---
 src/passes/Heap2Local.cpp | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index 9b46413a0d4..5615a3280ab 100644
--- a/src/passes/Heap2Local.cpp
+++ b/src/passes/Heap2Local.cpp
@@ -264,6 +264,7 @@ struct EscapeAnalyzer {
       auto seenBefore = !seen.emplace(parent).second;
       auto reachedInThisAllocation = reached.count(parent) > 0;
       if (seenBefore && !reachedInThisAllocation) {
+      // XXX struct.set can be from different allocations!
         return true;
       }
 

From 10a8b9d9d7af85b1b7615b5182b78bf2dd6f14a2 Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Mon, 24 Jun 2024 13:01:17 -0700
Subject: [PATCH 04/21] work

---
 src/passes/Heap2Local.cpp       | 26 +++++++++++++-------------
 test/lit/passes/heap2local.wast | 24 ++++++++++++++++++++----
 2 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index 5615a3280ab..8b64795a57c 100644
--- a/src/passes/Heap2Local.cpp
+++ b/src/passes/Heap2Local.cpp
@@ -248,6 +248,19 @@ struct EscapeAnalyzer {
       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 true;
+      }
+
+      // 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
@@ -268,19 +281,6 @@ struct EscapeAnalyzer {
         return true;
       }
 
-      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 true;
-      }
-
-      // The parent either fully consumes us, or flows us onwards; either way,
-      // we can proceed here, hopefully.
-      assert(interaction == ParentChildInteraction::FullyConsumes ||
-             interaction == ParentChildInteraction::Flows);
-
       // We can proceed, as the parent interacts with us properly, and we are
       // the only allocation to get here.
 
diff --git a/test/lit/passes/heap2local.wast b/test/lit/passes/heap2local.wast
index e6d5d4d88df..a037684ec92 100644
--- a/test/lit/passes/heap2local.wast
+++ b/test/lit/passes/heap2local.wast
@@ -2162,15 +2162,31 @@
 
   ;; CHECK:      (func $ref-eq-self (type $4) (result i32)
   ;; CHECK-NEXT:  (local $eq eqref)
-  ;; CHECK-NEXT:  (ref.eq
-  ;; CHECK-NEXT:   (local.tee $eq
-  ;; CHECK-NEXT:    (struct.new $struct.A
+  ;; CHECK-NEXT:  (local $1 i32)
+  ;; CHECK-NEXT:  (local $2 f64)
+  ;; CHECK-NEXT:  (local $3 i32)
+  ;; CHECK-NEXT:  (local $4 f64)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result nullref)
+  ;; CHECK-NEXT:    (local.set $3
   ;; CHECK-NEXT:     (i32.const 0)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $4
   ;; CHECK-NEXT:     (f64.const 0)
   ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $1
+  ;; CHECK-NEXT:     (local.get $3)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $2
+  ;; CHECK-NEXT:     (local.get $4)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (ref.null none)
   ;; CHECK-NEXT:   )
-  ;; CHECK-NEXT:   (local.get $eq)
   ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (ref.null none)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (i32.const 1)
   ;; CHECK-NEXT: )
   (func $ref-eq-self (result i32)
     (local $eq eqref)

From 43e5381caddd69323ac1530aecfcfcc34db43a53 Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Mon, 24 Jun 2024 13:11:12 -0700
Subject: [PATCH 05/21] Undo.opt

---
 src/passes/Heap2Local.cpp | 43 +++++----------------------------------
 1 file changed, 5 insertions(+), 38 deletions(-)

diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index 8b64795a57c..d32e0266363 100644
--- a/src/passes/Heap2Local.cpp
+++ b/src/passes/Heap2Local.cpp
@@ -170,16 +170,8 @@ namespace {
 // 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
-// optimization later (|seen| and |reached|).
+// optimization later (|reached|).
 struct EscapeAnalyzer {
-  // All the expressions that have already been seen by the optimizer, see the
-  // comment above on exclusivity: once we have seen something when analyzing
-  // one allocation, if we reach it again then we can exit early since seeing it
-  // a second time proves we lost exclusivity. We must track this across
-  // multiple instances of EscapeAnalyzer as each handles a particular
-  // allocation.
-  std::unordered_set<Expression*>& seen;
-
   // To find what escapes, we need to follow where values flow, both up to
   // parents, and via branches, and through locals.
   // TODO: for efficiency, only scan reference types in LocalGraph
@@ -190,13 +182,12 @@ struct EscapeAnalyzer {
   const PassOptions& passOptions;
   Module& wasm;
 
-  EscapeAnalyzer(std::unordered_set<Expression*>& seen,
-                 const LocalGraph& localGraph,
+  EscapeAnalyzer(const LocalGraph& localGraph,
                  const Parents& parents,
                  const BranchUtils::BranchTargets& branchTargets,
                  const PassOptions& passOptions,
                  Module& wasm)
-    : seen(seen), localGraph(localGraph), parents(parents),
+    : localGraph(localGraph), parents(parents),
       branchTargets(branchTargets), passOptions(passOptions), wasm(wasm) {}
 
   // We must track all the local.sets that write the allocation, to verify
@@ -261,26 +252,6 @@ struct EscapeAnalyzer {
       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
-      // the same thing more than once; what this check does is verify we don't
-      // 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).
-      //
-      // It is ok, however, to see a parent more than once, if the allocation
-      // flows to it from two children, as is the case for ref.eq. Likewise, for
-      // struct.set, where we also have different considerations for the two
-      // children (the reference does not escape, but the value does). We only
-      // stop here if we've seen this parent in a *different* allocation.
-      auto seenBefore = !seen.emplace(parent).second;
-      auto reachedInThisAllocation = reached.count(parent) > 0;
-      if (seenBefore && !reachedInThisAllocation) {
-      // XXX struct.set can be from different allocations!
-        return true;
-      }
-
       // We can proceed, as the parent interacts with us properly, and we are
       // the only allocation to get here.
 
@@ -1060,10 +1031,6 @@ struct Heap2Local {
     // flow to.
     localGraph.computeSetInfluences();
 
-    // All the expressions we have already looked at. We use this to avoid
-    // repeated work, see above.
-    std::unordered_set<Expression*> seen;
-
     // Find all the relevant allocations in the function: StructNew, ArrayNew,
     // ArrayNewFixed.
     struct AllocationFinder : public PostWalker<AllocationFinder> {
@@ -1123,7 +1090,7 @@ struct Heap2Local {
       }
 
       EscapeAnalyzer analyzer(
-        seen, localGraph, parents, branchTargets, passOptions, wasm);
+        localGraph, parents, branchTargets, passOptions, wasm);
       if (!analyzer.escapes(allocation)) {
         // Convert the allocation and all its uses into a struct. Then convert
         // the struct into locals.
@@ -1143,7 +1110,7 @@ struct Heap2Local {
       // Check for escaping, noting relevant information as we go. If this does
       // not escape, optimize it into locals.
       EscapeAnalyzer analyzer(
-        seen, localGraph, parents, branchTargets, passOptions, wasm);
+        localGraph, parents, branchTargets, passOptions, wasm);
       if (!analyzer.escapes(allocation)) {
         Struct2Local(allocation, analyzer, func, wasm);
       }

From d36281b919a03e87df10b92cf013220c0aed9f47 Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Mon, 24 Jun 2024 13:40:06 -0700
Subject: [PATCH 06/21] re-do

---
 src/passes/Heap2Local.cpp | 45 ++++++++++++++++++++++++++++++++++-----
 1 file changed, 40 insertions(+), 5 deletions(-)

diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index d32e0266363..ffc0baaa850 100644
--- a/src/passes/Heap2Local.cpp
+++ b/src/passes/Heap2Local.cpp
@@ -170,8 +170,16 @@ namespace {
 // 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
-// optimization later (|reached|).
+// optimization later (|seen| and |reached|).
 struct EscapeAnalyzer {
+  // All the expressions that have already been seen by the optimizer, see the
+  // comment above on exclusivity: once we have seen something when analyzing
+  // one allocation, if we reach it again then we can exit early since seeing it
+  // a second time proves we lost exclusivity. We must track this across
+  // multiple instances of EscapeAnalyzer as each handles a particular
+  // allocation.
+  std::unordered_set<Expression*>& seen;
+
   // To find what escapes, we need to follow where values flow, both up to
   // parents, and via branches, and through locals.
   // TODO: for efficiency, only scan reference types in LocalGraph
@@ -182,12 +190,13 @@ struct EscapeAnalyzer {
   const PassOptions& passOptions;
   Module& wasm;
 
-  EscapeAnalyzer(const LocalGraph& localGraph,
+  EscapeAnalyzer(std::unordered_set<Expression*>& seen,
+                 const LocalGraph& localGraph,
                  const Parents& parents,
                  const BranchUtils::BranchTargets& branchTargets,
                  const PassOptions& passOptions,
                  Module& wasm)
-    : localGraph(localGraph), parents(parents),
+    : seen(seen), localGraph(localGraph), parents(parents),
       branchTargets(branchTargets), passOptions(passOptions), wasm(wasm) {}
 
   // We must track all the local.sets that write the allocation, to verify
@@ -252,6 +261,28 @@ struct EscapeAnalyzer {
       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
+      // the same thing more than once; what this check does is verify we don't
+      // 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).
+      //
+      // It is ok, however, to see a parent more than once, if the allocation
+      // flows to it from two children, as is the case for ref.eq. Likewise, for
+      // struct.set, where we also have different considerations for the two
+      // children (the reference does not escape, but the value does), and in
+      // that case there may be different allocations for the children.
+      if (interaction == ParentChildInteraction::Flows) {
+        auto seenBefore = !seen.emplace(parent).second;
+        auto reachedInThisAllocation = reached.count(parent) > 0;
+        if (seenBefore && !reachedInThisAllocation) {
+        // XXX struct.set can be from different allocations!
+          return true;
+        }
+      }
+
       // We can proceed, as the parent interacts with us properly, and we are
       // the only allocation to get here.
 
@@ -1031,6 +1062,10 @@ struct Heap2Local {
     // flow to.
     localGraph.computeSetInfluences();
 
+    // All the expressions we have already looked at. We use this to avoid
+    // repeated work, see above.
+    std::unordered_set<Expression*> seen;
+
     // Find all the relevant allocations in the function: StructNew, ArrayNew,
     // ArrayNewFixed.
     struct AllocationFinder : public PostWalker<AllocationFinder> {
@@ -1090,7 +1125,7 @@ struct Heap2Local {
       }
 
       EscapeAnalyzer analyzer(
-        localGraph, parents, branchTargets, passOptions, wasm);
+        seen, localGraph, parents, branchTargets, passOptions, wasm);
       if (!analyzer.escapes(allocation)) {
         // Convert the allocation and all its uses into a struct. Then convert
         // the struct into locals.
@@ -1110,7 +1145,7 @@ struct Heap2Local {
       // Check for escaping, noting relevant information as we go. If this does
       // not escape, optimize it into locals.
       EscapeAnalyzer analyzer(
-        localGraph, parents, branchTargets, passOptions, wasm);
+        seen, localGraph, parents, branchTargets, passOptions, wasm);
       if (!analyzer.escapes(allocation)) {
         Struct2Local(allocation, analyzer, func, wasm);
       }

From 1c11a1e1276bd43d326555d075af8795208d3b49 Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Mon, 24 Jun 2024 13:45:30 -0700
Subject: [PATCH 07/21] Revert "re-do"

This reverts commit d36281b919a03e87df10b92cf013220c0aed9f47.
---
 src/passes/Heap2Local.cpp | 45 +++++----------------------------------
 1 file changed, 5 insertions(+), 40 deletions(-)

diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index ffc0baaa850..d32e0266363 100644
--- a/src/passes/Heap2Local.cpp
+++ b/src/passes/Heap2Local.cpp
@@ -170,16 +170,8 @@ namespace {
 // 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
-// optimization later (|seen| and |reached|).
+// optimization later (|reached|).
 struct EscapeAnalyzer {
-  // All the expressions that have already been seen by the optimizer, see the
-  // comment above on exclusivity: once we have seen something when analyzing
-  // one allocation, if we reach it again then we can exit early since seeing it
-  // a second time proves we lost exclusivity. We must track this across
-  // multiple instances of EscapeAnalyzer as each handles a particular
-  // allocation.
-  std::unordered_set<Expression*>& seen;
-
   // To find what escapes, we need to follow where values flow, both up to
   // parents, and via branches, and through locals.
   // TODO: for efficiency, only scan reference types in LocalGraph
@@ -190,13 +182,12 @@ struct EscapeAnalyzer {
   const PassOptions& passOptions;
   Module& wasm;
 
-  EscapeAnalyzer(std::unordered_set<Expression*>& seen,
-                 const LocalGraph& localGraph,
+  EscapeAnalyzer(const LocalGraph& localGraph,
                  const Parents& parents,
                  const BranchUtils::BranchTargets& branchTargets,
                  const PassOptions& passOptions,
                  Module& wasm)
-    : seen(seen), localGraph(localGraph), parents(parents),
+    : localGraph(localGraph), parents(parents),
       branchTargets(branchTargets), passOptions(passOptions), wasm(wasm) {}
 
   // We must track all the local.sets that write the allocation, to verify
@@ -261,28 +252,6 @@ struct EscapeAnalyzer {
       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
-      // the same thing more than once; what this check does is verify we don't
-      // 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).
-      //
-      // It is ok, however, to see a parent more than once, if the allocation
-      // flows to it from two children, as is the case for ref.eq. Likewise, for
-      // struct.set, where we also have different considerations for the two
-      // children (the reference does not escape, but the value does), and in
-      // that case there may be different allocations for the children.
-      if (interaction == ParentChildInteraction::Flows) {
-        auto seenBefore = !seen.emplace(parent).second;
-        auto reachedInThisAllocation = reached.count(parent) > 0;
-        if (seenBefore && !reachedInThisAllocation) {
-        // XXX struct.set can be from different allocations!
-          return true;
-        }
-      }
-
       // We can proceed, as the parent interacts with us properly, and we are
       // the only allocation to get here.
 
@@ -1062,10 +1031,6 @@ struct Heap2Local {
     // flow to.
     localGraph.computeSetInfluences();
 
-    // All the expressions we have already looked at. We use this to avoid
-    // repeated work, see above.
-    std::unordered_set<Expression*> seen;
-
     // Find all the relevant allocations in the function: StructNew, ArrayNew,
     // ArrayNewFixed.
     struct AllocationFinder : public PostWalker<AllocationFinder> {
@@ -1125,7 +1090,7 @@ struct Heap2Local {
       }
 
       EscapeAnalyzer analyzer(
-        seen, localGraph, parents, branchTargets, passOptions, wasm);
+        localGraph, parents, branchTargets, passOptions, wasm);
       if (!analyzer.escapes(allocation)) {
         // Convert the allocation and all its uses into a struct. Then convert
         // the struct into locals.
@@ -1145,7 +1110,7 @@ struct Heap2Local {
       // Check for escaping, noting relevant information as we go. If this does
       // not escape, optimize it into locals.
       EscapeAnalyzer analyzer(
-        seen, localGraph, parents, branchTargets, passOptions, wasm);
+        localGraph, parents, branchTargets, passOptions, wasm);
       if (!analyzer.escapes(allocation)) {
         Struct2Local(allocation, analyzer, func, wasm);
       }

From 19823a64bf53f5bd5c5bd8b16bf32c01cff9ed40 Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Mon, 24 Jun 2024 13:45:53 -0700
Subject: [PATCH 08/21] format

---
 src/passes/Heap2Local.cpp | 12 +++++-------
 1 file changed, 5 insertions(+), 7 deletions(-)

diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index d32e0266363..34fd9f54d83 100644
--- a/src/passes/Heap2Local.cpp
+++ b/src/passes/Heap2Local.cpp
@@ -187,8 +187,8 @@ struct EscapeAnalyzer {
                  const BranchUtils::BranchTargets& branchTargets,
                  const PassOptions& passOptions,
                  Module& wasm)
-    : localGraph(localGraph), parents(parents),
-      branchTargets(branchTargets), passOptions(passOptions), wasm(wasm) {}
+    : localGraph(localGraph), parents(parents), branchTargets(branchTargets),
+      passOptions(passOptions), wasm(wasm) {}
 
   // We must track all the local.sets that write the allocation, to verify
   // exclusivity.
@@ -717,11 +717,9 @@ struct Struct2Local : PostWalker<Struct2Local> {
     // not escape to any other place.
     int32_t result = analyzer.reached.count(curr->left) > 0 &&
                      analyzer.reached.count(curr->right) > 0;
-    auto* block = builder.makeBlock({
-      builder.makeDrop(curr->left),
-      builder.makeDrop(curr->right),
-      builder.makeConst(Literal(result))
-    });
+    auto* block = builder.makeBlock({builder.makeDrop(curr->left),
+                                     builder.makeDrop(curr->right),
+                                     builder.makeConst(Literal(result))});
     replaceCurrent(block);
   }
 

From 078b81313dee1f4e9960a6a039c74cd073cbee17 Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Mon, 24 Jun 2024 15:29:14 -0700
Subject: [PATCH 09/21] almost

---
 src/passes/Heap2Local.cpp       |  11 ++-
 test/lit/passes/heap2local.wast | 169 +++++++++++++++++++++++---------
 2 files changed, 129 insertions(+), 51 deletions(-)

diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index 34fd9f54d83..5c36c0ed951 100644
--- a/src/passes/Heap2Local.cpp
+++ b/src/passes/Heap2Local.cpp
@@ -714,13 +714,14 @@ 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.
+    // not escape to any other place. Note that we need no special handling for
+    // unreachability here: if one arm is unreachable, then the result does not
+    // matter.
     int32_t result = analyzer.reached.count(curr->left) > 0 &&
                      analyzer.reached.count(curr->right) > 0;
-    auto* block = builder.makeBlock({builder.makeDrop(curr->left),
-                                     builder.makeDrop(curr->right),
-                                     builder.makeConst(Literal(result))});
-    replaceCurrent(block);
+    // For simplicity, simply drop the RefEq and put a constant result after.
+    replaceCurrent(builder.makeSequence(builder.makeDrop(curr),
+                                        builder.makeConst(Literal(result))));
   }
 
   void visitRefAs(RefAs* curr) {
diff --git a/test/lit/passes/heap2local.wast b/test/lit/passes/heap2local.wast
index a037684ec92..7ffd90eb677 100644
--- a/test/lit/passes/heap2local.wast
+++ b/test/lit/passes/heap2local.wast
@@ -2089,25 +2089,25 @@
   ;; CHECK-NEXT:  (local $2 i32)
   ;; CHECK-NEXT:  (local $3 f64)
   ;; CHECK-NEXT:  (drop
-  ;; CHECK-NEXT:   (block (result nullref)
-  ;; CHECK-NEXT:    (local.set $2
-  ;; CHECK-NEXT:     (i32.const 0)
-  ;; CHECK-NEXT:    )
-  ;; CHECK-NEXT:    (local.set $3
-  ;; CHECK-NEXT:     (f64.const 0)
-  ;; CHECK-NEXT:    )
-  ;; CHECK-NEXT:    (local.set $0
-  ;; CHECK-NEXT:     (local.get $2)
-  ;; CHECK-NEXT:    )
-  ;; CHECK-NEXT:    (local.set $1
-  ;; CHECK-NEXT:     (local.get $3)
+  ;; CHECK-NEXT:   (ref.eq
+  ;; CHECK-NEXT:    (block (result nullref)
+  ;; CHECK-NEXT:     (local.set $2
+  ;; CHECK-NEXT:      (i32.const 0)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (local.set $3
+  ;; CHECK-NEXT:      (f64.const 0)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (local.set $0
+  ;; CHECK-NEXT:      (local.get $2)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (local.set $1
+  ;; CHECK-NEXT:      (local.get $3)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (ref.null none)
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:    (ref.null none)
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
-  ;; CHECK-NEXT:  (drop
-  ;; CHECK-NEXT:   (ref.null none)
-  ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (i32.const 0)
   ;; CHECK-NEXT: )
   (func $ref-eq (result i32)
@@ -2128,23 +2128,23 @@
   ;; CHECK-NEXT:  (local $3 i32)
   ;; CHECK-NEXT:  (local $4 f64)
   ;; CHECK-NEXT:  (drop
-  ;; CHECK-NEXT:   (ref.null none)
-  ;; CHECK-NEXT:  )
-  ;; CHECK-NEXT:  (drop
-  ;; CHECK-NEXT:   (block (result nullref)
-  ;; CHECK-NEXT:    (local.set $3
-  ;; CHECK-NEXT:     (i32.const 0)
-  ;; CHECK-NEXT:    )
-  ;; CHECK-NEXT:    (local.set $4
-  ;; CHECK-NEXT:     (f64.const 0)
-  ;; CHECK-NEXT:    )
-  ;; CHECK-NEXT:    (local.set $1
-  ;; CHECK-NEXT:     (local.get $3)
-  ;; CHECK-NEXT:    )
-  ;; CHECK-NEXT:    (local.set $2
-  ;; CHECK-NEXT:     (local.get $4)
-  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   (ref.eq
   ;; CHECK-NEXT:    (ref.null none)
+  ;; CHECK-NEXT:    (block (result nullref)
+  ;; CHECK-NEXT:     (local.set $3
+  ;; CHECK-NEXT:      (i32.const 0)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (local.set $4
+  ;; CHECK-NEXT:      (f64.const 0)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (local.set $1
+  ;; CHECK-NEXT:      (local.get $3)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (local.set $2
+  ;; CHECK-NEXT:      (local.get $4)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (ref.null none)
+  ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (i32.const 0)
@@ -2167,25 +2167,25 @@
   ;; CHECK-NEXT:  (local $3 i32)
   ;; CHECK-NEXT:  (local $4 f64)
   ;; CHECK-NEXT:  (drop
-  ;; CHECK-NEXT:   (block (result nullref)
-  ;; CHECK-NEXT:    (local.set $3
-  ;; CHECK-NEXT:     (i32.const 0)
-  ;; CHECK-NEXT:    )
-  ;; CHECK-NEXT:    (local.set $4
-  ;; CHECK-NEXT:     (f64.const 0)
-  ;; CHECK-NEXT:    )
-  ;; CHECK-NEXT:    (local.set $1
-  ;; CHECK-NEXT:     (local.get $3)
-  ;; CHECK-NEXT:    )
-  ;; CHECK-NEXT:    (local.set $2
-  ;; CHECK-NEXT:     (local.get $4)
+  ;; CHECK-NEXT:   (ref.eq
+  ;; CHECK-NEXT:    (block (result nullref)
+  ;; CHECK-NEXT:     (local.set $3
+  ;; CHECK-NEXT:      (i32.const 0)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (local.set $4
+  ;; CHECK-NEXT:      (f64.const 0)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (local.set $1
+  ;; CHECK-NEXT:      (local.get $3)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (local.set $2
+  ;; CHECK-NEXT:      (local.get $4)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (ref.null none)
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:    (ref.null none)
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
-  ;; CHECK-NEXT:  (drop
-  ;; CHECK-NEXT:   (ref.null none)
-  ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (i32.const 1)
   ;; CHECK-NEXT: )
   (func $ref-eq-self (result i32)
@@ -2202,6 +2202,83 @@
       (local.get $eq)
     )
   )
+
+  ;; CHECK:      (func $ref-eq-unreachable (type $4) (result i32)
+  ;; CHECK-NEXT:  (local $0 i32)
+  ;; CHECK-NEXT:  (local $1 f64)
+  ;; CHECK-NEXT:  (local $2 i32)
+  ;; CHECK-NEXT:  (local $3 f64)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (ref.eq
+  ;; CHECK-NEXT:    (block (result nullref)
+  ;; CHECK-NEXT:     (local.set $2
+  ;; CHECK-NEXT:      (i32.const 0)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (local.set $3
+  ;; CHECK-NEXT:      (f64.const 0)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (local.set $0
+  ;; CHECK-NEXT:      (local.get $2)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (local.set $1
+  ;; CHECK-NEXT:      (local.get $3)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (ref.null none)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (unreachable)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (i32.const 0)
+  ;; CHECK-NEXT: )
+  (func $ref-eq-unreachable (result i32)
+    ;; When a child is unreachable, the result does not matter (but we should
+    ;; still emit validating code).
+    (ref.eq
+      (struct.new $struct.A
+        (i32.const 0)
+        (f64.const 0)
+      )
+      (unreachable)
+    )
+  )
+
+  ;; CHECK:      (func $ref-eq-unreachable-flipped (type $4) (result i32)
+  ;; CHECK-NEXT:  (local $0 i32)
+  ;; CHECK-NEXT:  (local $1 f64)
+  ;; CHECK-NEXT:  (local $2 i32)
+  ;; CHECK-NEXT:  (local $3 f64)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (ref.eq
+  ;; CHECK-NEXT:    (unreachable)
+  ;; CHECK-NEXT:    (block (result nullref)
+  ;; CHECK-NEXT:     (local.set $2
+  ;; CHECK-NEXT:      (i32.const 0)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (local.set $3
+  ;; CHECK-NEXT:      (f64.const 0)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (local.set $0
+  ;; CHECK-NEXT:      (local.get $2)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (local.set $1
+  ;; CHECK-NEXT:      (local.get $3)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (ref.null none)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (i32.const 0)
+  ;; CHECK-NEXT: )
+  (func $ref-eq-unreachable-flipped (result i32)
+    ;; As above, but with children flipped.
+    (ref.eq
+      (unreachable)
+      (struct.new $struct.A
+        (i32.const 0)
+        (f64.const 0)
+      )
+    )
+  )
 )
 
 (module

From 090af2d67fcb592980f71f6f50307cf2d8b1ca05 Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Mon, 24 Jun 2024 15:42:42 -0700
Subject: [PATCH 10/21] work

---
 src/passes/Heap2Local.cpp       | 10 +++--
 test/lit/passes/heap2local.wast | 66 +++++++++++++++------------------
 2 files changed, 37 insertions(+), 39 deletions(-)

diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index 5c36c0ed951..39c77bec489 100644
--- a/src/passes/Heap2Local.cpp
+++ b/src/passes/Heap2Local.cpp
@@ -712,11 +712,15 @@ struct Struct2Local : PostWalker<Struct2Local> {
       return;
     }
 
+    if (curr->type == Type::unreachable) {
+      // The result does not matter. Leave things as they are (and let DCE
+      // handle it).
+      return;
+    }
+
     // 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. Note that we need no special handling for
-    // unreachability here: if one arm is unreachable, then the result does not
-    // matter.
+    // not escape to any other place.
     int32_t result = analyzer.reached.count(curr->left) > 0 &&
                      analyzer.reached.count(curr->right) > 0;
     // For simplicity, simply drop the RefEq and put a constant result after.
diff --git a/test/lit/passes/heap2local.wast b/test/lit/passes/heap2local.wast
index 7ffd90eb677..816cb13533c 100644
--- a/test/lit/passes/heap2local.wast
+++ b/test/lit/passes/heap2local.wast
@@ -2208,27 +2208,24 @@
   ;; CHECK-NEXT:  (local $1 f64)
   ;; CHECK-NEXT:  (local $2 i32)
   ;; CHECK-NEXT:  (local $3 f64)
-  ;; CHECK-NEXT:  (drop
-  ;; CHECK-NEXT:   (ref.eq
-  ;; CHECK-NEXT:    (block (result nullref)
-  ;; CHECK-NEXT:     (local.set $2
-  ;; CHECK-NEXT:      (i32.const 0)
-  ;; CHECK-NEXT:     )
-  ;; CHECK-NEXT:     (local.set $3
-  ;; CHECK-NEXT:      (f64.const 0)
-  ;; CHECK-NEXT:     )
-  ;; CHECK-NEXT:     (local.set $0
-  ;; CHECK-NEXT:      (local.get $2)
-  ;; CHECK-NEXT:     )
-  ;; CHECK-NEXT:     (local.set $1
-  ;; CHECK-NEXT:      (local.get $3)
-  ;; CHECK-NEXT:     )
-  ;; CHECK-NEXT:     (ref.null none)
+  ;; CHECK-NEXT:  (ref.eq
+  ;; CHECK-NEXT:   (block (result nullref)
+  ;; CHECK-NEXT:    (local.set $2
+  ;; CHECK-NEXT:     (i32.const 0)
   ;; CHECK-NEXT:    )
-  ;; CHECK-NEXT:    (unreachable)
+  ;; CHECK-NEXT:    (local.set $3
+  ;; CHECK-NEXT:     (f64.const 0)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $0
+  ;; CHECK-NEXT:     (local.get $2)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $1
+  ;; CHECK-NEXT:     (local.get $3)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (ref.null none)
   ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (unreachable)
   ;; CHECK-NEXT:  )
-  ;; CHECK-NEXT:  (i32.const 0)
   ;; CHECK-NEXT: )
   (func $ref-eq-unreachable (result i32)
     ;; When a child is unreachable, the result does not matter (but we should
@@ -2247,27 +2244,24 @@
   ;; CHECK-NEXT:  (local $1 f64)
   ;; CHECK-NEXT:  (local $2 i32)
   ;; CHECK-NEXT:  (local $3 f64)
-  ;; CHECK-NEXT:  (drop
-  ;; CHECK-NEXT:   (ref.eq
-  ;; CHECK-NEXT:    (unreachable)
-  ;; CHECK-NEXT:    (block (result nullref)
-  ;; CHECK-NEXT:     (local.set $2
-  ;; CHECK-NEXT:      (i32.const 0)
-  ;; CHECK-NEXT:     )
-  ;; CHECK-NEXT:     (local.set $3
-  ;; CHECK-NEXT:      (f64.const 0)
-  ;; CHECK-NEXT:     )
-  ;; CHECK-NEXT:     (local.set $0
-  ;; CHECK-NEXT:      (local.get $2)
-  ;; CHECK-NEXT:     )
-  ;; CHECK-NEXT:     (local.set $1
-  ;; CHECK-NEXT:      (local.get $3)
-  ;; CHECK-NEXT:     )
-  ;; CHECK-NEXT:     (ref.null none)
+  ;; CHECK-NEXT:  (ref.eq
+  ;; CHECK-NEXT:   (unreachable)
+  ;; CHECK-NEXT:   (block (result nullref)
+  ;; CHECK-NEXT:    (local.set $2
+  ;; CHECK-NEXT:     (i32.const 0)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $3
+  ;; CHECK-NEXT:     (f64.const 0)
   ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $0
+  ;; CHECK-NEXT:     (local.get $2)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $1
+  ;; CHECK-NEXT:     (local.get $3)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (ref.null none)
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
-  ;; CHECK-NEXT:  (i32.const 0)
   ;; CHECK-NEXT: )
   (func $ref-eq-unreachable-flipped (result i32)
     ;; As above, but with children flipped.

From 60050be9abc2660a2a9cf97f9b83ccdf642b3f79 Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Tue, 25 Jun 2024 09:08:41 -0700
Subject: [PATCH 11/21] fix

---
 test/lit/passes/heap2local.wast | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/test/lit/passes/heap2local.wast b/test/lit/passes/heap2local.wast
index 816cb13533c..3a922af322c 100644
--- a/test/lit/passes/heap2local.wast
+++ b/test/lit/passes/heap2local.wast
@@ -2129,7 +2129,7 @@
   ;; CHECK-NEXT:  (local $4 f64)
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (ref.eq
-  ;; CHECK-NEXT:    (ref.null none)
+  ;; CHECK-NEXT:    (local.get $other)
   ;; CHECK-NEXT:    (block (result nullref)
   ;; CHECK-NEXT:     (local.set $3
   ;; CHECK-NEXT:      (i32.const 0)
@@ -2152,7 +2152,7 @@
   (func $ref-eq-flip (param $other eqref) (result i32)
     ;; As above, but flipped, and compared to a local.
     (ref.eq
-      (ref.null eq)
+      (local.get $other)
       (struct.new $struct.A
         (i32.const 0)
         (f64.const 0)

From dea6db54ddda7f1d8dc9d2360caacd829baf6397 Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Tue, 25 Jun 2024 10:55:52 -0700
Subject: [PATCH 12/21] add two more tests

---
 test/lit/passes/heap2local.wast | 80 +++++++++++++++++++++++++++++++++
 1 file changed, 80 insertions(+)

diff --git a/test/lit/passes/heap2local.wast b/test/lit/passes/heap2local.wast
index 3a922af322c..f1f629ffc3f 100644
--- a/test/lit/passes/heap2local.wast
+++ b/test/lit/passes/heap2local.wast
@@ -38,6 +38,8 @@
 
   ;; CHECK:      (type $12 (func (param eqref) (result i32)))
 
+  ;; CHECK:      (type $13 (func (param eqref eqref) (result i32)))
+
   ;; CHECK:      (func $simple (type $1)
   ;; CHECK-NEXT:  (local $0 i32)
   ;; CHECK-NEXT:  (local $1 f64)
@@ -2273,6 +2275,84 @@
       )
     )
   )
+
+  ;; CHECK:      (func $ref-eq-unrelated (type $13) (param $x eqref) (param $y eqref) (result i32)
+  ;; CHECK-NEXT:  (ref.eq
+  ;; CHECK-NEXT:   (local.get $x)
+  ;; CHECK-NEXT:   (local.get $y)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $ref-eq-unrelated (param $x eqref) (param $y eqref) (result i32)
+    ;; We know nothing about either arm, and do nothing.
+    (ref.eq
+      (local.get $x)
+      (local.get $y)
+    )
+  )
+
+  ;; CHECK:      (func $ref-eq-two (type $4) (result i32)
+  ;; CHECK-NEXT:  (local $0 i32)
+  ;; CHECK-NEXT:  (local $1 f64)
+  ;; CHECK-NEXT:  (local $2 i32)
+  ;; CHECK-NEXT:  (local $3 f64)
+  ;; CHECK-NEXT:  (local $4 i32)
+  ;; CHECK-NEXT:  (local $5 f64)
+  ;; CHECK-NEXT:  (local $6 i32)
+  ;; CHECK-NEXT:  (local $7 f64)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result i32)
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (ref.eq
+  ;; CHECK-NEXT:      (block (result nullref)
+  ;; CHECK-NEXT:       (local.set $2
+  ;; CHECK-NEXT:        (i32.const 0)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (local.set $3
+  ;; CHECK-NEXT:        (f64.const 0)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (local.set $0
+  ;; CHECK-NEXT:        (local.get $2)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (local.set $1
+  ;; CHECK-NEXT:        (local.get $3)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (ref.null none)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (block (result nullref)
+  ;; CHECK-NEXT:       (local.set $6
+  ;; CHECK-NEXT:        (i32.const 1)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (local.set $7
+  ;; CHECK-NEXT:        (f64.const 2.2)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (local.set $4
+  ;; CHECK-NEXT:        (local.get $6)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (local.set $5
+  ;; CHECK-NEXT:        (local.get $7)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (ref.null none)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i32.const 0)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (i32.const 0)
+  ;; CHECK-NEXT: )
+  (func $ref-eq-two (result i32)
+    ;; Two separate allocations. We can optimize them away, and the result is 0.
+    (ref.eq
+      (struct.new $struct.A
+        (i32.const 0)
+        (f64.const 0)
+      )
+      (struct.new $struct.A
+        (i32.const 1)
+        (f64.const 2.2)
+      )
+    )
+  )
 )
 
 (module

From 9029ba2a7469e69171fd42058dd693f21cd7069d Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Tue, 25 Jun 2024 11:35:19 -0700
Subject: [PATCH 13/21] improve test

---
 test/lit/passes/heap2local.wast | 31 ++++++++++++++++++++++++++++++-
 1 file changed, 30 insertions(+), 1 deletion(-)

diff --git a/test/lit/passes/heap2local.wast b/test/lit/passes/heap2local.wast
index f1f629ffc3f..bd613e18c4b 100644
--- a/test/lit/passes/heap2local.wast
+++ b/test/lit/passes/heap2local.wast
@@ -2277,13 +2277,42 @@
   )
 
   ;; CHECK:      (func $ref-eq-unrelated (type $13) (param $x eqref) (param $y eqref) (result i32)
+  ;; CHECK-NEXT:  (local $2 i32)
+  ;; CHECK-NEXT:  (local $3 f64)
+  ;; CHECK-NEXT:  (local $4 i32)
+  ;; CHECK-NEXT:  (local $5 f64)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result nullref)
+  ;; CHECK-NEXT:    (local.set $4
+  ;; CHECK-NEXT:     (i32.const 0)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $5
+  ;; CHECK-NEXT:     (f64.const 0)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $2
+  ;; CHECK-NEXT:     (local.get $4)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $3
+  ;; CHECK-NEXT:     (local.get $5)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (ref.null none)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (ref.eq
   ;; CHECK-NEXT:   (local.get $x)
   ;; CHECK-NEXT:   (local.get $y)
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT: )
   (func $ref-eq-unrelated (param $x eqref) (param $y eqref) (result i32)
-    ;; We know nothing about either arm, and do nothing.
+    ;; We know nothing about either ref.eq arm, and do nothing, despite
+    ;; another allocation in the function (which ensures we enter the
+    ;; optimization part of the pass; that other allocation can be removed).
+    (drop
+      (struct.new $struct.A
+        (i32.const 0)
+        (f64.const 0)
+      )
+    )
     (ref.eq
       (local.get $x)
       (local.get $y)

From f305e4f635b803b947494a259c40d833dc5a991f Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Tue, 25 Jun 2024 11:41:23 -0700
Subject: [PATCH 14/21] more

---
 src/passes/Heap2Local.cpp | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index 39c77bec489..6093c8d8a2e 100644
--- a/src/passes/Heap2Local.cpp
+++ b/src/passes/Heap2Local.cpp
@@ -349,6 +349,12 @@ struct EscapeAnalyzer {
       void visitLocalSet(LocalSet* curr) { escapes = false; }
 
       // Reference operations. TODO add more
+      void visitRefIsNull(RefIsNull* curr) {
+        // The reference is compared to null, but nothing more.
+        escapes = false;
+        fullyConsumes = true;
+      }
+
       void visitRefEq(RefEq* curr) {
         // The reference is compared for identity, but nothing more.
         escapes = false;
@@ -707,6 +713,17 @@ struct Struct2Local : PostWalker<Struct2Local> {
     replaceCurrent(builder.makeBlock(contents));
   }
 
+  void visitRefIsNull(RefIsNull* curr) {
+    if (!analyzer.reached.count(curr)) {
+      return;
+    }
+
+    // The result must be 0, since the allocation is not null. Drop the RefIs
+    // and append that.
+    replaceCurrent(builder.makeSequence(builder.makeDrop(curr),
+                                        builder.makeConst(Literal(int32_t(0)))));
+  }
+
   void visitRefEq(RefEq* curr) {
     if (!analyzer.reached.count(curr)) {
       return;

From 8d26f56312c134a3fcbb8243dd840ad1a13b5f6a Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Tue, 25 Jun 2024 11:44:03 -0700
Subject: [PATCH 15/21] test

---
 test/lit/passes/heap2local.wast | 53 +++++++++++++++++++++++++++++++++
 1 file changed, 53 insertions(+)

diff --git a/test/lit/passes/heap2local.wast b/test/lit/passes/heap2local.wast
index bd613e18c4b..d91a895431f 100644
--- a/test/lit/passes/heap2local.wast
+++ b/test/lit/passes/heap2local.wast
@@ -40,6 +40,8 @@
 
   ;; CHECK:      (type $13 (func (param eqref eqref) (result i32)))
 
+  ;; CHECK:      (type $14 (func (param anyref) (result i32)))
+
   ;; CHECK:      (func $simple (type $1)
   ;; CHECK-NEXT:  (local $0 i32)
   ;; CHECK-NEXT:  (local $1 f64)
@@ -2382,6 +2384,57 @@
       )
     )
   )
+
+  ;; CHECK:      (func $ref-is (type $14) (param $x anyref) (result i32)
+  ;; CHECK-NEXT:  (local $1 i32)
+  ;; CHECK-NEXT:  (local $2 f64)
+  ;; CHECK-NEXT:  (local $3 i32)
+  ;; CHECK-NEXT:  (local $4 f64)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (ref.is_null
+  ;; CHECK-NEXT:    (local.get $x)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (block (result i32)
+  ;; CHECK-NEXT:   (drop
+  ;; CHECK-NEXT:    (ref.is_null
+  ;; CHECK-NEXT:     (block (result nullref)
+  ;; CHECK-NEXT:      (local.set $3
+  ;; CHECK-NEXT:       (i32.const 0)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (local.set $4
+  ;; CHECK-NEXT:       (f64.const 0)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (local.set $1
+  ;; CHECK-NEXT:       (local.get $3)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (local.set $2
+  ;; CHECK-NEXT:       (local.get $4)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (ref.null none)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (i32.const 0)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $ref-is (param $x anyref) (result i32)
+    ;; A ref.is that we can do nothing for, and should not modify, even though
+    ;; we optimize later.
+    (drop
+      (ref.is_null
+        (local.get $x)
+      )
+    )
+    ;; The result is 0 as the allocation is not null, and we can remove the
+    ;; allocation.
+    (ref.is_null
+      (struct.new $struct.A
+        (i32.const 0)
+        (f64.const 0)
+      )
+    )
+  )
 )
 
 (module

From cffde1529b70c3945413ce530b8715ef0339f581 Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Tue, 25 Jun 2024 11:46:14 -0700
Subject: [PATCH 16/21] todo

---
 src/passes/Heap2Local.cpp | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index 6093c8d8a2e..4e580c2b9e6 100644
--- a/src/passes/Heap2Local.cpp
+++ b/src/passes/Heap2Local.cpp
@@ -382,6 +382,8 @@ struct EscapeAnalyzer {
         }
       }
 
+      // TODO: RefTest
+
       // GC operations.
       void visitStructSet(StructSet* curr) {
         // The reference does not escape (but the value is stored to memory and

From 66f8b13134918fafbea77612d33fc6575b64ec7b Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Tue, 25 Jun 2024 11:57:57 -0700
Subject: [PATCH 17/21] work

---
 src/passes/Heap2Local.cpp       | 22 +++++++++++--
 test/lit/passes/heap2local.wast | 57 +++++++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+), 2 deletions(-)

diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index 4e580c2b9e6..e786f37d605 100644
--- a/src/passes/Heap2Local.cpp
+++ b/src/passes/Heap2Local.cpp
@@ -373,6 +373,11 @@ struct EscapeAnalyzer {
         }
       }
 
+      void visitRefTest(RefTest* curr) {
+        escapes = false;
+        fullyConsumes = true;
+      }
+
       void visitRefCast(RefCast* curr) {
         // As it is our allocation that flows through here, we need to
         // check that the cast will not trap, so that we can continue
@@ -382,8 +387,6 @@ struct EscapeAnalyzer {
         }
       }
 
-      // TODO: RefTest
-
       // GC operations.
       void visitStructSet(StructSet* curr) {
         // The reference does not escape (but the value is stored to memory and
@@ -758,6 +761,21 @@ struct Struct2Local : PostWalker<Struct2Local> {
     replaceCurrent(curr->value);
   }
 
+  void visitRefTest(RefTest* curr) {
+    if (!analyzer.reached.count(curr)) {
+      return;
+    }
+
+    // This test operates on the allocation, which means we can compute whether
+    // it will succeed statically. We do not even need
+    // GCTypeUtils::evaluateCastCheck because we know the allocation's type
+    // precisely (it cannot be a strict subtype of the type - it is the type).
+    int32_t result = Type::isSubType(allocation->type, curr->castType);
+    // For simplicity, simply drop the RefEq and put a constant result after.
+    replaceCurrent(builder.makeSequence(builder.makeDrop(curr),
+                                        builder.makeConst(Literal(result))));
+  }
+
   void visitRefCast(RefCast* curr) {
     if (!analyzer.reached.count(curr)) {
       return;
diff --git a/test/lit/passes/heap2local.wast b/test/lit/passes/heap2local.wast
index d91a895431f..c58ea14270d 100644
--- a/test/lit/passes/heap2local.wast
+++ b/test/lit/passes/heap2local.wast
@@ -7,6 +7,8 @@
   ;; CHECK:      (type $struct.A (struct (field (mut i32)) (field (mut f64))))
   (type $struct.A (struct (field (mut i32)) (field (mut f64))))
 
+  (type $struct.B (sub $struct.A (struct (field (mut i32)) (field (mut f64)))))
+
   ;; CHECK:      (type $1 (func))
 
   ;; CHECK:      (type $2 (func (result f64)))
@@ -2435,6 +2437,61 @@
       )
     )
   )
+
+  (func $ref-test (param $x anyref)
+    ;; This cast must succeed (it tests the exact type), and we can remove the
+    ;; allocation.
+    (drop
+      (ref.test (ref $struct.A)
+        (struct.new $struct.A
+          (i32.const 0)
+          (f64.const 0)
+        )
+      )
+    )
+    ;; Testing a supertype also works.
+    (drop
+      (ref.test (ref null $struct.A)
+        (struct.new $struct.A
+          (i32.const 1)
+          (f64.const 2.2)
+        )
+      )
+    )
+    (drop
+      (ref.test (ref null any)
+        (struct.new $struct.A
+          (i32.const 3)
+          (f64.const 4.4)
+        )
+      )
+    )
+  )
+  (func $ref-test-bad
+    ;; We know nothing about this cast.
+    (drop
+      (ref.test (ref $struct.A)
+        (local.get $x)
+      )
+    )
+    ;; These casts must fail, but we can remove the allocation.
+    (drop
+      (ref.test (ref $struct.B)
+        (struct.new $struct.A
+          (i32.const 0)
+          (f64.const 0)
+        )
+      )
+    )
+    (drop
+      (ref.test (ref null $struct.B)
+        (struct.new $struct.A
+          (i32.const 1)
+          (f64.const 2.2)
+        )
+      )
+    )
+  )
 )
 
 (module

From 72c9b605bbbfaafc854c8a574f259889da3c2442 Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Tue, 25 Jun 2024 12:00:27 -0700
Subject: [PATCH 18/21] test

---
 test/lit/passes/heap2local.wast | 188 +++++++++++++++++++++++++++++---
 1 file changed, 170 insertions(+), 18 deletions(-)

diff --git a/test/lit/passes/heap2local.wast b/test/lit/passes/heap2local.wast
index c58ea14270d..ed2af717d98 100644
--- a/test/lit/passes/heap2local.wast
+++ b/test/lit/passes/heap2local.wast
@@ -4,10 +4,8 @@
 ;; RUN: foreach %s %t wasm-opt -all --remove-unused-names --heap2local -S -o - | filecheck %s
 
 (module
-  ;; CHECK:      (type $struct.A (struct (field (mut i32)) (field (mut f64))))
-  (type $struct.A (struct (field (mut i32)) (field (mut f64))))
-
-  (type $struct.B (sub $struct.A (struct (field (mut i32)) (field (mut f64)))))
+  ;; CHECK:      (type $struct.A (sub (struct (field (mut i32)) (field (mut f64)))))
+  (type $struct.A (sub (struct (field (mut i32)) (field (mut f64)))))
 
   ;; CHECK:      (type $1 (func))
 
@@ -23,6 +21,9 @@
 
   ;; CHECK:      (type $7 (func (param i32) (result f64)))
 
+  ;; CHECK:      (type $struct.B (sub $struct.A (struct (field (mut i32)) (field (mut f64)))))
+  (type $struct.B (sub $struct.A (struct (field (mut i32)) (field (mut f64)))))
+
   ;; CHECK:      (type $struct.packed (struct (field (mut i8)) (field (mut i32))))
   (type $struct.packed (struct (field (mut i8)) (field (mut i32))))
 
@@ -32,17 +33,19 @@
 
   (type $struct.nonnullable (struct (field (ref $struct.A))))
 
-  ;; CHECK:      (type $9 (func (param (ref null $struct.recursive))))
+  ;; CHECK:      (type $10 (func (param (ref null $struct.recursive))))
+
+  ;; CHECK:      (type $11 (func (param (ref $struct.A))))
 
-  ;; CHECK:      (type $10 (func (param (ref $struct.A))))
+  ;; CHECK:      (type $12 (func (param i32)))
 
-  ;; CHECK:      (type $11 (func (param i32)))
+  ;; CHECK:      (type $13 (func (param eqref) (result i32)))
 
-  ;; CHECK:      (type $12 (func (param eqref) (result i32)))
+  ;; CHECK:      (type $14 (func (param eqref eqref) (result i32)))
 
-  ;; CHECK:      (type $13 (func (param eqref eqref) (result i32)))
+  ;; CHECK:      (type $15 (func (param anyref) (result i32)))
 
-  ;; CHECK:      (type $14 (func (param anyref) (result i32)))
+  ;; CHECK:      (type $16 (func (param anyref)))
 
   ;; CHECK:      (func $simple (type $1)
   ;; CHECK-NEXT:  (local $0 i32)
@@ -948,7 +951,7 @@
     )
   )
 
-  ;; CHECK:      (func $set-value (type $9) (param $struct.recursive (ref null $struct.recursive))
+  ;; CHECK:      (func $set-value (type $10) (param $struct.recursive (ref null $struct.recursive))
   ;; CHECK-NEXT:  (local $ref (ref null $struct.recursive))
   ;; CHECK-NEXT:  (struct.set $struct.recursive 0
   ;; CHECK-NEXT:   (local.get $struct.recursive)
@@ -1063,7 +1066,7 @@
     )
   )
 
-  ;; CHECK:      (func $non-nullable (type $10) (param $a (ref $struct.A))
+  ;; CHECK:      (func $non-nullable (type $11) (param $a (ref $struct.A))
   ;; CHECK-NEXT:  (local $1 (ref $struct.A))
   ;; CHECK-NEXT:  (local $2 (ref $struct.A))
   ;; CHECK-NEXT:  (drop
@@ -1095,7 +1098,7 @@
     )
   )
 
-  ;; CHECK:      (func $before-loop-use-multi (type $11) (param $x i32)
+  ;; CHECK:      (func $before-loop-use-multi (type $12) (param $x i32)
   ;; CHECK-NEXT:  (local $ref (ref null $struct.A))
   ;; CHECK-NEXT:  (local $2 i32)
   ;; CHECK-NEXT:  (local $3 f64)
@@ -2128,7 +2131,7 @@
     )
   )
 
-  ;; CHECK:      (func $ref-eq-flip (type $12) (param $other eqref) (result i32)
+  ;; CHECK:      (func $ref-eq-flip (type $13) (param $other eqref) (result i32)
   ;; CHECK-NEXT:  (local $1 i32)
   ;; CHECK-NEXT:  (local $2 f64)
   ;; CHECK-NEXT:  (local $3 i32)
@@ -2280,7 +2283,7 @@
     )
   )
 
-  ;; CHECK:      (func $ref-eq-unrelated (type $13) (param $x eqref) (param $y eqref) (result i32)
+  ;; CHECK:      (func $ref-eq-unrelated (type $14) (param $x eqref) (param $y eqref) (result i32)
   ;; CHECK-NEXT:  (local $2 i32)
   ;; CHECK-NEXT:  (local $3 f64)
   ;; CHECK-NEXT:  (local $4 i32)
@@ -2387,7 +2390,7 @@
     )
   )
 
-  ;; CHECK:      (func $ref-is (type $14) (param $x anyref) (result i32)
+  ;; CHECK:      (func $ref-is (type $15) (param $x anyref) (result i32)
   ;; CHECK-NEXT:  (local $1 i32)
   ;; CHECK-NEXT:  (local $2 f64)
   ;; CHECK-NEXT:  (local $3 i32)
@@ -2438,7 +2441,93 @@
     )
   )
 
-  (func $ref-test (param $x anyref)
+  ;; CHECK:      (func $ref-test (type $1)
+  ;; CHECK-NEXT:  (local $0 i32)
+  ;; CHECK-NEXT:  (local $1 f64)
+  ;; CHECK-NEXT:  (local $2 i32)
+  ;; CHECK-NEXT:  (local $3 f64)
+  ;; CHECK-NEXT:  (local $4 i32)
+  ;; CHECK-NEXT:  (local $5 f64)
+  ;; CHECK-NEXT:  (local $6 i32)
+  ;; CHECK-NEXT:  (local $7 f64)
+  ;; CHECK-NEXT:  (local $8 i32)
+  ;; CHECK-NEXT:  (local $9 f64)
+  ;; CHECK-NEXT:  (local $10 i32)
+  ;; CHECK-NEXT:  (local $11 f64)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result i32)
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (ref.test (ref $struct.A)
+  ;; CHECK-NEXT:      (block (result nullref)
+  ;; CHECK-NEXT:       (local.set $2
+  ;; CHECK-NEXT:        (i32.const 0)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (local.set $3
+  ;; CHECK-NEXT:        (f64.const 0)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (local.set $0
+  ;; CHECK-NEXT:        (local.get $2)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (local.set $1
+  ;; CHECK-NEXT:        (local.get $3)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (ref.null none)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i32.const 1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result i32)
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (ref.test (ref $struct.A)
+  ;; CHECK-NEXT:      (block (result nullref)
+  ;; CHECK-NEXT:       (local.set $6
+  ;; CHECK-NEXT:        (i32.const 1)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (local.set $7
+  ;; CHECK-NEXT:        (f64.const 2.2)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (local.set $4
+  ;; CHECK-NEXT:        (local.get $6)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (local.set $5
+  ;; CHECK-NEXT:        (local.get $7)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (ref.null none)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i32.const 1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result i32)
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (ref.test (ref $struct.A)
+  ;; CHECK-NEXT:      (block (result nullref)
+  ;; CHECK-NEXT:       (local.set $10
+  ;; CHECK-NEXT:        (i32.const 3)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (local.set $11
+  ;; CHECK-NEXT:        (f64.const 4.4)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (local.set $8
+  ;; CHECK-NEXT:        (local.get $10)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (local.set $9
+  ;; CHECK-NEXT:        (local.get $11)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (ref.null none)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i32.const 1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $ref-test
     ;; This cast must succeed (it tests the exact type), and we can remove the
     ;; allocation.
     (drop
@@ -2467,7 +2556,70 @@
       )
     )
   )
-  (func $ref-test-bad
+  ;; CHECK:      (func $ref-test-bad (type $16) (param $x anyref)
+  ;; CHECK-NEXT:  (local $1 i32)
+  ;; CHECK-NEXT:  (local $2 f64)
+  ;; CHECK-NEXT:  (local $3 i32)
+  ;; CHECK-NEXT:  (local $4 f64)
+  ;; CHECK-NEXT:  (local $5 i32)
+  ;; CHECK-NEXT:  (local $6 f64)
+  ;; CHECK-NEXT:  (local $7 i32)
+  ;; CHECK-NEXT:  (local $8 f64)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (ref.test (ref $struct.A)
+  ;; CHECK-NEXT:    (local.get $x)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result i32)
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (ref.test (ref $struct.B)
+  ;; CHECK-NEXT:      (block (result nullref)
+  ;; CHECK-NEXT:       (local.set $3
+  ;; CHECK-NEXT:        (i32.const 0)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (local.set $4
+  ;; CHECK-NEXT:        (f64.const 0)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (local.set $1
+  ;; CHECK-NEXT:        (local.get $3)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (local.set $2
+  ;; CHECK-NEXT:        (local.get $4)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (ref.null none)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i32.const 0)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result i32)
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (ref.test (ref $struct.B)
+  ;; CHECK-NEXT:      (block (result nullref)
+  ;; CHECK-NEXT:       (local.set $7
+  ;; CHECK-NEXT:        (i32.const 1)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (local.set $8
+  ;; CHECK-NEXT:        (f64.const 2.2)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (local.set $5
+  ;; CHECK-NEXT:        (local.get $7)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (local.set $6
+  ;; CHECK-NEXT:        (local.get $8)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (ref.null none)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i32.const 0)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $ref-test-bad (param $x anyref)
     ;; We know nothing about this cast.
     (drop
       (ref.test (ref $struct.A)

From fd71b16139495f6190e2eeb8d967f907fa821201 Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Tue, 25 Jun 2024 12:00:39 -0700
Subject: [PATCH 19/21] formt

---
 src/passes/Heap2Local.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index e786f37d605..acf4e5cb394 100644
--- a/src/passes/Heap2Local.cpp
+++ b/src/passes/Heap2Local.cpp
@@ -725,8 +725,8 @@ struct Struct2Local : PostWalker<Struct2Local> {
 
     // The result must be 0, since the allocation is not null. Drop the RefIs
     // and append that.
-    replaceCurrent(builder.makeSequence(builder.makeDrop(curr),
-                                        builder.makeConst(Literal(int32_t(0)))));
+    replaceCurrent(builder.makeSequence(
+      builder.makeDrop(curr), builder.makeConst(Literal(int32_t(0)))));
   }
 
   void visitRefEq(RefEq* curr) {

From 4193d58a4b0ea2ac86a4a03fd3af675882358edc Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Tue, 25 Jun 2024 13:35:20 -0700
Subject: [PATCH 20/21] test

---
 src/passes/Heap2Local.cpp       |  24 +++++-
 test/lit/passes/heap2local.wast | 140 +++++++++++++++++++++++++++++---
 2 files changed, 152 insertions(+), 12 deletions(-)

diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index acf4e5cb394..9e249d438e2 100644
--- a/src/passes/Heap2Local.cpp
+++ b/src/passes/Heap2Local.cpp
@@ -857,12 +857,15 @@ struct Array2Struct : PostWalker<Array2Struct> {
   EscapeAnalyzer& analyzer;
   Function* func;
   Builder builder;
+  // The original type of the allocation, before we turn it into a struct.
+  Type originalType;
 
   Array2Struct(Expression* allocation,
                EscapeAnalyzer& analyzer,
                Function* func,
                Module& wasm)
-    : allocation(allocation), analyzer(analyzer), func(func), builder(wasm) {
+    : allocation(allocation), analyzer(analyzer), func(func), builder(wasm),
+      originalType(allocation->type) {
 
     // Build the struct type we need: as many fields as the size of the array,
     // all of the same type as the array's element.
@@ -1026,6 +1029,25 @@ struct Array2Struct : PostWalker<Array2Struct> {
     noteCurrentIsReached();
   }
 
+  // Some additional operations need special handling
+  void visitRefTest(RefTest* curr) {
+    if (!analyzer.reached.count(curr)) {
+      return;
+    }
+
+    // When we ref.test an array allocation, we cannot simply turn the array
+    // into a struct, as then the test will behave different. (Note that this is
+    // not a problem for ref.*cast*, as the cast simply goes away when the value
+    // flows through, and we verify it will do so in the escape analysis.) To
+    // handle this, check if the test succeeds or not, and write out the outcome
+    // here (similar to Struct2Local::visitRefTest). Note that we test on
+    // |originalType| here and not |allocation->type|, as the allocation has
+    // been turned into a struct.
+    int32_t result = Type::isSubType(originalType, curr->castType);
+    replaceCurrent(builder.makeSequence(builder.makeDrop(curr),
+                                        builder.makeConst(Literal(result))));
+  }
+
   // Get the value in an expression we know must contain a constant index.
   Index getIndex(Expression* curr) {
     return curr->cast<Const>()->value.getUnsigned();
diff --git a/test/lit/passes/heap2local.wast b/test/lit/passes/heap2local.wast
index ed2af717d98..6879671743d 100644
--- a/test/lit/passes/heap2local.wast
+++ b/test/lit/passes/heap2local.wast
@@ -2924,9 +2924,9 @@
 
   ;; CHECK:      (type $1 (struct (field (mut i32))))
 
-  ;; CHECK:      (type $2 (func (result i32)))
+  ;; CHECK:      (type $2 (func))
 
-  ;; CHECK:      (type $3 (func))
+  ;; CHECK:      (type $3 (func (result i32)))
 
   ;; CHECK:      (type $4 (func (param i32) (result i32)))
 
@@ -2940,7 +2940,9 @@
 
   ;; CHECK:      (type $9 (struct (field (mut i32)) (field (mut i32)) (field (mut i32)) (field (mut i32)) (field (mut i32)) (field (mut i32)) (field (mut i32)) (field (mut i32)) (field (mut i32)) (field (mut i32)) (field (mut i32)) (field (mut i32)) (field (mut i32)) (field (mut i32)) (field (mut i32)) (field (mut i32)) (field (mut i32)) (field (mut i32)) (field (mut i32))))
 
-  ;; CHECK:      (func $array.new_default (type $3)
+  ;; CHECK:      (type $10 (func (param anyref)))
+
+  ;; CHECK:      (func $array.new_default (type $2)
   ;; CHECK-NEXT:  (local $temp (ref $array))
   ;; CHECK-NEXT:  (local $1 i32)
   ;; CHECK-NEXT:  (local $2 i32)
@@ -3080,7 +3082,7 @@
     )
   )
 
-  ;; CHECK:      (func $array.new (type $2) (result i32)
+  ;; CHECK:      (func $array.new (type $3) (result i32)
   ;; CHECK-NEXT:  (local $temp (ref $array))
   ;; CHECK-NEXT:  (local $1 i32)
   ;; CHECK-NEXT:  (local $2 i32)
@@ -3262,7 +3264,7 @@
     )
   )
 
-  ;; CHECK:      (func $array.local.super (type $3)
+  ;; CHECK:      (func $array.local.super (type $2)
   ;; CHECK-NEXT:  (local $temp anyref)
   ;; CHECK-NEXT:  (local $1 i32)
   ;; CHECK-NEXT:  (local $2 i32)
@@ -3311,7 +3313,7 @@
     )
   )
 
-  ;; CHECK:      (func $array.folded (type $2) (result i32)
+  ;; CHECK:      (func $array.folded (type $3) (result i32)
   ;; CHECK-NEXT:  (local $0 i32)
   ;; CHECK-NEXT:  (local $1 i32)
   ;; CHECK-NEXT:  (local $2 i32)
@@ -3392,7 +3394,7 @@
     )
   )
 
-  ;; CHECK:      (func $array.folded.multiple (type $3)
+  ;; CHECK:      (func $array.folded.multiple (type $2)
   ;; CHECK-NEXT:  (local $0 i32)
   ;; CHECK-NEXT:  (local $1 i32)
   ;; CHECK-NEXT:  (local $2 i32)
@@ -3645,7 +3647,7 @@
     )
   )
 
-  ;; CHECK:      (func $array.nested.refinalize.get (type $2) (result i32)
+  ;; CHECK:      (func $array.nested.refinalize.get (type $3) (result i32)
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (block (result nullref)
   ;; CHECK-NEXT:    (ref.null none)
@@ -3665,7 +3667,7 @@
     )
   )
 
-  ;; CHECK:      (func $array.nested.refinalize.set (type $3)
+  ;; CHECK:      (func $array.nested.refinalize.set (type $2)
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (block (result nullref)
   ;; CHECK-NEXT:    (ref.null none)
@@ -3687,7 +3689,7 @@
     )
   )
 
-  ;; CHECK:      (func $array.flowing.type (type $2) (result i32)
+  ;; CHECK:      (func $array.flowing.type (type $3) (result i32)
   ;; CHECK-NEXT:  (local $temp (ref $array))
   ;; CHECK-NEXT:  (local $1 i32)
   ;; CHECK-NEXT:  (local $2 i32)
@@ -3773,13 +3775,129 @@
     )
   )
 
-  ;; CHECK:      (func $get-i32 (type $2) (result i32)
+  ;; CHECK:      (func $get-i32 (type $3) (result i32)
   ;; CHECK-NEXT:  (i32.const 1337)
   ;; CHECK-NEXT: )
   (func $get-i32 (result i32)
     ;; Helper for the above.
     (i32.const 1337)
   )
+
+  ;; CHECK:      (func $ref-test (type $2)
+  ;; CHECK-NEXT:  (local $0 i32)
+  ;; CHECK-NEXT:  (local $1 i32)
+  ;; CHECK-NEXT:  (local $2 i32)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result i32)
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (block (result i32)
+  ;; CHECK-NEXT:      (drop
+  ;; CHECK-NEXT:       (ref.test (ref $array)
+  ;; CHECK-NEXT:        (block (result nullref)
+  ;; CHECK-NEXT:         (local.set $0
+  ;; CHECK-NEXT:          (i32.const 0)
+  ;; CHECK-NEXT:         )
+  ;; CHECK-NEXT:         (ref.null none)
+  ;; CHECK-NEXT:        )
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (i32.const 0)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i32.const 1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result i32)
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (block (result i32)
+  ;; CHECK-NEXT:      (drop
+  ;; CHECK-NEXT:       (ref.test (ref $array)
+  ;; CHECK-NEXT:        (block (result nullref)
+  ;; CHECK-NEXT:         (local.set $1
+  ;; CHECK-NEXT:          (i32.const 0)
+  ;; CHECK-NEXT:         )
+  ;; CHECK-NEXT:         (local.set $2
+  ;; CHECK-NEXT:          (i32.const 0)
+  ;; CHECK-NEXT:         )
+  ;; CHECK-NEXT:         (ref.null none)
+  ;; CHECK-NEXT:        )
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (i32.const 0)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i32.const 1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $ref-test
+    ;; This cast must succeed (it tests the exact type), and we can remove the
+    ;; allocation.
+    (drop
+      (ref.test (ref $array)
+        (array.new_default $array
+          (i32.const 1)
+        )
+      )
+    )
+    ;; Testing a supertype also works.
+    (drop
+      (ref.test (ref null any)
+        (array.new_default $array
+          (i32.const 2)
+        )
+      )
+    )
+  )
+  ;; CHECK:      (func $ref-test-bad (type $10) (param $x anyref)
+  ;; CHECK-NEXT:  (local $1 i32)
+  ;; CHECK-NEXT:  (local $2 i32)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (ref.test (ref $array)
+  ;; CHECK-NEXT:    (local.get $x)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result i32)
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (block (result i32)
+  ;; CHECK-NEXT:      (drop
+  ;; CHECK-NEXT:       (ref.test (ref none)
+  ;; CHECK-NEXT:        (block (result nullref)
+  ;; CHECK-NEXT:         (local.set $1
+  ;; CHECK-NEXT:          (i32.const 0)
+  ;; CHECK-NEXT:         )
+  ;; CHECK-NEXT:         (local.set $2
+  ;; CHECK-NEXT:          (i32.const 0)
+  ;; CHECK-NEXT:         )
+  ;; CHECK-NEXT:         (ref.null none)
+  ;; CHECK-NEXT:        )
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (i32.const 0)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i32.const 0)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $ref-test-bad (param $x anyref)
+    ;; We know nothing about this cast.
+    (drop
+      (ref.test (ref $array)
+        (local.get $x)
+      )
+    )
+    ;; This cast fails, but we can remove the allocation.
+    (drop
+      (ref.test (ref struct)
+        (array.new_default $array
+          (i32.const 2)
+        )
+      )
+    )
+  )
 )
 
 ;; Arrays with reference values.

From 1eb9848118df102f4d2e25c4484501b85427d768 Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Tue, 25 Jun 2024 16:34:16 -0700
Subject: [PATCH 21/21] fix

---
 src/passes/Heap2Local.cpp       |   6 +-
 test/lit/passes/heap2local.wast | 215 +++++++++++++++-----------------
 2 files changed, 103 insertions(+), 118 deletions(-)

diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index 9e249d438e2..7aec23ccf1a 100644
--- a/src/passes/Heap2Local.cpp
+++ b/src/passes/Heap2Local.cpp
@@ -771,8 +771,10 @@ struct Struct2Local : PostWalker<Struct2Local> {
     // GCTypeUtils::evaluateCastCheck because we know the allocation's type
     // precisely (it cannot be a strict subtype of the type - it is the type).
     int32_t result = Type::isSubType(allocation->type, curr->castType);
-    // For simplicity, simply drop the RefEq and put a constant result after.
-    replaceCurrent(builder.makeSequence(builder.makeDrop(curr),
+    // Remove the RefTest and leave only its reference child. If we kept it,
+    // we'd need to refinalize (as the input to the test changes, since the
+    // reference becomes a null, which has a different type).
+    replaceCurrent(builder.makeSequence(builder.makeDrop(curr->ref),
                                         builder.makeConst(Literal(result))));
   }
 
diff --git a/test/lit/passes/heap2local.wast b/test/lit/passes/heap2local.wast
index 6879671743d..d642aabb3da 100644
--- a/test/lit/passes/heap2local.wast
+++ b/test/lit/passes/heap2local.wast
@@ -7,6 +7,8 @@
   ;; CHECK:      (type $struct.A (sub (struct (field (mut i32)) (field (mut f64)))))
   (type $struct.A (sub (struct (field (mut i32)) (field (mut f64)))))
 
+  (type $struct.B (sub $struct.A (struct (field (mut i32)) (field (mut f64)))))
+
   ;; CHECK:      (type $1 (func))
 
   ;; CHECK:      (type $2 (func (result f64)))
@@ -21,9 +23,6 @@
 
   ;; CHECK:      (type $7 (func (param i32) (result f64)))
 
-  ;; CHECK:      (type $struct.B (sub $struct.A (struct (field (mut i32)) (field (mut f64)))))
-  (type $struct.B (sub $struct.A (struct (field (mut i32)) (field (mut f64)))))
-
   ;; CHECK:      (type $struct.packed (struct (field (mut i8)) (field (mut i32))))
   (type $struct.packed (struct (field (mut i8)) (field (mut i32))))
 
@@ -33,19 +32,19 @@
 
   (type $struct.nonnullable (struct (field (ref $struct.A))))
 
-  ;; CHECK:      (type $10 (func (param (ref null $struct.recursive))))
+  ;; CHECK:      (type $9 (func (param (ref null $struct.recursive))))
 
-  ;; CHECK:      (type $11 (func (param (ref $struct.A))))
+  ;; CHECK:      (type $10 (func (param (ref $struct.A))))
 
-  ;; CHECK:      (type $12 (func (param i32)))
+  ;; CHECK:      (type $11 (func (param i32)))
 
-  ;; CHECK:      (type $13 (func (param eqref) (result i32)))
+  ;; CHECK:      (type $12 (func (param eqref) (result i32)))
 
-  ;; CHECK:      (type $14 (func (param eqref eqref) (result i32)))
+  ;; CHECK:      (type $13 (func (param eqref eqref) (result i32)))
 
-  ;; CHECK:      (type $15 (func (param anyref) (result i32)))
+  ;; CHECK:      (type $14 (func (param anyref) (result i32)))
 
-  ;; CHECK:      (type $16 (func (param anyref)))
+  ;; CHECK:      (type $15 (func (param anyref)))
 
   ;; CHECK:      (func $simple (type $1)
   ;; CHECK-NEXT:  (local $0 i32)
@@ -951,7 +950,7 @@
     )
   )
 
-  ;; CHECK:      (func $set-value (type $10) (param $struct.recursive (ref null $struct.recursive))
+  ;; CHECK:      (func $set-value (type $9) (param $struct.recursive (ref null $struct.recursive))
   ;; CHECK-NEXT:  (local $ref (ref null $struct.recursive))
   ;; CHECK-NEXT:  (struct.set $struct.recursive 0
   ;; CHECK-NEXT:   (local.get $struct.recursive)
@@ -1066,7 +1065,7 @@
     )
   )
 
-  ;; CHECK:      (func $non-nullable (type $11) (param $a (ref $struct.A))
+  ;; CHECK:      (func $non-nullable (type $10) (param $a (ref $struct.A))
   ;; CHECK-NEXT:  (local $1 (ref $struct.A))
   ;; CHECK-NEXT:  (local $2 (ref $struct.A))
   ;; CHECK-NEXT:  (drop
@@ -1098,7 +1097,7 @@
     )
   )
 
-  ;; CHECK:      (func $before-loop-use-multi (type $12) (param $x i32)
+  ;; CHECK:      (func $before-loop-use-multi (type $11) (param $x i32)
   ;; CHECK-NEXT:  (local $ref (ref null $struct.A))
   ;; CHECK-NEXT:  (local $2 i32)
   ;; CHECK-NEXT:  (local $3 f64)
@@ -2131,7 +2130,7 @@
     )
   )
 
-  ;; CHECK:      (func $ref-eq-flip (type $13) (param $other eqref) (result i32)
+  ;; CHECK:      (func $ref-eq-flip (type $12) (param $other eqref) (result i32)
   ;; CHECK-NEXT:  (local $1 i32)
   ;; CHECK-NEXT:  (local $2 f64)
   ;; CHECK-NEXT:  (local $3 i32)
@@ -2283,7 +2282,7 @@
     )
   )
 
-  ;; CHECK:      (func $ref-eq-unrelated (type $14) (param $x eqref) (param $y eqref) (result i32)
+  ;; CHECK:      (func $ref-eq-unrelated (type $13) (param $x eqref) (param $y eqref) (result i32)
   ;; CHECK-NEXT:  (local $2 i32)
   ;; CHECK-NEXT:  (local $3 f64)
   ;; CHECK-NEXT:  (local $4 i32)
@@ -2390,7 +2389,7 @@
     )
   )
 
-  ;; CHECK:      (func $ref-is (type $15) (param $x anyref) (result i32)
+  ;; CHECK:      (func $ref-is (type $14) (param $x anyref) (result i32)
   ;; CHECK-NEXT:  (local $1 i32)
   ;; CHECK-NEXT:  (local $2 f64)
   ;; CHECK-NEXT:  (local $3 i32)
@@ -2457,22 +2456,20 @@
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (block (result i32)
   ;; CHECK-NEXT:    (drop
-  ;; CHECK-NEXT:     (ref.test (ref $struct.A)
-  ;; CHECK-NEXT:      (block (result nullref)
-  ;; CHECK-NEXT:       (local.set $2
-  ;; CHECK-NEXT:        (i32.const 0)
-  ;; CHECK-NEXT:       )
-  ;; CHECK-NEXT:       (local.set $3
-  ;; CHECK-NEXT:        (f64.const 0)
-  ;; CHECK-NEXT:       )
-  ;; CHECK-NEXT:       (local.set $0
-  ;; CHECK-NEXT:        (local.get $2)
-  ;; CHECK-NEXT:       )
-  ;; CHECK-NEXT:       (local.set $1
-  ;; CHECK-NEXT:        (local.get $3)
-  ;; CHECK-NEXT:       )
-  ;; CHECK-NEXT:       (ref.null none)
+  ;; CHECK-NEXT:     (block (result nullref)
+  ;; CHECK-NEXT:      (local.set $2
+  ;; CHECK-NEXT:       (i32.const 0)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (local.set $3
+  ;; CHECK-NEXT:       (f64.const 0)
   ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (local.set $0
+  ;; CHECK-NEXT:       (local.get $2)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (local.set $1
+  ;; CHECK-NEXT:       (local.get $3)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (ref.null none)
   ;; CHECK-NEXT:     )
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:    (i32.const 1)
@@ -2481,22 +2478,20 @@
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (block (result i32)
   ;; CHECK-NEXT:    (drop
-  ;; CHECK-NEXT:     (ref.test (ref $struct.A)
-  ;; CHECK-NEXT:      (block (result nullref)
-  ;; CHECK-NEXT:       (local.set $6
-  ;; CHECK-NEXT:        (i32.const 1)
-  ;; CHECK-NEXT:       )
-  ;; CHECK-NEXT:       (local.set $7
-  ;; CHECK-NEXT:        (f64.const 2.2)
-  ;; CHECK-NEXT:       )
-  ;; CHECK-NEXT:       (local.set $4
-  ;; CHECK-NEXT:        (local.get $6)
-  ;; CHECK-NEXT:       )
-  ;; CHECK-NEXT:       (local.set $5
-  ;; CHECK-NEXT:        (local.get $7)
-  ;; CHECK-NEXT:       )
-  ;; CHECK-NEXT:       (ref.null none)
+  ;; CHECK-NEXT:     (block (result nullref)
+  ;; CHECK-NEXT:      (local.set $6
+  ;; CHECK-NEXT:       (i32.const 1)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (local.set $7
+  ;; CHECK-NEXT:       (f64.const 2.2)
   ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (local.set $4
+  ;; CHECK-NEXT:       (local.get $6)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (local.set $5
+  ;; CHECK-NEXT:       (local.get $7)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (ref.null none)
   ;; CHECK-NEXT:     )
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:    (i32.const 1)
@@ -2505,22 +2500,20 @@
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (block (result i32)
   ;; CHECK-NEXT:    (drop
-  ;; CHECK-NEXT:     (ref.test (ref $struct.A)
-  ;; CHECK-NEXT:      (block (result nullref)
-  ;; CHECK-NEXT:       (local.set $10
-  ;; CHECK-NEXT:        (i32.const 3)
-  ;; CHECK-NEXT:       )
-  ;; CHECK-NEXT:       (local.set $11
-  ;; CHECK-NEXT:        (f64.const 4.4)
-  ;; CHECK-NEXT:       )
-  ;; CHECK-NEXT:       (local.set $8
-  ;; CHECK-NEXT:        (local.get $10)
-  ;; CHECK-NEXT:       )
-  ;; CHECK-NEXT:       (local.set $9
-  ;; CHECK-NEXT:        (local.get $11)
-  ;; CHECK-NEXT:       )
-  ;; CHECK-NEXT:       (ref.null none)
+  ;; CHECK-NEXT:     (block (result nullref)
+  ;; CHECK-NEXT:      (local.set $10
+  ;; CHECK-NEXT:       (i32.const 3)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (local.set $11
+  ;; CHECK-NEXT:       (f64.const 4.4)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (local.set $8
+  ;; CHECK-NEXT:       (local.get $10)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (local.set $9
+  ;; CHECK-NEXT:       (local.get $11)
   ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (ref.null none)
   ;; CHECK-NEXT:     )
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:    (i32.const 1)
@@ -2556,7 +2549,7 @@
       )
     )
   )
-  ;; CHECK:      (func $ref-test-bad (type $16) (param $x anyref)
+  ;; CHECK:      (func $ref-test-bad (type $15) (param $x anyref)
   ;; CHECK-NEXT:  (local $1 i32)
   ;; CHECK-NEXT:  (local $2 f64)
   ;; CHECK-NEXT:  (local $3 i32)
@@ -2573,22 +2566,20 @@
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (block (result i32)
   ;; CHECK-NEXT:    (drop
-  ;; CHECK-NEXT:     (ref.test (ref $struct.B)
-  ;; CHECK-NEXT:      (block (result nullref)
-  ;; CHECK-NEXT:       (local.set $3
-  ;; CHECK-NEXT:        (i32.const 0)
-  ;; CHECK-NEXT:       )
-  ;; CHECK-NEXT:       (local.set $4
-  ;; CHECK-NEXT:        (f64.const 0)
-  ;; CHECK-NEXT:       )
-  ;; CHECK-NEXT:       (local.set $1
-  ;; CHECK-NEXT:        (local.get $3)
-  ;; CHECK-NEXT:       )
-  ;; CHECK-NEXT:       (local.set $2
-  ;; CHECK-NEXT:        (local.get $4)
-  ;; CHECK-NEXT:       )
-  ;; CHECK-NEXT:       (ref.null none)
+  ;; CHECK-NEXT:     (block (result nullref)
+  ;; CHECK-NEXT:      (local.set $3
+  ;; CHECK-NEXT:       (i32.const 0)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (local.set $4
+  ;; CHECK-NEXT:       (f64.const 0)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (local.set $1
+  ;; CHECK-NEXT:       (local.get $3)
   ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (local.set $2
+  ;; CHECK-NEXT:       (local.get $4)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (ref.null none)
   ;; CHECK-NEXT:     )
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:    (i32.const 0)
@@ -2597,22 +2588,20 @@
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (block (result i32)
   ;; CHECK-NEXT:    (drop
-  ;; CHECK-NEXT:     (ref.test (ref $struct.B)
-  ;; CHECK-NEXT:      (block (result nullref)
-  ;; CHECK-NEXT:       (local.set $7
-  ;; CHECK-NEXT:        (i32.const 1)
-  ;; CHECK-NEXT:       )
-  ;; CHECK-NEXT:       (local.set $8
-  ;; CHECK-NEXT:        (f64.const 2.2)
-  ;; CHECK-NEXT:       )
-  ;; CHECK-NEXT:       (local.set $5
-  ;; CHECK-NEXT:        (local.get $7)
-  ;; CHECK-NEXT:       )
-  ;; CHECK-NEXT:       (local.set $6
-  ;; CHECK-NEXT:        (local.get $8)
-  ;; CHECK-NEXT:       )
-  ;; CHECK-NEXT:       (ref.null none)
+  ;; CHECK-NEXT:     (block (result nullref)
+  ;; CHECK-NEXT:      (local.set $7
+  ;; CHECK-NEXT:       (i32.const 1)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (local.set $8
+  ;; CHECK-NEXT:       (f64.const 2.2)
   ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (local.set $5
+  ;; CHECK-NEXT:       (local.get $7)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (local.set $6
+  ;; CHECK-NEXT:       (local.get $8)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (ref.null none)
   ;; CHECK-NEXT:     )
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:    (i32.const 0)
@@ -3792,13 +3781,11 @@
   ;; CHECK-NEXT:    (drop
   ;; CHECK-NEXT:     (block (result i32)
   ;; CHECK-NEXT:      (drop
-  ;; CHECK-NEXT:       (ref.test (ref $array)
-  ;; CHECK-NEXT:        (block (result nullref)
-  ;; CHECK-NEXT:         (local.set $0
-  ;; CHECK-NEXT:          (i32.const 0)
-  ;; CHECK-NEXT:         )
-  ;; CHECK-NEXT:         (ref.null none)
+  ;; CHECK-NEXT:       (block (result nullref)
+  ;; CHECK-NEXT:        (local.set $0
+  ;; CHECK-NEXT:         (i32.const 0)
   ;; CHECK-NEXT:        )
+  ;; CHECK-NEXT:        (ref.null none)
   ;; CHECK-NEXT:       )
   ;; CHECK-NEXT:      )
   ;; CHECK-NEXT:      (i32.const 0)
@@ -3812,16 +3799,14 @@
   ;; CHECK-NEXT:    (drop
   ;; CHECK-NEXT:     (block (result i32)
   ;; CHECK-NEXT:      (drop
-  ;; CHECK-NEXT:       (ref.test (ref $array)
-  ;; CHECK-NEXT:        (block (result nullref)
-  ;; CHECK-NEXT:         (local.set $1
-  ;; CHECK-NEXT:          (i32.const 0)
-  ;; CHECK-NEXT:         )
-  ;; CHECK-NEXT:         (local.set $2
-  ;; CHECK-NEXT:          (i32.const 0)
-  ;; CHECK-NEXT:         )
-  ;; CHECK-NEXT:         (ref.null none)
+  ;; CHECK-NEXT:       (block (result nullref)
+  ;; CHECK-NEXT:        (local.set $1
+  ;; CHECK-NEXT:         (i32.const 0)
+  ;; CHECK-NEXT:        )
+  ;; CHECK-NEXT:        (local.set $2
+  ;; CHECK-NEXT:         (i32.const 0)
   ;; CHECK-NEXT:        )
+  ;; CHECK-NEXT:        (ref.null none)
   ;; CHECK-NEXT:       )
   ;; CHECK-NEXT:      )
   ;; CHECK-NEXT:      (i32.const 0)
@@ -3863,16 +3848,14 @@
   ;; CHECK-NEXT:    (drop
   ;; CHECK-NEXT:     (block (result i32)
   ;; CHECK-NEXT:      (drop
-  ;; CHECK-NEXT:       (ref.test (ref none)
-  ;; CHECK-NEXT:        (block (result nullref)
-  ;; CHECK-NEXT:         (local.set $1
-  ;; CHECK-NEXT:          (i32.const 0)
-  ;; CHECK-NEXT:         )
-  ;; CHECK-NEXT:         (local.set $2
-  ;; CHECK-NEXT:          (i32.const 0)
-  ;; CHECK-NEXT:         )
-  ;; CHECK-NEXT:         (ref.null none)
+  ;; CHECK-NEXT:       (block (result nullref)
+  ;; CHECK-NEXT:        (local.set $1
+  ;; CHECK-NEXT:         (i32.const 0)
+  ;; CHECK-NEXT:        )
+  ;; CHECK-NEXT:        (local.set $2
+  ;; CHECK-NEXT:         (i32.const 0)
   ;; CHECK-NEXT:        )
+  ;; CHECK-NEXT:        (ref.null none)
   ;; CHECK-NEXT:       )
   ;; CHECK-NEXT:      )
   ;; CHECK-NEXT:      (i32.const 0)