OptimizeInstructions: Refinalize after a cast removal (#4611)

Casts can replace a type with a subtype, which normally has no downsides, but
in a corner case of struct types it can lead to us needing to refinalize higher up
too, see details in the comment.

We have avoided any Refinalize calls in OptimizeInstructions, but the case
handled here requires it sadly. I considered moving it to another pass, but this
is a peephole optimization so there isn't really a better place.

2 files changed, 81 insertions, 3 deletions

diff --git a/src/passes/OptimizeInstructions.cpp b/src/passes/OptimizeInstructions.cpp
index bde60e093..4b499213f 100644
--- a/src/passes/OptimizeInstructions.cpp
+++ b/src/passes/OptimizeInstructions.cpp
@@ -208,6 +208,9 @@ struct OptimizeInstructions
 
   bool fastMath;
 
+  // In rare cases we make a change to a type, and will do a refinalize.
+  bool refinalize = false;
+
   void doWalkFunction(Function* func) {
     fastMath = getPassOptions().fastMath;
 
@@ -221,6 +224,10 @@ struct OptimizeInstructions
     // Main walk.
     super::doWalkFunction(func);
 
+    if (refinalize) {
+      ReFinalize().walkFunctionInModule(func, getModule());
+    }
+
     // Final optimizations.
     {
       FinalOptimizer optimizer(getPassOptions());
@@ -1622,6 +1629,25 @@ struct OptimizeInstructions
                                           passOptions));
         } else {
           replaceCurrent(curr->ref);
+
+          // We must refinalize here, as we may be returning a more specific
+          // type, which can alter the parent. For example:
+          //
+          //  (struct.get $parent 0
+          //   (ref.cast_static $parent
+          //    (local.get $child)
+          //   )
+          //  )
+          //
+          // Try to cast a $child to its parent, $parent. That always works,
+          // so the cast can be removed.
+          // Then once the cast is removed, the outer struct.get
+          // will have a reference with a different type, making it a
+          // (struct.get $child ..) instead of $parent.
+          // But if $parent and $child have different types on field 0 (the
+          // child may have a more refined one) then the struct.get must be
+          // refinalized so the IR node has the expected type.
+          refinalize = true;
         }
         return;
       }
diff --git a/test/lit/passes/optimize-instructions-gc-iit.wast b/test/lit/passes/optimize-instructions-gc-iit.wast
index 37bc6cb69..d0638860f 100644
--- a/test/lit/passes/optimize-instructions-gc-iit.wast
+++ b/test/lit/passes/optimize-instructions-gc-iit.wast
@@ -1,10 +1,10 @@
 ;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
-;; RUN: wasm-opt %s --optimize-instructions --ignore-implicit-traps --enable-reference-types --enable-gc -S -o - \
+;; RUN: foreach %s %t wasm-opt --optimize-instructions --ignore-implicit-traps --enable-reference-types --enable-gc -S -o - \
 ;; RUN:   | filecheck %s
-;; RUN: wasm-opt %s --optimize-instructions --ignore-implicit-traps --enable-reference-types --enable-gc --nominal -S -o - \
+;; RUN: foreach %s %t wasm-opt --optimize-instructions --ignore-implicit-traps --enable-reference-types --enable-gc --nominal -S -o - \
 ;; RUN:   | filecheck %s --check-prefix NOMNL
 ;; Also test trapsNeverHappen (with nominal; no need for both type system modes).
-;; RUN: wasm-opt %s --optimize-instructions --traps-never-happen --enable-reference-types --enable-gc --nominal -S -o - \
+;; RUN: foreach %s %t wasm-opt --optimize-instructions --traps-never-happen --enable-reference-types --enable-gc --nominal -S -o - \
 ;; RUN:   | filecheck %s --check-prefix NOMNL-TNH
 
 (module
@@ -367,3 +367,55 @@
     )
   )
 )
+
+(module
+  ;; CHECK:      (type $B (struct (field (ref null $A))))
+
+  ;; CHECK:      (type $A (struct ))
+  ;; NOMNL:      (type $C (struct_subtype (field (ref null $D)) $B))
+
+  ;; NOMNL:      (type $D (struct_subtype  $A))
+
+  ;; NOMNL:      (type $A (struct_subtype  data))
+  ;; NOMNL-TNH:      (type $C (struct_subtype (field (ref null $D)) $B))
+
+  ;; NOMNL-TNH:      (type $D (struct_subtype  $A))
+
+  ;; NOMNL-TNH:      (type $A (struct_subtype  data))
+  (type $A (struct_subtype  data))
+  ;; NOMNL:      (type $B (struct_subtype (field (ref null $A)) $A))
+  ;; NOMNL-TNH:      (type $B (struct_subtype (field (ref null $A)) $A))
+  (type $B (struct_subtype (field (ref null $A)) $A))
+  (type $C (struct_subtype (field (ref null $D)) $B))
+  (type $D (struct_subtype  $A))
+
+  ;; CHECK:      (func $test (param $C (ref $B)) (result anyref)
+  ;; CHECK-NEXT:  (struct.get $B 0
+  ;; CHECK-NEXT:   (local.get $C)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  ;; NOMNL:      (func $test (type $ref|$C|_=>_anyref) (param $C (ref $C)) (result anyref)
+  ;; NOMNL-NEXT:  (struct.get $C 0
+  ;; NOMNL-NEXT:   (local.get $C)
+  ;; NOMNL-NEXT:  )
+  ;; NOMNL-NEXT: )
+  ;; NOMNL-TNH:      (func $test (type $ref|$C|_=>_anyref) (param $C (ref $C)) (result anyref)
+  ;; NOMNL-TNH-NEXT:  (struct.get $C 0
+  ;; NOMNL-TNH-NEXT:   (local.get $C)
+  ;; NOMNL-TNH-NEXT:  )
+  ;; NOMNL-TNH-NEXT: )
+  (func $test (param $C (ref $C)) (result anyref)
+    (struct.get $B 0
+      (ref.cast_static $B ;; Try to cast a $C to its parent, $B. That always
+                          ;; works, so the cast can be removed.
+                          ;; Then once the cast is removed, the outer struct.get
+                          ;; will have a reference with a different type,
+                          ;; making it a (struct.get $C ..) instead of $B.
+                          ;; But $B and $C have different types on field 0, and
+                          ;; so the struct.get must be refinalized so the node
+                          ;; has the expected type.
+        (local.get $C)
+      )
+    )
+  )
+)