[Wasm GC] Fix Heap2Local + non-nullable locals (#4017)

Given

(local $x (ref $foo))
(local.set $x ..)
(.. (local.get $x))

If we remove the local.set but not the get, then we end up with

(local $x (ref $foo))
(.. (local.get $x))

It looks like the local.get reads the initial value of a non-nullable local,
which is not allowed.

In practice, this would crash in precompute-propagate which would
try to propagate the initial value to the get. Add an assertion there with
a clear message, as until we have full validation of non-nullable locals
(and the spec for that is in flux), that pass is where bugs will end up
being noticed.

To fix this, replace the get as well. We can replace it with a null
for simplicity; it will never be used anyhow.

This also uncovered a small bug with reached not containing all
the things we reached - it was missing local.gets.

4 files changed, 65 insertions, 41 deletions

diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index a2665958e..ac60c5cac 100644
--- a/src/passes/Heap2Local.cpp
+++ b/src/passes/Heap2Local.cpp
@@ -297,12 +297,6 @@ struct Heap2LocalOptimizer {
 
       // We don't need any sets of the reference to any of the locals it
       // originally was written to.
-      //
-      // Note that after we remove the sets, other passes can easily remove the
-      // gets, and so we do not bother to do anything for them. (Also, in
-      // general it is not trivial to replace the gets - we'd need something of
-      // the same type, but the type might be a non-nullable reference type in
-      // the case of a parameter, and in the future maybe of some locals.)
       if (curr->isTee()) {
         replaceCurrent(curr->value);
       } else {
@@ -310,6 +304,28 @@ struct Heap2LocalOptimizer {
       }
     }
 
+    void visitLocalGet(LocalGet* curr) {
+      if (!reached.count(curr)) {
+        return;
+      }
+
+      // Uses of this get will drop it, so the value does not matter. Replace it
+      // with something else, which avoids issues with non-nullability (when
+      // non-nullable locals are enabled), which could happen like this:
+      //
+      //   (local $x (ref $foo))
+      //   (local.set $x ..)
+      //   (.. (local.get $x))
+      //
+      // If we remove the set but not the get then the get would appear to read
+      // the default value of a non-nullable local, which is not allowed.
+      //
+      // For simplicity, replace the get with a null. We anyhow have null types
+      // in the places where our allocation was earlier, see notes on
+      // visitBlock, and so using a null here adds no extra complexity.
+      replaceCurrent(builder.makeRefNull(curr->type.getHeapType()));
+    }
+
     void visitBreak(Break* curr) {
       if (!reached.count(curr)) {
         return;
@@ -390,8 +406,7 @@ struct Heap2LocalOptimizer {
       // 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)));
+      contents.push_back(builder.makeRefNull(allocation->type.getHeapType()));
       replaceCurrent(builder.makeBlock(contents));
     }
 
@@ -533,8 +548,10 @@ struct Heap2LocalOptimizer {
       }
 
       // If we got to here, then we can continue to hope that we can optimize
-      // this allocation. Mark the parent as reached by it, and continue.
+      // this allocation. Mark the parent and child as reached by it, and
+      // continue.
       rewriter.reached.insert(parent);
+      rewriter.reached.insert(child);
     }
 
     // We finished the loop over the flows. Do the final checks.
diff --git a/src/passes/Precompute.cpp b/src/passes/Precompute.cpp
index a832c9e85..44227a50f 100644
--- a/src/passes/Precompute.cpp
+++ b/src/passes/Precompute.cpp
@@ -350,8 +350,10 @@ private:
           Literals curr;
           if (set == nullptr) {
             if (getFunction()->isVar(get->index)) {
-              curr =
-                Literal::makeZeros(getFunction()->getLocalType(get->index));
+              auto localType = getFunction()->getLocalType(get->index);
+              assert(!localType.isNonNullable() &&
+                     "Non-nullable locals must not use the default value");
+              curr = Literal::makeZeros(localType);
             } else {
               // it's a param, so it's hopeless
               values = {};
diff --git a/src/wasm-builder.h b/src/wasm-builder.h
index 876474223..3eb263580 100644
--- a/src/wasm-builder.h
+++ b/src/wasm-builder.h
@@ -618,6 +618,11 @@ public:
     ret->finalize();
     return ret;
   }
+  RefNull* makeRefNull(HeapType type) {
+    auto* ret = wasm.allocator.alloc<RefNull>();
+    ret->finalize(Type(type, Nullable));
+    return ret;
+  }
   RefNull* makeRefNull(Type type) {
     auto* ret = wasm.allocator.alloc<RefNull>();
     ret->finalize(type);
diff --git a/test/lit/passes/heap2local.wast b/test/lit/passes/heap2local.wast
index f7120906e..61b18f0c4 100644
--- a/test/lit/passes/heap2local.wast
+++ b/test/lit/passes/heap2local.wast
@@ -318,7 +318,7 @@
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (block
   ;; CHECK-NEXT:   (drop
-  ;; CHECK-NEXT:    (local.get $ref)
+  ;; CHECK-NEXT:    (ref.null $struct.A)
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:   (local.set $1
   ;; CHECK-NEXT:    (i32.const 1)
@@ -361,7 +361,7 @@
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (block (result f64)
   ;; CHECK-NEXT:   (drop
-  ;; CHECK-NEXT:    (local.get $ref)
+  ;; CHECK-NEXT:    (ref.null $struct.A)
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:   (local.get $2)
   ;; CHECK-NEXT:  )
@@ -404,11 +404,11 @@
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (drop
-  ;; CHECK-NEXT:   (local.get $ref)
+  ;; CHECK-NEXT:   (ref.null $struct.A)
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (block (result f64)
   ;; CHECK-NEXT:   (drop
-  ;; CHECK-NEXT:    (local.get $ref)
+  ;; CHECK-NEXT:    (ref.null $struct.A)
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:   (local.get $2)
   ;; CHECK-NEXT:  )
@@ -480,13 +480,13 @@
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (block (result (ref null $struct.A))
   ;; CHECK-NEXT:    (block (result (ref null $struct.A))
-  ;; CHECK-NEXT:     (local.get $ref)
+  ;; CHECK-NEXT:     (ref.null $struct.A)
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (block (result f64)
   ;; CHECK-NEXT:   (drop
-  ;; CHECK-NEXT:    (local.get $ref)
+  ;; CHECK-NEXT:    (ref.null $struct.A)
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:   (local.get $2)
   ;; CHECK-NEXT:  )
@@ -609,11 +609,11 @@
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (drop
-  ;; CHECK-NEXT:   (local.get $ref)
+  ;; CHECK-NEXT:   (ref.null $struct.A)
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (block (result f64)
   ;; CHECK-NEXT:   (drop
-  ;; CHECK-NEXT:    (local.get $ref)
+  ;; CHECK-NEXT:    (ref.null $struct.A)
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:   (local.get $2)
   ;; CHECK-NEXT:  )
@@ -654,12 +654,12 @@
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (drop
-  ;; CHECK-NEXT:   (local.get $ref)
+  ;; CHECK-NEXT:   (ref.null $struct.A)
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (block (result i32)
   ;; CHECK-NEXT:    (drop
-  ;; CHECK-NEXT:     (local.get $ref)
+  ;; CHECK-NEXT:     (ref.null $struct.A)
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:    (local.get $2)
   ;; CHECK-NEXT:   )
@@ -667,7 +667,7 @@
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (block (result f64)
   ;; CHECK-NEXT:    (drop
-  ;; CHECK-NEXT:     (local.get $ref-2)
+  ;; CHECK-NEXT:     (ref.null $struct.A)
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:    (local.get $3)
   ;; CHECK-NEXT:   )
@@ -719,12 +719,12 @@
   ;; CHECK-NEXT:  (if
   ;; CHECK-NEXT:   (local.get $x)
   ;; CHECK-NEXT:   (drop
-  ;; CHECK-NEXT:    (local.get $ref)
+  ;; CHECK-NEXT:    (ref.null $struct.A)
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (block (result f64)
   ;; CHECK-NEXT:   (drop
-  ;; CHECK-NEXT:    (local.get $ref)
+  ;; CHECK-NEXT:    (ref.null $struct.A)
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:   (local.get $3)
   ;; CHECK-NEXT:  )
@@ -773,7 +773,7 @@
   ;; CHECK-NEXT:     (call $send-ref
   ;; CHECK-NEXT:      (ref.null $struct.A)
   ;; CHECK-NEXT:     )
-  ;; CHECK-NEXT:     (local.get $ref)
+  ;; CHECK-NEXT:     (ref.null $struct.A)
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:   (local.get $2)
@@ -945,7 +945,7 @@
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (block (result (ref null $struct.recursive))
   ;; CHECK-NEXT:    (drop
-  ;; CHECK-NEXT:     (local.get $0)
+  ;; CHECK-NEXT:     (ref.null $struct.recursive)
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:    (local.get $1)
   ;; CHECK-NEXT:   )
@@ -1102,7 +1102,7 @@
   ;; CHECK-NEXT:   (drop
   ;; CHECK-NEXT:    (block (result i32)
   ;; CHECK-NEXT:     (drop
-  ;; CHECK-NEXT:      (local.get $ref)
+  ;; CHECK-NEXT:      (ref.null $struct.A)
   ;; CHECK-NEXT:     )
   ;; CHECK-NEXT:     (local.get $2)
   ;; CHECK-NEXT:    )
@@ -1112,14 +1112,14 @@
   ;; CHECK-NEXT:    (drop
   ;; CHECK-NEXT:     (block (result f64)
   ;; CHECK-NEXT:      (drop
-  ;; CHECK-NEXT:       (local.get $ref)
+  ;; CHECK-NEXT:       (ref.null $struct.A)
   ;; CHECK-NEXT:      )
   ;; CHECK-NEXT:      (local.get $3)
   ;; CHECK-NEXT:     )
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:    (block
   ;; CHECK-NEXT:     (drop
-  ;; CHECK-NEXT:      (local.get $ref)
+  ;; CHECK-NEXT:      (ref.null $struct.A)
   ;; CHECK-NEXT:     )
   ;; CHECK-NEXT:     (local.set $3
   ;; CHECK-NEXT:      (f64.const 42)
@@ -1129,13 +1129,13 @@
   ;; CHECK-NEXT:   (loop $inner
   ;; CHECK-NEXT:    (block
   ;; CHECK-NEXT:     (drop
-  ;; CHECK-NEXT:      (local.get $ref)
+  ;; CHECK-NEXT:      (ref.null $struct.A)
   ;; CHECK-NEXT:     )
   ;; CHECK-NEXT:     (local.set $2
   ;; CHECK-NEXT:      (i32.add
   ;; CHECK-NEXT:       (block (result i32)
   ;; CHECK-NEXT:        (drop
-  ;; CHECK-NEXT:         (local.get $ref)
+  ;; CHECK-NEXT:         (ref.null $struct.A)
   ;; CHECK-NEXT:        )
   ;; CHECK-NEXT:        (local.get $2)
   ;; CHECK-NEXT:       )
@@ -1152,7 +1152,7 @@
   ;; CHECK-NEXT:    (drop
   ;; CHECK-NEXT:     (block (result i32)
   ;; CHECK-NEXT:      (drop
-  ;; CHECK-NEXT:       (local.get $ref)
+  ;; CHECK-NEXT:       (ref.null $struct.A)
   ;; CHECK-NEXT:      )
   ;; CHECK-NEXT:      (local.get $2)
   ;; CHECK-NEXT:     )
@@ -1160,7 +1160,7 @@
   ;; CHECK-NEXT:    (drop
   ;; CHECK-NEXT:     (block (result f64)
   ;; CHECK-NEXT:      (drop
-  ;; CHECK-NEXT:       (local.get $ref)
+  ;; CHECK-NEXT:       (ref.null $struct.A)
   ;; CHECK-NEXT:      )
   ;; CHECK-NEXT:      (local.get $3)
   ;; CHECK-NEXT:     )
@@ -1340,7 +1340,7 @@
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (block (result i32)
   ;; CHECK-NEXT:    (drop
-  ;; CHECK-NEXT:     (local.get $ref)
+  ;; CHECK-NEXT:     (ref.null $struct.A)
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:    (local.get $1)
   ;; CHECK-NEXT:   )
@@ -1362,7 +1362,7 @@
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (block (result i32)
   ;; CHECK-NEXT:    (drop
-  ;; CHECK-NEXT:     (local.get $ref)
+  ;; CHECK-NEXT:     (ref.null $struct.A)
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:    (local.get $3)
   ;; CHECK-NEXT:   )
@@ -1432,7 +1432,7 @@
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (block (result i32)
   ;; CHECK-NEXT:    (drop
-  ;; CHECK-NEXT:     (local.get $ref1)
+  ;; CHECK-NEXT:     (ref.null $struct.A)
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:    (local.get $2)
   ;; CHECK-NEXT:   )
@@ -1440,7 +1440,7 @@
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (block (result i32)
   ;; CHECK-NEXT:    (drop
-  ;; CHECK-NEXT:     (local.get $ref2)
+  ;; CHECK-NEXT:     (ref.null $struct.A)
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:    (local.get $4)
   ;; CHECK-NEXT:   )
@@ -1580,7 +1580,7 @@
   ;; CHECK-NEXT:    (block $block (result (ref null $struct.A))
   ;; CHECK-NEXT:     (drop
   ;; CHECK-NEXT:      (br_if $block
-  ;; CHECK-NEXT:       (local.get $0)
+  ;; CHECK-NEXT:       (ref.null $struct.A)
   ;; CHECK-NEXT:       (i32.const 0)
   ;; CHECK-NEXT:      )
   ;; CHECK-NEXT:     )
@@ -1790,7 +1790,7 @@
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (block
   ;; CHECK-NEXT:   (drop
-  ;; CHECK-NEXT:    (local.get $ref)
+  ;; CHECK-NEXT:    (ref.null $struct.A)
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:   (local.set $1
   ;; CHECK-NEXT:    (i32.const 1)
@@ -1845,11 +1845,11 @@
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (drop
-  ;; CHECK-NEXT:   (local.get $ref)
+  ;; CHECK-NEXT:   (ref.null $struct.A)
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (block (result i32)
   ;; CHECK-NEXT:   (drop
-  ;; CHECK-NEXT:    (local.get $ref)
+  ;; CHECK-NEXT:    (ref.null $struct.A)
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:   (local.get $1)
   ;; CHECK-NEXT:  )