1 files changed, 8 insertions, 2 deletions
diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index ac6fb49a6..fc3c18db0 100644
--- a/src/passes/Heap2Local.cpp
+++ b/src/passes/Heap2Local.cpp
@@ -274,12 +274,14 @@ struct Heap2LocalOptimizer {
     // left to other passes, like getting rid of dropped code without side
     // effects.
 
-    void visitBlock(Block* curr) {
+    // Adjust the type that flows through an expression, updating that type as
+    // necessary.
+    void adjustTypeFlowingThrough(Expression* curr) {
       if (!reached.count(curr)) {
         return;
       }
 
-      // Our allocation passes through this block. We must turn its type into a
+      // Our allocation passes through this expr. 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
@@ -290,6 +292,10 @@ struct Heap2LocalOptimizer {
       curr->type = Type(curr->type.getHeapType(), Nullable);
     }
 
+    void visitBlock(Block* curr) { adjustTypeFlowingThrough(curr); }
+
+    void visitLoop(Loop* curr) { adjustTypeFlowingThrough(curr); }
+
     void visitLocalSet(LocalSet* curr) {
       if (!reached.count(curr)) {
         return;