[Wasm GC] Fix TypeRefining on fallthrough values via tee (#4900)

A rather tricky corner case: we normally look at fallthrough values for copies of
fields, so when we try to refine a field, we ignore stuff like this:

a.x = b.x;

That copies the same field on the same type to itself, so refining is not limited by
it. But if we have something else in the middle, and that thing cannot change
type, then it is a problem, like this:

(struct.set
  (..ref..)
  (local.tee $temp
    (struct.get)))

tee has the type of the local, which does not change in this pass. So we can't
look at just the fallthrough here and skip the tee: after refining the field, the
tee's old type might not fit in the field's new type.

We could perhaps add casts to fix things up, but those may have too big a
cost. For now, just ignore the fallthrough.

4 files changed, 57 insertions, 10 deletions

diff --git a/src/ir/properties.h b/src/ir/properties.h
index 4bd66b1d5..1d2937d81 100644
--- a/src/ir/properties.h
+++ b/src/ir/properties.h
@@ -240,17 +240,29 @@ inline Index getZeroExtBits(Expression* curr) {
 // way to the final value falling through, potentially through multiple
 // intermediate expressions.
 //
+// Behavior wrt tee/br_if is customizable, since in some cases we do not want to
+// look through them (for example, the type of a tee is related to the local,
+// not the value, so if we returned the fallthrough of the tee we'd have a
+// possible difference between the type in the IR and the type of the value,
+// which some cases care about; the same for a br_if, whose type is related to
+// the branch target).
+//
 // TODO: Receive a Module instead of FeatureSet, to pass to EffectAnalyzer?
-inline Expression* getImmediateFallthrough(Expression* curr,
-                                           const PassOptions& passOptions,
-                                           Module& module) {
+
+enum class FallthroughBehavior { AllowTeeBrIf, NoTeeBrIf };
+
+inline Expression* getImmediateFallthrough(
+  Expression* curr,
+  const PassOptions& passOptions,
+  Module& module,
+  FallthroughBehavior behavior = FallthroughBehavior::AllowTeeBrIf) {
   // If the current node is unreachable, there is no value
   // falling through.
   if (curr->type == Type::unreachable) {
     return curr;
   }
   if (auto* set = curr->dynCast<LocalSet>()) {
-    if (set->isTee()) {
+    if (set->isTee() && behavior == FallthroughBehavior::AllowTeeBrIf) {
       return set->value;
     }
   } else if (auto* block = curr->dynCast<Block>()) {
@@ -270,7 +282,8 @@ inline Expression* getImmediateFallthrough(Expression* curr,
       }
     }
   } else if (auto* br = curr->dynCast<Break>()) {
-    if (br->condition && br->value) {
+    if (br->condition && br->value &&
+        behavior == FallthroughBehavior::AllowTeeBrIf) {
       return br->value;
     }
   } else if (auto* tryy = curr->dynCast<Try>()) {
@@ -289,11 +302,13 @@ inline Expression* getImmediateFallthrough(Expression* curr,
 
 // Similar to getImmediateFallthrough, but looks through multiple children to
 // find the final value that falls through.
-inline Expression* getFallthrough(Expression* curr,
-                                  const PassOptions& passOptions,
-                                  Module& module) {
+inline Expression* getFallthrough(
+  Expression* curr,
+  const PassOptions& passOptions,
+  Module& module,
+  FallthroughBehavior behavior = FallthroughBehavior::AllowTeeBrIf) {
   while (1) {
-    auto* next = getImmediateFallthrough(curr, passOptions, module);
+    auto* next = getImmediateFallthrough(curr, passOptions, module, behavior);
     if (next == curr) {
       return curr;
     }
diff --git a/src/ir/struct-utils.h b/src/ir/struct-utils.h
index ba5617de6..9d02bb779 100644
--- a/src/ir/struct-utils.h
+++ b/src/ir/struct-utils.h
@@ -191,7 +191,10 @@ struct StructScanner
     // (otherwise, we'd need to consider both the type actually written and the
     // type of the fallthrough, somehow).
     auto* fallthrough = Properties::getFallthrough(
-      expr, this->getPassOptions(), *this->getModule());
+      expr,
+      this->getPassOptions(),
+      *this->getModule(),
+      static_cast<SubType*>(this)->getFallthroughBehavior());
     if (fallthrough->type == expr->type) {
       expr = fallthrough;
     }
@@ -205,6 +208,11 @@ struct StructScanner
     static_cast<SubType*>(this)->noteExpression(expr, type, index, info);
   }
 
+  Properties::FallthroughBehavior getFallthroughBehavior() {
+    // By default, look at and use tee&br_if fallthrough values.
+    return Properties::FallthroughBehavior::AllowTeeBrIf;
+  }
+
   FunctionStructValuesMap<T>& functionNewInfos;
   FunctionStructValuesMap<T>& functionSetGetInfos;
 };
diff --git a/src/passes/TypeRefining.cpp b/src/passes/TypeRefining.cpp
index 258feae32..ac95d01fe 100644
--- a/src/passes/TypeRefining.cpp
+++ b/src/passes/TypeRefining.cpp
@@ -72,6 +72,24 @@ struct FieldInfoScanner
   void noteRead(HeapType type, Index index, FieldInfo& info) {
     // Nothing to do for a read, we just care about written values.
   }
+
+  Properties::FallthroughBehavior getFallthroughBehavior() {
+    // Looking at fallthrough values may be dangerous here, because it ignores
+    // intermediate steps. Consider this:
+    //
+    // (struct.set $T 0
+    //   (local.tee
+    //     (struct.get $T 0)))
+    //
+    // This is a copy of a field to itself - normally something we can ignore
+    // (see above). But in this case, if we refine the field then that will
+    // update the struct.get, but the local.tee will still have the old type,
+    // which may not be refined enough. We could in theory always fix this up
+    // using casts later, but those casts may be expensive (especially ref.casts
+    // as opposed to ref.as_non_null), so for now just ignore tee fallthroughs.
+    // TODO: investigate more
+    return Properties::FallthroughBehavior::NoTeeBrIf;
+  }
 };
 
 struct TypeRefining : public Pass {
diff --git a/src/wasm/wasm.cpp b/src/wasm/wasm.cpp
index 574eb5c47..20d1b643e 100644
--- a/src/wasm/wasm.cpp
+++ b/src/wasm/wasm.cpp
@@ -259,6 +259,12 @@ void Break::finalize() {
     if (condition->type == Type::unreachable) {
       type = Type::unreachable;
     } else if (value) {
+      // N.B. This is not correct wrt the spec, which mandates that it be the
+      // type of the block we target. In practice this does not matter because
+      // the br_if return value is not really used in the wild. To fix this,
+      // we'd need to do something like what we do for local.tee's type, which
+      // is to fix it up in a way that is aware of function-level context and
+      // not just the instruction itself (which would be a pain).
       type = value->type;
     } else {
       type = Type::none;