[Wasm GC] Heap2Local: Replace the allocation with null (#3893)

Previously we would try to stop using the allocation as much as possible,
for example not writing it to locals any more, and leaving it to other passes
to actually remove it (and remove gets of those locals etc.). This seemed
simpler and more modular, but does not actually work in some cases as
the fuzzer has found. Specifically, if we stop writing our allocation to
locals, then if we do a (ref.as_non_null (local.get ..)) of that, then we
will trap on the null present in the local.

Instead, this changes our rewriting to do slightly more work, but it is
simpler in the end. We replace the allocation with a null, and replace
all the places that use it accordingly, for example, updating types to be
nullable, and removing RefAsNonNulls, etc. This literally gets rid of the
allocation and all the places it flows to (leaving less for other passes to
do later).

1 files changed, 54 insertions, 24 deletions

diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index afdcab06f..a2665958e 100644
--- a/src/passes/Heap2Local.cpp
+++ b/src/passes/Heap2Local.cpp
@@ -265,6 +265,31 @@ struct Heap2LocalOptimizer {
       walk(func->body);
     }
 
+    // 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
+    // remove all uses of it as much as possible, using the information we have
+    // (for example, when our allocation reaches a RefAsNonNull we can simply
+    // remove that operation as we know it would not throw). Some things are
+    // left to other passes, like getting rid of dropped code without side
+    // effects.
+
+    void visitBlock(Block* curr) {
+      if (!reached.count(curr)) {
+        return;
+      }
+
+      // Our allocation passes through this block. We must turn its type into a
+      // nullable one, because we will remove things like RefAsNonNull of it,
+      // which means we may no longer have a non-nullable value as our input,
+      // and we could fail to validate. It is safe to make this change in terms
+      // of our parent, since we know very specifically that only safe things
+      // will end up using our value, like a StructGet or a Drop, which do not
+      // care about non-nullability.
+      assert(curr->type.isRef());
+      curr->type = Type(curr->type.getHeapType(), Nullable);
+    }
+
     void visitLocalSet(LocalSet* curr) {
       if (!reached.count(curr)) {
         return;
@@ -285,16 +310,21 @@ struct Heap2LocalOptimizer {
       }
     }
 
+    void visitBreak(Break* curr) {
+      if (!reached.count(curr)) {
+        return;
+      }
+
+      // Breaks that our allocation flows through may change type, as we now
+      // have a nullable type there.
+      curr->finalize();
+    }
+
     void visitStructNew(StructNew* curr) {
       if (curr != allocation) {
         return;
       }
 
-      // We do not remove the allocation itself here, rather we make it
-      // unnecessary, and then depend on other optimizations to clean up. (We
-      // cannot simply remove it because we need to replace it with something of
-      // the same non-nullable type.)
-
       // First, assign the initial values to the new locals.
       std::vector<Expression*> contents;
 
@@ -340,21 +370,12 @@ struct Heap2LocalOptimizer {
             builder.makeLocalGet(tempIndexes[i], fields[i].type)));
         }
 
-        // Read the values in the allocation (we don't need to, as the
-        // allocation is not used after our optimization, but we need something
-        // with the right type anyhow).
-        for (Index i = 0; i < tempIndexes.size(); i++) {
-          allocation->operands[i] =
-            builder.makeLocalGet(localIndexes[i], fields[i].type);
-        }
-
         // TODO Check if the nondefault case does not increase code size in some
         //      cases. A heap allocation that implicitly sets the default values
         //      is smaller than multiple explicit settings of locals to
         //      defaults.
       } else {
-        // Set the default values, and replace the allocation with a block that
-        // first does that, then contains the allocation.
+        // Set the default values.
         // Note that we must assign the defaults because we might be in a loop,
         // that is, there might be a previous value.
         for (Index i = 0; i < localIndexes.size(); i++) {
@@ -364,12 +385,27 @@ struct Heap2LocalOptimizer {
         }
       }
 
-      // Put the allocation itself at the end of the block, so the block has the
-      // exact same type as the allocation it replaces.
-      contents.push_back(allocation);
+      // Drop the RTT (as it may have side effects; leave it to other passes).
+      contents.push_back(builder.makeDrop(allocation->rtt));
+      // Replace the allocation with a null reference. This changes the type
+      // from non-nullable to nullable, but as we optimize away the code that
+      // the allocation reaches, we will handle that.
+      contents.push_back(
+        builder.makeRefNull(Type(allocation->type.getHeapType(), Nullable)));
       replaceCurrent(builder.makeBlock(contents));
     }
 
+    void visitRefAs(RefAs* curr) {
+      if (!reached.count(curr)) {
+        return;
+      }
+
+      // It is safe to optimize out this RefAsNonNull, since we proved it
+      // contains our allocation, and so cannot trap.
+      assert(curr->op == RefAsNonNull);
+      replaceCurrent(curr->value);
+    }
+
     void visitStructSet(StructSet* curr) {
       if (!reached.count(curr)) {
         return;
@@ -570,12 +606,6 @@ struct Heap2LocalOptimizer {
           // null (so there is no trap), and we can continue to (hopefully)
           // optimize this allocation.
           escapes = false;
-
-          // Note that while we can look through this operation, we cannot get
-          // rid of it later, as its parent might depend on receiving a
-          // non-nullable type. So we will leave the RefAsNonNull as it is,
-          // even if we do optimize the allocation, and we depend on other
-          // passes to remove the RefAsNonNull.
         }
       }