SimplifyLocals: ReFinalize when needed (#4878)

2 files changed, 63 insertions, 0 deletions

diff --git a/src/passes/SimplifyLocals.cpp b/src/passes/SimplifyLocals.cpp
index 159f09963..ee9e43b0b 100644
--- a/src/passes/SimplifyLocals.cpp
+++ b/src/passes/SimplifyLocals.cpp
@@ -53,6 +53,7 @@
 #include <ir/linear-execution.h>
 #include <ir/local-utils.h>
 #include <ir/manipulation.h>
+#include <ir/utils.h>
 #include <pass.h>
 #include <wasm-builder.h>
 #include <wasm-traversal.h>
@@ -120,6 +121,9 @@ struct SimplifyLocals
   // local => # of local.gets for it
   LocalGetCounter getCounter;
 
+  // In rare cases we make a change to a type that requires a refinalize.
+  bool refinalize = false;
+
   static void
   doNoteNonLinear(SimplifyLocals<allowTee, allowStructure, allowNesting>* self,
                   Expression** currp) {
@@ -254,6 +258,23 @@ struct SimplifyLocals
       if (oneUse) {
         // with just one use, we can sink just the value
         this->replaceCurrent(set->value);
+
+        // We are replacing a local.get with the value of the local.set. That
+        // may require a refinalize in certain cases, like this:
+        //
+        //  (struct.get $X 0
+        //    (local.get $x)
+        //  )
+        //
+        // If we replace the local.get with a more refined type then the
+        // struct.get may read a more refined type (if the subtype has a more
+        // refined type for that particular field). Note that this cannot happen
+        // in the other arm of this if-else, where we replace the local.get with
+        // a tee, since tees have the type of the local, so no types change
+        // there.
+        if (set->value->type != curr->type) {
+          refinalize = true;
+        }
       } else {
         this->replaceCurrent(set);
         assert(!set->isTee());
@@ -889,6 +910,10 @@ struct SimplifyLocals
         }
       }
     } while (anotherCycle);
+
+    if (refinalize) {
+      ReFinalize().walkFunctionInModule(func, this->getModule());
+    }
   }
 
   bool runMainOptimizations(Function* func) {
diff --git a/test/lit/passes/simplify-locals-gc.wast b/test/lit/passes/simplify-locals-gc.wast
index 33e4d1bd5..a930cae08 100644
--- a/test/lit/passes/simplify-locals-gc.wast
+++ b/test/lit/passes/simplify-locals-gc.wast
@@ -13,6 +13,17 @@
   ;; NOMNL:      (type $struct-immutable (struct_subtype (field i32) data))
   (type $struct-immutable (struct (field i32)))
 
+  ;; CHECK:      (type $B (struct (field dataref)))
+
+  ;; CHECK:      (type $A (struct (field (ref null data))))
+  ;; NOMNL:      (type $A (struct_subtype (field (ref null data)) data))
+  (type $A (struct_subtype (field (ref null data)) data))
+
+  ;; $B is a subtype of $A, and its field has a more refined type (it is non-
+  ;; nullable).
+  ;; NOMNL:      (type $B (struct_subtype (field dataref) $A))
+  (type $B (struct_subtype (field (ref data)) $A))
+
   ;; Writes to heap objects cannot be reordered with reads.
   ;; CHECK:      (func $no-reorder-past-write (param $x (ref $struct)) (result i32)
   ;; CHECK-NEXT:  (local $temp i32)
@@ -214,4 +225,31 @@
     )
    )
   )
+
+  ;; CHECK:      (func $needs-refinalize (param $b (ref $B)) (result anyref)
+  ;; CHECK-NEXT:  (local $a (ref null $A))
+  ;; CHECK-NEXT:  (nop)
+  ;; CHECK-NEXT:  (struct.get $B 0
+  ;; CHECK-NEXT:   (local.get $b)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  ;; NOMNL:      (func $needs-refinalize (type $ref|$B|_=>_anyref) (param $b (ref $B)) (result anyref)
+  ;; NOMNL-NEXT:  (local $a (ref null $A))
+  ;; NOMNL-NEXT:  (nop)
+  ;; NOMNL-NEXT:  (struct.get $B 0
+  ;; NOMNL-NEXT:   (local.get $b)
+  ;; NOMNL-NEXT:  )
+  ;; NOMNL-NEXT: )
+  (func $needs-refinalize (param $b (ref $B)) (result anyref)
+    (local $a (ref null $A))
+    (local.set $a
+      (local.get $b)
+    )
+    ;; This begins as a struct.get of $A, but after we move the set's value onto
+    ;; the get, we'll be reading from $B. $B's field has a more refined type, so
+    ;; we must update the type of the struct.get using refinalize.
+    (struct.get $A 0
+      (local.get $a)
+    )
+  )
 )