Improve br_on* optimizations (#5887)

Optimize both the known-null and known-non-null cases for BrOnNull and
BrOnNonNull and optimize for more cast behaviors such as SuccessOnlyIfNonNull
and Unreachable for BrOnCast and BrOnCastFail. Leave optimizing
SuccessOnlyIfNull to future work, since that's more complicated. Use type
information from fallthrough values to inform all the optimizations.

1 files changed, 141 insertions, 35 deletions

diff --git a/src/passes/RemoveUnusedBrs.cpp b/src/passes/RemoveUnusedBrs.cpp
index dc7741c0e..4fda7afb2 100644
--- a/src/passes/RemoveUnusedBrs.cpp
+++ b/src/passes/RemoveUnusedBrs.cpp
@@ -699,8 +699,11 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
     }
 
     struct Optimizer : public PostWalker<Optimizer> {
+      PassOptions& passOptions;
       bool worked = false;
 
+      Optimizer(PassOptions& passOptions) : passOptions(passOptions) {}
+
       void visitBrOn(BrOn* curr) {
         // Ignore unreachable BrOns which we cannot improve anyhow. Note that
         // we must check the ref field manually, as we may be changing types as
@@ -712,56 +715,159 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
           return;
         }
 
-        // First, check for a possible null which would prevent optimizations on
-        // null checks.
-        // TODO: Use the fallthrough to determine in more cases that we
-        // definitely have a null.
-        auto refType = curr->ref->type;
-        if (refType.isNullable() &&
-            (curr->op == BrOnNull || curr->op == BrOnNonNull)) {
-          return;
-        }
+        Builder builder(*getModule());
+
+        Type refType =
+          Properties::getFallthroughType(curr->ref, passOptions, *getModule());
+        assert(refType.isRef());
+
+        // When we optimize based on all the fallthrough type information
+        // available, we may need to insert a cast to maintain validity. For
+        // example, in this case we know the cast will succeed, but it would be
+        // invalid to send curr->ref directly:
+        //
+        //   (br_on_cast $l anyref i31ref
+        //     (block (result anyref)
+        //       (i31.new ...)))
+        //
+        // We could just always do the cast and leave removing the casts to
+        // OptimizeInstructions, but it's simple enough to avoid unnecessary
+        // casting here.
+        auto maybeCast = [&](Expression* expr, Type type) -> Expression* {
+          assert(expr->type.isRef() && type.isRef());
+          if (Type::isSubType(expr->type, type)) {
+            return expr;
+          }
+          if (HeapType::isSubType(expr->type.getHeapType(),
+                                  type.getHeapType())) {
+            return builder.makeRefAs(RefAsNonNull, expr);
+          }
+          return builder.makeRefCast(expr, type);
+        };
 
         if (curr->op == BrOnNull) {
-          assert(refType.isNonNullable());
-          // This cannot be null, so the br is never taken, and the non-null
-          // value flows through.
-          replaceCurrent(curr->ref);
-          worked = true;
+          if (refType.isNull()) {
+            // The branch will definitely be taken.
+            replaceCurrent(builder.makeSequence(builder.makeDrop(curr->ref),
+                                                builder.makeBreak(curr->name)));
+            worked = true;
+            return;
+          }
+          if (refType.isNonNullable()) {
+            // The branch will definitely not be taken.
+            replaceCurrent(maybeCast(curr->ref, curr->type));
+            worked = true;
+            return;
+          }
           return;
         }
+
         if (curr->op == BrOnNonNull) {
-          assert(refType.isNonNullable());
-          // This cannot be null, so the br is always taken.
-          replaceCurrent(
-            Builder(*getModule()).makeBreak(curr->name, curr->ref));
-          worked = true;
+          if (refType.isNull()) {
+            // Definitely not taken.
+            replaceCurrent(builder.makeDrop(curr->ref));
+            worked = true;
+            return;
+          }
+          if (refType.isNonNullable()) {
+            // Definitely taken.
+            replaceCurrent(builder.makeBreak(
+              curr->name, maybeCast(curr->ref, curr->getSentType())));
+            worked = true;
+            return;
+          }
           return;
         }
 
-        // Check if the type is the kind we are checking for.
+        // Improve the cast target type as much as possible given what we know
+        // about the input. Unlike in BrOn::finalize(), we consider type
+        // information from all the fallthrough values here. We can continue to
+        // further optimizations after this, and those optimizations might even
+        // benefit from this improvement.
+        auto glb = Type::getGreatestLowerBound(curr->castType, refType);
+        if (glb != Type::unreachable && glb != curr->castType) {
+          curr->castType = glb;
+          curr->finalize();
+          worked = true;
+        }
+
+        // Depending on what we know about the cast results, we may be able to
+        // optimize.
         auto result = GCTypeUtils::evaluateCastCheck(refType, curr->castType);
         if (curr->op == BrOnCastFail) {
           result = GCTypeUtils::flipEvaluationResult(result);
         }
 
-        if (result == GCTypeUtils::Success) {
-          // The cast succeeds, so we can switch from BrOn to a simple br that
-          // is always taken.
-          replaceCurrent(
-            Builder(*getModule()).makeBreak(curr->name, curr->ref));
-          worked = true;
-        } else if (result == GCTypeUtils::Failure ||
-                   result == GCTypeUtils::Unreachable) {
-          // The cast fails, so the branch is never taken, and the value just
-          // flows through. Or, the cast cannot even be reached, so it does not
-          // matter what we do, and we can handle it as a failure.
-          replaceCurrent(curr->ref);
-          worked = true;
+        switch (result) {
+          case GCTypeUtils::Unknown:
+            // Anything could happen, so we cannot optimize.
+            return;
+          case GCTypeUtils::Success: {
+            replaceCurrent(builder.makeBreak(
+              curr->name, maybeCast(curr->ref, curr->getSentType())));
+            worked = true;
+            return;
+          }
+          case GCTypeUtils::Failure: {
+            replaceCurrent(maybeCast(curr->ref, curr->type));
+            worked = true;
+            return;
+          }
+          case GCTypeUtils::SuccessOnlyIfNull: {
+            // TODO: optimize this case using the following replacement, which
+            // avoids using any scratch locals and only does a single null
+            // check, but does require generating a fresh label:
+            //
+            //   (br_on_cast $l (ref null $X) (ref null $Y)
+            //     (...)
+            //   )
+            //     =>
+            //   (block $l' (result (ref $X))
+            //     (br_on_non_null $l' ;; reuses `curr`
+            //       (...)
+            //     )
+            //     (br $l
+            //       (ref.null bot<X>)
+            //     )
+            //   )
+            return;
+          }
+          case GCTypeUtils::SuccessOnlyIfNonNull: {
+            // Perform this replacement:
+            //
+            //   (br_on_cast $l (ref null $X') (ref $X))
+            //     (...)
+            //   )
+            //     =>
+            //   (block (result (ref bot<X>))
+            //     (br_on_non_null $l ;; reuses `curr`
+            //       (...)
+            //     (ref.null bot<X>)
+            //   )
+            curr->ref = maybeCast(
+              curr->ref, Type(curr->getSentType().getHeapType(), Nullable));
+            curr->op = BrOnNonNull;
+            curr->castType = Type::none;
+            curr->type = Type::none;
+
+            assert(curr->ref->type.isRef());
+            auto* refNull = builder.makeRefNull(curr->ref->type.getHeapType());
+            replaceCurrent(builder.makeBlock({curr, refNull}, refNull->type));
+            worked = true;
+            return;
+          }
+          case GCTypeUtils::Unreachable: {
+            // The cast is never executed, possibly because its input type is
+            // uninhabitable. Replace it with unreachable.
+            auto* drop = builder.makeDrop(curr->ref);
+            auto* unreachable = ExpressionManipulator::unreachable(curr);
+            replaceCurrent(builder.makeBlock({drop, unreachable}));
+            worked = true;
+            return;
+          }
         }
-        // TODO: Handle SuccessOnlyIfNull and SuccessOnlyIfNonNull.
       }
-    } optimizer;
+    } optimizer(getPassOptions());
 
     optimizer.setModule(getModule());
     optimizer.doWalkFunction(func);