[Wasm GC] Fix CoalesceLocals on tees that receive a refined type (#5289)

Same testcase as in #5287 but in another pass.

2 files changed, 73 insertions, 3 deletions

diff --git a/src/passes/CoalesceLocals.cpp b/src/passes/CoalesceLocals.cpp
index 558be07ff..062449e95 100644
--- a/src/passes/CoalesceLocals.cpp
+++ b/src/passes/CoalesceLocals.cpp
@@ -523,7 +523,10 @@ void CoalesceLocals::applyIndices(std::vector<Index>& indices,
           }
         }
         if (auto* subSet = set->value->dynCast<LocalSet>()) {
-          if (subSet->index == set->index) {
+          // Only do so if not only the index matches but also the type. If the
+          // inner type is more refined, leave that for other passes.
+          if (subSet->index == set->index &&
+              subSet->value->type == subSet->type) {
             set->value = subSet->value;
             continue;
           }
@@ -550,12 +553,28 @@ void CoalesceLocals::applyIndices(std::vector<Index>& indices,
         if (!action.effective) {
           // value may have no side effects, further optimizations can eliminate
           // it
-          *action.origin = set->value;
+          auto* value = set->value;
           if (!set->isTee()) {
             // we need to drop it
             Drop* drop = ExpressionManipulator::convert<LocalSet, Drop>(set);
-            drop->value = *action.origin;
+            drop->value = value;
             *action.origin = drop;
+          } else {
+            // This is a tee, and so, as earlier in this function, we must be
+            // careful of subtyping. Above we simply avoided the problem by
+            // leaving it for other passes, but we do want to remove ineffective
+            // stores - nothing else does that as well as this pass. Instead,
+            // create a block to cast back to the original type, which avoids
+            // changing types here, and leave it to other passes to refine types
+            // and remove the block.
+            auto originalType = (*action.origin)->type;
+            if (originalType != set->value->type) {
+              (*action.origin) =
+                Builder(*getModule()).makeBlock({set->value}, originalType);
+            } else {
+              // No special handling, just use the value.
+              *action.origin = set->value;
+            }
           }
           continue;
         }
diff --git a/test/lit/passes/coalesce-locals-gc.wast b/test/lit/passes/coalesce-locals-gc.wast
index 5947aae11..f8694ce86 100644
--- a/test/lit/passes/coalesce-locals-gc.wast
+++ b/test/lit/passes/coalesce-locals-gc.wast
@@ -8,8 +8,16 @@
 ;; testcases.
 
 (module
+ ;; CHECK:      (type $A (struct (field dataref)))
+
  ;; CHECK:      (type $array (array (mut i8)))
  (type $array (array (mut i8)))
+
+ (type $A (struct_subtype (field (ref null data)) data))
+
+ ;; CHECK:      (type $B (struct (field (ref data))))
+ (type $B (struct_subtype (field (ref data)) $A))
+
  ;; CHECK:      (global $global (ref null $array) (ref.null none))
  (global $global (ref null $array) (ref.null $array))
 
@@ -165,4 +173,47 @@
    (local.get $null-i31)
   )
  )
+
+ ;; CHECK:      (func $remove-tee-refinalize (param $0 (ref null $A)) (param $1 (ref null $B)) (result dataref)
+ ;; CHECK-NEXT:  (struct.get $A 0
+ ;; CHECK-NEXT:   (block (result (ref null $A))
+ ;; CHECK-NEXT:    (local.get $1)
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-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. We want to avoid changing
+  ;; types here, so we'll wrap it in a block, and leave further improvements
+  ;; for other passes.
+  (struct.get $A 0
+   (local.tee $a
+    (local.get $b)
+   )
+  )
+ )
+
+ ;; CHECK:      (func $remove-tee-refinalize-2 (param $0 (ref null $A)) (param $1 (ref null $B)) (result dataref)
+ ;; CHECK-NEXT:  (struct.get $A 0
+ ;; CHECK-NEXT:   (block (result (ref null $A))
+ ;; CHECK-NEXT:    (local.get $1)
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $remove-tee-refinalize-2
+  (param $a (ref null $A))
+  (param $b (ref null $B))
+  (result (ref null data))
+  ;; As above, but with an extra tee in the middle. The result should be the
+  ;; same.
+  (struct.get $A 0
+   (local.tee $a
+    (local.tee $a
+     (local.get $b)
+    )
+   )
+  )
+ )
 )