Heap2Local: Refinalize if we end up refining (#5879)

We shouldn't need to in the general case, but the fuzzer found a corner case
where we do need to, see the explanation + testcase, but basically Heap2Local
replaces struct fields with locals, and the locals should have the same types,
but if a field was somehow less refined for some reason, then the locals could
actually be more refined. (And a field could be less refined if we read it from a
typed that was under-refined due to a tee or such.)

2 files changed, 79 insertions, 6 deletions

diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index 7e2b02add..84037e0dd 100644
--- a/src/passes/Heap2Local.cpp
+++ b/src/passes/Heap2Local.cpp
@@ -233,11 +233,12 @@ struct Heap2LocalOptimizer {
   struct Rewriter : PostWalker<Rewriter> {
     StructNew* allocation;
     Function* func;
+    Module* module;
     Builder builder;
     const FieldList& fields;
 
     Rewriter(StructNew* allocation, Function* func, Module* module)
-      : allocation(allocation), func(func), builder(*module),
+      : allocation(allocation), func(func), module(module), builder(*module),
         fields(allocation->type.getHeapType().getStruct().fields) {}
 
     // We must track all the local.sets that write the allocation, to verify
@@ -252,6 +253,9 @@ struct Heap2LocalOptimizer {
     // Maps indexes in the struct to the local index that will replace them.
     std::vector<Index> localIndexes;
 
+    // In rare cases we may need to refinalize, see below.
+    bool refinalize = false;
+
     void applyOptimization() {
       // Allocate locals to store the allocation's fields in.
       for (auto field : fields) {
@@ -260,6 +264,10 @@ struct Heap2LocalOptimizer {
 
       // Replace the things we need to using the visit* methods.
       walk(func->body);
+
+      if (refinalize) {
+        ReFinalize().walkFunctionInModule(func, module);
+      }
     }
 
     // Rewrite the code in visit* methods. The general approach taken is to
@@ -439,10 +447,25 @@ struct Heap2LocalOptimizer {
         return;
       }
 
-      replaceCurrent(
-        builder.makeSequence(builder.makeDrop(curr->ref),
-                             builder.makeLocalGet(localIndexes[curr->index],
-                                                  fields[curr->index].type)));
+      auto type = fields[curr->index].type;
+      if (type != curr->type) {
+        // Normally we are just replacing a struct.get with a local.get of a
+        // local that was created to have the same type as the struct's field,
+        // but in some cases we may refine, if the struct.get's reference type
+        // is less refined than the reference that actually arrives, like here:
+        //
+        //  (struct.get $parent 0
+        //    (block (ref $parent)
+        //      (struct.new $child)))
+        //
+        // We allocated locals for the field of the child, and are replacing a
+        // get of the parent field with a local of the same type as the child's,
+        // which may be more refined.
+        refinalize = true;
+      }
+      replaceCurrent(builder.makeSequence(
+        builder.makeDrop(curr->ref),
+        builder.makeLocalGet(localIndexes[curr->index], type)));
     }
   };
 
diff --git a/test/lit/passes/heap2local.wast b/test/lit/passes/heap2local.wast
index ca6439a8a..090ab4de0 100644
--- a/test/lit/passes/heap2local.wast
+++ b/test/lit/passes/heap2local.wast
@@ -1,7 +1,7 @@
 ;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
 
 ;; (remove-unused-names allows the pass to see that blocks flow values)
-;; RUN: wasm-opt %s -all --remove-unused-names --heap2local -S -o - | filecheck %s
+;; 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))))
@@ -1896,3 +1896,53 @@
     )
   )
 )
+
+(module
+  ;; CHECK:      (type $A (struct (field (ref null $A))))
+  (type $A (struct (field (ref null $A))))
+  ;; CHECK:      (type $B (sub $A (struct (field (ref $A)))))
+  (type $B (sub $A (struct (field (ref $A)))))
+
+  ;; CHECK:      (func $func (type $none_=>_anyref) (result anyref)
+  ;; CHECK-NEXT:  (local $a (ref $A))
+  ;; CHECK-NEXT:  (local $1 (ref $A))
+  ;; CHECK-NEXT:  (local $2 (ref $A))
+  ;; CHECK-NEXT:  (ref.cast $B
+  ;; CHECK-NEXT:   (block (result (ref $A))
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (block (result nullref)
+  ;; CHECK-NEXT:      (local.set $2
+  ;; CHECK-NEXT:       (struct.new $A
+  ;; CHECK-NEXT:        (ref.null none)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (local.set $1
+  ;; CHECK-NEXT:       (local.get $2)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (ref.null none)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.get $1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $func (result anyref)
+    (local $a (ref $A))
+    ;; Refinalization will be needed here, as a struct.new of $B can be
+    ;; optimized, and that reference flows through a tee, which loses type
+    ;; precision, into a local.get of $A. After heap2local we'll end up using a
+    ;; local of the type of $B's field which is more precise than $A's, and the
+    ;; cast must be updated to be non-nullable.
+    (ref.cast null $B
+      (struct.get $A 0
+        (local.tee $a
+          (struct.new $B
+            (struct.new $A
+              (ref.null none)
+            )
+          )
+        )
+      )
+    )
+  )
+)