[Wasm GC] DeadArgumentElimination: Update tees after refining param types (#4031)

3 files changed, 51 insertions, 3 deletions

diff --git a/src/passes/DeadArgumentElimination.cpp b/src/passes/DeadArgumentElimination.cpp
index 7ba1be9f8..9d8fa5b0d 100644
--- a/src/passes/DeadArgumentElimination.cpp
+++ b/src/passes/DeadArgumentElimination.cpp
@@ -588,7 +588,8 @@ private:
     // In terms of parameters, we can do this. However, we must also check
     // local operations in the body, as if the parameter is reused and written
     // to, then those types must be taken into account as well.
-    for (auto* set : FindAll<LocalSet>(func->body).list) {
+    FindAll<LocalSet> sets(func->body);
+    for (auto* set : sets.list) {
       auto index = set->index;
       if (func->isParam(index) &&
           !Type::isSubType(set->value->type, newParamTypes[index])) {
@@ -602,7 +603,7 @@ private:
       return;
     }
 
-    // We can do this! Update the types, including the types of gets.
+    // We can do this! Update the types, including the types of gets and tees.
     func->setParams(newParams);
     for (auto* get : FindAll<LocalGet>(func->body).list) {
       auto index = get->index;
@@ -610,6 +611,16 @@ private:
         get->type = func->getLocalType(index);
       }
     }
+    for (auto* set : sets.list) {
+      auto index = set->index;
+      if (func->isParam(index) && set->isTee()) {
+        set->type = func->getLocalType(index);
+        set->finalize();
+      }
+    }
+
+    // Propagate the new get and set types outwards.
+    ReFinalize().walkFunctionInModule(func, module);
   }
 
   // See if the types returned from a function allow us to define a more refined
diff --git a/src/passes/LocalSubtyping.cpp b/src/passes/LocalSubtyping.cpp
index 9c6da1131..472463f44 100644
--- a/src/passes/LocalSubtyping.cpp
+++ b/src/passes/LocalSubtyping.cpp
@@ -134,8 +134,9 @@ struct LocalSubtyping : public WalkerPass<PostWalker<LocalSubtyping>> {
           // NB: These tee updates will not be needed if the type of tees
           //     becomes that of their value, in the spec.
           for (auto* set : setsForLocal[i]) {
-            if (set->isTee() && set->type != Type::unreachable) {
+            if (set->isTee()) {
               set->type = newType;
+              set->finalize();
             }
           }
         }
diff --git a/test/lit/passes/dae-gc.wast b/test/lit/passes/dae-gc.wast
index 230f757fc..0ea322ce8 100644
--- a/test/lit/passes/dae-gc.wast
+++ b/test/lit/passes/dae-gc.wast
@@ -213,6 +213,42 @@
   (local.set $y (ref.null ${i32_i64}))
  )
 
+ ;; CHECK:      (func $call-various-params-tee
+ ;; CHECK-NEXT:  (call $various-params-tee
+ ;; CHECK-NEXT:   (ref.null ${i32})
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $call-various-params-tee
+  ;; The argument gets {i32}, which allows us to refine.
+  (call $various-params-tee
+   (ref.null ${i32})
+  )
+ )
+ ;; CHECK:      (func $various-params-tee (param $x (ref null ${i32}))
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (local.get $x)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (block $block (result (ref null ${i32}))
+ ;; CHECK-NEXT:    (local.tee $x
+ ;; CHECK-NEXT:     (ref.null ${i32_i64})
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $various-params-tee (param $x (ref null ${}))
+  ;; "Use" the locals to avoid other optimizations kicking in.
+  (drop (local.get $x))
+  ;; Write to $x in a way that allows us to make the type more specific. We
+  ;; must also update the type of the tee (if we do not, a validation error
+  ;; would occur), and that will also cause the block's type to update as well.
+  (drop
+   (block (result (ref null ${}))
+    (local.tee $x (ref.null ${i32_i64}))
+   )
+  )
+ )
+
  ;; CHECK:      (func $call-various-params-null
  ;; CHECK-NEXT:  (call $various-params-null
  ;; CHECK-NEXT:   (ref.as_non_null