[Wasm GC] Refinalize in UnneededSetRemover when necessary (#5287)

2 files changed, 41 insertions, 6 deletions

diff --git a/src/ir/local-utils.h b/src/ir/local-utils.h
index ef74a35eb..153ebd3f3 100644
--- a/src/ir/local-utils.h
+++ b/src/ir/local-utils.h
@@ -19,6 +19,7 @@
 
 #include <ir/effects.h>
 #include <ir/manipulation.h>
+#include <ir/utils.h>
 
 namespace wasm {
 
@@ -61,9 +62,14 @@ struct UnneededSetRemover : public PostWalker<UnneededSetRemover> {
     : passOptions(passOptions), localGetCounter(&localGetCounter),
       module(module) {
     walk(func->body);
+
+    if (refinalize) {
+      ReFinalize().walkFunctionInModule(func, &module);
+    }
   }
 
   bool removed = false;
+  bool refinalize = false;
 
   void visitLocalSet(LocalSet* curr) {
     // If no possible uses, remove.
@@ -94,6 +100,10 @@ struct UnneededSetRemover : public PostWalker<UnneededSetRemover> {
     auto* value = set->value;
     if (set->isTee()) {
       replaceCurrent(value);
+      if (value->type != set->type) {
+        // The value is more refined, so we'll need to refinalize.
+        refinalize = true;
+      }
     } else if (EffectAnalyzer(passOptions, module, set->value)
                  .hasSideEffects()) {
       Drop* drop = ExpressionManipulator::convert<LocalSet, Drop>(set);
diff --git a/test/lit/passes/simplify-locals-gc.wast b/test/lit/passes/simplify-locals-gc.wast
index 80bbe6d90..35f90387d 100644
--- a/test/lit/passes/simplify-locals-gc.wast
+++ b/test/lit/passes/simplify-locals-gc.wast
@@ -5,19 +5,19 @@
 ;; RUN:   | filecheck %s --check-prefix=NOMNL
 
 (module
+  ;; CHECK:      (type $B (struct (field (ref data))))
+
   ;; CHECK:      (type $struct (struct (field (mut i32))))
+  ;; NOMNL:      (type $A (struct_subtype (field dataref) data))
+
+  ;; NOMNL:      (type $B (struct_subtype (field (ref data)) $A))
+
   ;; NOMNL:      (type $struct (struct_subtype (field (mut i32)) data))
   (type $struct (struct (field (mut i32))))
 
-  ;; CHECK:      (type $B (struct (field (ref data))))
-
   ;; CHECK:      (type $A (struct (field dataref)))
 
   ;; CHECK:      (type $struct-immutable (struct (field i32)))
-  ;; NOMNL:      (type $A (struct_subtype (field dataref) data))
-
-  ;; NOMNL:      (type $B (struct_subtype (field (ref data)) $A))
-
   ;; NOMNL:      (type $struct-immutable (struct_subtype (field i32) data))
   (type $struct-immutable (struct (field i32)))
 
@@ -702,4 +702,29 @@
   (func $use-any (param $any anyref)
     ;; Helper function for the above.
   )
+
+  ;; CHECK:      (func $remove-tee-refinalize (param $a (ref null $A)) (param $b (ref null $B)) (result dataref)
+  ;; CHECK-NEXT:  (struct.get $B 0
+  ;; CHECK-NEXT:   (local.get $b)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  ;; NOMNL:      (func $remove-tee-refinalize (type $ref?|$A|_ref?|$B|_=>_dataref) (param $a (ref null $A)) (param $b (ref null $B)) (result dataref)
+  ;; NOMNL-NEXT:  (struct.get $B 0
+  ;; NOMNL-NEXT:   (local.get $b)
+  ;; NOMNL-NEXT:  )
+  ;; NOMNL-NEXT: )
+  (func $remove-tee-refinalize
+    (param $a (ref null $A))
+    (param $b (ref null $B))
+    (result (ref null data))
+
+    ;; The local.tee receives a $B and flows out an $A. After we remove it (it is
+    ;; obviously unnecessary), the struct.get will be reading from the more
+    ;; refined type $B.
+    (struct.get $A 0
+      (local.tee $a
+        (local.get $b)
+      )
+    )
+  )
 )