OptimizeCasts: Also handle local.tee (#6507)

Converts the following:

 (some.operation
   (ref.cast .. (local.tee $ref ..))
   (local.get $ref)
 )

into:

 (some.operation
   (local.tee $temp
     (ref.cast .. (local.tee $ref ..))
   )
   (local.get $temp)
 )

2 files changed, 63 insertions, 27 deletions

diff --git a/src/passes/OptimizeCasts.cpp b/src/passes/OptimizeCasts.cpp
index 7877bea9f..5edde4d2f 100644
--- a/src/passes/OptimizeCasts.cpp
+++ b/src/passes/OptimizeCasts.cpp
@@ -93,21 +93,10 @@
 // RefAs with ExternInternalize and ExternExternalize are not considered casts
 // when obtaining fallthroughs, and so are ignored.
 //
-// TODO: 1. Look past individual basic blocks? This may be worth considering
-//          given the pattern of a cast appearing in an if condition that is
-//          then used in an if arm, for example, where simple dominance shows
-//          the cast can be reused.
-// TODO: 2. Look at LocalSet as well and not just Get. That would add some
-//          overlap with the other passes mentioned above (SimplifyLocals and
-//          RedundantSetElimination also track sets and can switch a get to use
-//          a better set's index when that refines the type). But once we do the
-//          first two TODOs above then we'd be adding some novel things here,
-//          as we could optimize "backwards" as well (TODO 1) and past basic
-//          blocks (TODO 2, though RedundantSetElimination does that as well).
-//          However, we should consider whether improving those other passes
-//          might make more sense (as it would help more than casts, if we could
-//          make them operate "backwards" and/or past basic blocks).
-//
+// TODO: Look past individual basic blocks? This may be worth considering
+//       given the pattern of a cast appearing in an if condition that is
+//       then used in an if arm, for example, where simple dominance shows
+//       the cast can be reused.
 
 #include "ir/effects.h"
 #include "ir/linear-execution.h"
@@ -453,20 +442,30 @@ struct BestCastFinder : public LinearExecutionWalker<BestCastFinder> {
   void visitRefCast(RefCast* curr) { handleRefinement(curr); }
 
   void handleRefinement(Expression* curr) {
-    auto* fallthrough = Properties::getFallthrough(curr, options, *getModule());
+    auto* teeFallthrough = Properties::getFallthrough(
+      curr, options, *getModule(), Properties::FallthroughBehavior::NoTeeBrIf);
+    if (auto* tee = teeFallthrough->dynCast<LocalSet>()) {
+      updateBestCast(curr, tee->index);
+    }
+    auto* fallthrough =
+      Properties::getFallthrough(teeFallthrough, options, *getModule());
     if (auto* get = fallthrough->dynCast<LocalGet>()) {
-      auto*& bestCast = mostCastedGets[get->index];
-      if (!bestCast) {
-        // This is the first.
-        bestCast = curr;
-        return;
-      }
+      updateBestCast(curr, get->index);
+    }
+  }
 
-      // See if we are better than the current best.
-      if (curr->type != bestCast->type &&
-          Type::isSubType(curr->type, bestCast->type)) {
-        bestCast = curr;
-      }
+  void updateBestCast(Expression* curr, Index index) {
+    auto*& bestCast = mostCastedGets[index];
+    if (!bestCast) {
+      // This is the first.
+      bestCast = curr;
+      return;
+    }
+
+    // See if we are better than the current best.
+    if (curr->type != bestCast->type &&
+        Type::isSubType(curr->type, bestCast->type)) {
+      bestCast = curr;
     }
   }
 };
diff --git a/test/lit/passes/optimize-casts.wast b/test/lit/passes/optimize-casts.wast
index 512a04e9d..a10540f5d 100644
--- a/test/lit/passes/optimize-casts.wast
+++ b/test/lit/passes/optimize-casts.wast
@@ -1352,6 +1352,43 @@
     )
   )
 
+  ;; CHECK:      (func $local-tee (type $2) (param $x (ref struct))
+  ;; CHECK-NEXT:  (local $y (ref struct))
+  ;; CHECK-NEXT:  (local $2 (ref $A))
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (local.tee $2
+  ;; CHECK-NEXT:    (ref.cast (ref $A)
+  ;; CHECK-NEXT:     (local.tee $y
+  ;; CHECK-NEXT:      (local.get $x)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (local.get $2)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (local.get $2)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $local-tee (param $x (ref struct))
+    (local $y (ref struct))
+    ;; We should use the cast value after it has been computed, in both gets.
+    (drop
+      (ref.cast (ref $A)
+        (local.tee $y
+          (local.get $x)
+        )
+      )
+    )
+    (drop
+      (local.get $x)
+    )
+    (drop
+      (local.get $y)
+    )
+  )
+
   ;; CHECK:      (func $get (type $11) (result (ref struct))
   ;; CHECK-NEXT:  (unreachable)
   ;; CHECK-NEXT: )