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;