3 files changed, 35 insertions, 11 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);