2 files changed, 56 insertions, 26 deletions

diff --git a/src/ir/properties.h b/src/ir/properties.h
index cddcbc5f2..e7e96507c 100644
--- a/src/ir/properties.h
+++ b/src/ir/properties.h
@@ -255,34 +255,35 @@ inline Index getZeroExtBits(Expression* curr) {
 
 enum class FallthroughBehavior { AllowTeeBrIf, NoTeeBrIf };
 
-inline Expression* getImmediateFallthrough(
-  Expression* curr,
+inline Expression** getImmediateFallthroughPtr(
+  Expression** currp,
   const PassOptions& passOptions,
   Module& module,
   FallthroughBehavior behavior = FallthroughBehavior::AllowTeeBrIf) {
+  auto* curr = *currp;
   // If the current node is unreachable, there is no value
   // falling through.
   if (curr->type == Type::unreachable) {
-    return curr;
+    return currp;
   }
   if (auto* set = curr->dynCast<LocalSet>()) {
     if (set->isTee() && behavior == FallthroughBehavior::AllowTeeBrIf) {
-      return set->value;
+      return &set->value;
     }
   } else if (auto* block = curr->dynCast<Block>()) {
     // if no name, we can't be broken to, and then can look at the fallthrough
     if (!block->name.is() && block->list.size() > 0) {
-      return block->list.back();
+      return &block->list.back();
     }
   } else if (auto* loop = curr->dynCast<Loop>()) {
-    return loop->body;
+    return &loop->body;
   } else if (auto* iff = curr->dynCast<If>()) {
     if (iff->ifFalse) {
       // Perhaps just one of the two actually returns.
       if (iff->ifTrue->type == Type::unreachable) {
-        return iff->ifFalse;
+        return &iff->ifFalse;
       } else if (iff->ifFalse->type == Type::unreachable) {
-        return iff->ifTrue;
+        return &iff->ifTrue;
       }
     }
   } else if (auto* br = curr->dynCast<Break>()) {
@@ -302,25 +303,33 @@ inline Expression* getImmediateFallthrough(
         behavior == FallthroughBehavior::AllowTeeBrIf &&
         EffectAnalyzer::canReorder(
           passOptions, module, br->condition, br->value)) {
-      return br->value;
+      return &br->value;
     }
   } else if (auto* tryy = curr->dynCast<Try>()) {
     if (!EffectAnalyzer(passOptions, module, tryy->body).throws()) {
-      return tryy->body;
+      return &tryy->body;
     }
   } else if (auto* as = curr->dynCast<RefCast>()) {
-    return as->ref;
+    return &as->ref;
   } else if (auto* as = curr->dynCast<RefAs>()) {
     // Extern conversions are not casts and actually produce new values.
     // Treating them as fallthroughs would lead to misoptimizations of
     // subsequent casts.
     if (as->op != ExternInternalize && as->op != ExternExternalize) {
-      return as->value;
+      return &as->value;
     }
   } else if (auto* br = curr->dynCast<BrOn>()) {
-    return br->ref;
+    return &br->ref;
   }
-  return curr;
+  return currp;
+}
+
+inline Expression* getImmediateFallthrough(
+  Expression* curr,
+  const PassOptions& passOptions,
+  Module& module,
+  FallthroughBehavior behavior = FallthroughBehavior::AllowTeeBrIf) {
+  return *getImmediateFallthroughPtr(&curr, passOptions, module, behavior);
 }
 
 // Similar to getImmediateFallthrough, but looks through multiple children to
diff --git a/src/passes/OptimizeInstructions.cpp b/src/passes/OptimizeInstructions.cpp
index b725ef03f..6537e9b56 100644
--- a/src/passes/OptimizeInstructions.cpp
+++ b/src/passes/OptimizeInstructions.cpp
@@ -1929,23 +1929,44 @@ struct OptimizeInstructions
     Builder builder(*getModule());
     auto nullType = curr->type.getHeapType().getBottom();
     {
-      auto* ref = curr->ref;
+      auto** refp = &curr->ref;
       while (1) {
+        auto* ref = *refp;
+
         auto result = GCTypeUtils::evaluateCastCheck(ref->type, curr->type);
 
         if (result == GCTypeUtils::Success) {
-          // The cast will succeed, but we can't just remove the cast and
-          // replace it with `ref` because the intermediate expressions might
-          // have had side effects. We can replace the cast with a drop followed
-          // by a direct return of the value, though.
-          //
-          // TODO: Do this for non-null values as well by storing the value to
-          // return in a tee.
+          // The cast will succeed. This can only happen if the ref is a subtype
+          // of the cast instruction, which means we can replace the cast with
+          // the ref.
+          assert(Type::isSubType(ref->type, curr->type));
+          if (curr->type != ref->type) {
+            refinalize = true;
+          }
+          // If there were no intermediate expressions, we can just skip the
+          // cast.
+          if (ref == curr->ref) {
+            replaceCurrent(ref);
+            return;
+          }
+          // Otherwise we can't just remove the cast and replace it with `ref`
+          // because the intermediate expressions might have had side effects.
+          // We can replace the cast with a drop followed by a direct return of
+          // the value, though.
           if (ref->type.isNull()) {
+            // We can materialize the resulting null value directly.
             replaceCurrent(builder.makeSequence(builder.makeDrop(curr->ref),
                                                 builder.makeRefNull(nullType)));
             return;
           }
+          // We need to use a tee to return the value since we can't materialize
+          // it directly.
+          auto scratch = builder.addVar(getFunction(), ref->type);
+          *refp = builder.makeLocalTee(scratch, ref, ref->type);
+          replaceCurrent(
+            builder.makeSequence(builder.makeDrop(curr->ref),
+                                 builder.makeLocalGet(scratch, ref->type)));
+          return;
         } else if (result == GCTypeUtils::Failure) {
           // This cast cannot succeed, so it will trap.
           // Make sure to emit a block with the same type as us; leave updating
@@ -1965,10 +1986,10 @@ struct OptimizeInstructions
           return;
         }
 
-        auto* last = ref;
-        ref = Properties::getImmediateFallthrough(
-          ref, getPassOptions(), *getModule());
-        if (ref == last) {
+        auto** last = refp;
+        refp = Properties::getImmediateFallthroughPtr(
+          refp, getPassOptions(), *getModule());
+        if (refp == last) {
           break;
         }
       }