[Wasm GC] Optimize RefIs and RefAs when the type lets us (#3725)

This is similar to the optimization of BrOn in #3719 and #3724. When the
type tells us the kind of input we have, we can tell at compile time what
result we'll get, like ref.is_func of something with type (ref func) will
always return 1, etc.

There are some code size and perf tradeoffs that should be looked into
more and are marked as TODOs.

3 files changed, 500 insertions, 1 deletions

diff --git a/src/ir/gc-type-utils.h b/src/ir/gc-type-utils.h
index 168c9529e..d59193193 100644
--- a/src/ir/gc-type-utils.h
+++ b/src/ir/gc-type-utils.h
@@ -42,7 +42,7 @@ enum EvaluationResult {
 // (like br_on_func checks if it is a function), see if type info lets us
 // determine that at compile time.
 // This ignores nullability - it just checks the kind.
-EvaluationResult evaluateKindCheck(Expression* curr) {
+inline EvaluationResult evaluateKindCheck(Expression* curr) {
   Kind expected;
   Expression* child;
 
@@ -66,6 +66,42 @@ EvaluationResult evaluateKindCheck(Expression* curr) {
         WASM_UNREACHABLE("unhandled BrOn");
     }
     child = br->ref;
+  } else if (auto* is = curr->dynCast<RefIs>()) {
+    switch (is->op) {
+      // We don't check nullability here.
+      case RefIsNull:
+        return Unknown;
+      case RefIsFunc:
+        expected = Func;
+        break;
+      case RefIsData:
+        expected = Data;
+        break;
+      case RefIsI31:
+        expected = I31;
+        break;
+      default:
+        WASM_UNREACHABLE("unhandled BrOn");
+    }
+    child = is->value;
+  } else if (auto* as = curr->dynCast<RefAs>()) {
+    switch (as->op) {
+      // We don't check nullability here.
+      case RefAsNonNull:
+        return Unknown;
+      case RefAsFunc:
+        expected = Func;
+        break;
+      case RefAsData:
+        expected = Data;
+        break;
+      case RefAsI31:
+        expected = I31;
+        break;
+      default:
+        WASM_UNREACHABLE("unhandled BrOn");
+    }
+    child = as->value;
   } else {
     WASM_UNREACHABLE("invalid input to evaluateKindCheck");
   }
diff --git a/src/passes/OptimizeInstructions.cpp b/src/passes/OptimizeInstructions.cpp
index be89257a1..972ffd794 100644
--- a/src/passes/OptimizeInstructions.cpp
+++ b/src/passes/OptimizeInstructions.cpp
@@ -26,6 +26,7 @@
 #include <ir/bits.h>
 #include <ir/cost.h>
 #include <ir/effects.h>
+#include <ir/gc-type-utils.h>
 #include <ir/literal-utils.h>
 #include <ir/load-utils.h>
 #include <ir/manipulation.h>
@@ -1005,6 +1006,98 @@ struct OptimizeInstructions
     }
   }
 
+  void visitRefIs(RefIs* curr) {
+    if (curr->type == Type::unreachable) {
+      return;
+    }
+
+    // Optimizating RefIs is not that obvious, since even if we know the result
+    // evaluates to 0 or 1 then the replacement may not actually save code size,
+    // since RefIsNull is a single byte (the others are 2), while adding a Const
+    // of 0 would be two bytes. Other factors are that we can remove the input
+    // and the added drop on it if it has no side effects, and that replacing
+    // with a constant may allow further optimizations later. For now, replace
+    // with a constant, but this warrants more investigation. TODO
+
+    Builder builder(*getModule());
+
+    auto nonNull = !curr->value->type.isNullable();
+
+    if (curr->op == RefIsNull) {
+      if (nonNull) {
+        replaceCurrent(builder.makeSequence(
+          builder.makeDrop(curr->value),
+          builder.makeConst(Literal::makeZero(Type::i32))));
+      }
+      return;
+    }
+
+    // Check if the type is the kind we are checking for.
+    auto result = GCTypeUtils::evaluateKindCheck(curr);
+
+    if (result != GCTypeUtils::Unknown) {
+      // We know the kind. Now we must also take into account nullability.
+      if (nonNull) {
+        // We know the entire result.
+        replaceCurrent(
+          builder.makeSequence(builder.makeDrop(curr->value),
+                               builder.makeConst(Literal::makeFromInt32(
+                                 result == GCTypeUtils::Success, Type::i32))));
+      } else {
+        // The value may be null. Leave only a check for that.
+        curr->op = RefIsNull;
+        if (result == GCTypeUtils::Success) {
+          // The input is of the right kind. If it is not null then the result
+          // is 1, and otherwise it is 0, so we need to flip the result of
+          // RefIsNull.
+          // Note that even after adding an eqz here we do not regress code size
+          // as RefIsNull is a single byte while the others are two. So we keep
+          // code size identical. However, in theory this may be more work, if
+          // a VM considers ref.is_X to be as fast as ref.is_null, and if eqz is
+          // not free, so this is worth more investigation. TODO
+          replaceCurrent(builder.makeUnary(EqZInt32, curr));
+        } else {
+          // The input is of the wrong kind. In this case if it is null we
+          // return zero because of that, and if it is not then we return zero
+          // because of the kind, so the result is always the same.
+          assert(result == GCTypeUtils::Failure);
+          replaceCurrent(builder.makeSequence(
+            builder.makeDrop(curr->value),
+            builder.makeConst(Literal::makeZero(Type::i32))));
+        }
+      }
+    }
+  }
+
+  void visitRefAs(RefAs* curr) {
+    if (curr->type == Type::unreachable) {
+      return;
+    }
+
+    // Check if the type is the kind we are checking for.
+    auto result = GCTypeUtils::evaluateKindCheck(curr);
+
+    if (result == GCTypeUtils::Success) {
+      // We know the kind is correct, so all that is left is a check for
+      // non-nullability, which we do lower down.
+      curr->op = RefAsNonNull;
+    } else if (result == GCTypeUtils::Failure) {
+      // This is the wrong kind, so it will trap. The binaryen optimizer does
+      // not differentiate traps, so we can perform a replacement here. We
+      // replace 2 bytes of ref.as_* with one byte of unreachable and one of a
+      // drop, which is no worse, and the value and the drop can be optimized
+      // out later if the value has no side effects.
+      Builder builder(*getModule());
+      replaceCurrent(builder.makeSequence(builder.makeDrop(curr->value),
+                                          builder.makeUnreachable()));
+      return;
+    }
+
+    if (curr->op == RefAsNonNull && !curr->value->type.isNullable()) {
+      replaceCurrent(curr->value);
+    }
+  }
+
   Index getMaxBitsForLocal(LocalGet* get) {
     // check what we know about the local
     return localInfo[get->index].maxBits;
diff --git a/test/lit/passes/optimize-instructions-gc.wast b/test/lit/passes/optimize-instructions-gc.wast
index 96cd19583..5c469b617 100644
--- a/test/lit/passes/optimize-instructions-gc.wast
+++ b/test/lit/passes/optimize-instructions-gc.wast
@@ -90,4 +90,374 @@
       (i32.const 0x123) ;; data over 0xff is unnecessary
     )
   )
+
+  ;; ref.is_null is not needed on a non-nullable value, and if something is
+  ;; a func we don't need that either etc. if we know the result
+  ;; CHECK:      (func $unneeded_is (param $struct (ref $struct)) (param $func (ref func)) (param $data dataref) (param $i31 i31ref)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result i32)
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (local.get $struct)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i32.const 0)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result i32)
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (local.get $func)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i32.const 1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result i32)
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (local.get $data)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i32.const 1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result i32)
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (local.get $i31)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i32.const 1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $unneeded_is
+    (param $struct (ref $struct))
+    (param $func (ref func))
+    (param $data (ref data))
+    (param $i31 (ref i31))
+    (drop
+      (ref.is_null (local.get $struct))
+    )
+    (drop
+       (ref.is_func (local.get $func))
+     )
+    (drop
+      (ref.is_data (local.get $data))
+    )
+    (drop
+      (ref.is_i31 (local.get $i31))
+    )
+  )
+
+  ;; similar to $unneeded_is, but the values are nullable. we can at least
+  ;; leave just the null check.
+  ;; CHECK:      (func $unneeded_is_null (param $struct (ref null $struct)) (param $func funcref) (param $data (ref null data)) (param $i31 (ref null i31))
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (ref.is_null
+  ;; CHECK-NEXT:    (local.get $struct)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (i32.eqz
+  ;; CHECK-NEXT:    (ref.is_null
+  ;; CHECK-NEXT:     (local.get $func)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (i32.eqz
+  ;; CHECK-NEXT:    (ref.is_null
+  ;; CHECK-NEXT:     (local.get $data)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (i32.eqz
+  ;; CHECK-NEXT:    (ref.is_null
+  ;; CHECK-NEXT:     (local.get $i31)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $unneeded_is_null
+   (param $struct (ref null $struct))
+    (param $func (ref null func))
+    (param $data (ref null data))
+    (param $i31 (ref null i31))
+    (drop
+      (ref.is_null (local.get $struct))
+    )
+    (drop
+      (ref.is_func (local.get $func))
+    )
+    (drop
+      (ref.is_data (local.get $data))
+    )
+    (drop
+      (ref.is_i31 (local.get $i31))
+    )
+  )
+
+  ;; similar to $unneeded_is, but the values are of mixed kind (is_func of
+  ;; data, etc.). regardless of nullability the result here is always 0.
+  ;; CHECK:      (func $unneeded_is_bad_kinds (param $func funcref) (param $data (ref null data)) (param $i31 (ref null i31))
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result i32)
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (local.get $data)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i32.const 0)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result i32)
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (local.get $i31)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i32.const 0)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result i32)
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (local.get $func)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i32.const 0)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result i32)
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (ref.as_non_null
+  ;; CHECK-NEXT:      (local.get $data)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i32.const 0)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result i32)
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (ref.as_non_null
+  ;; CHECK-NEXT:      (local.get $i31)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i32.const 0)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result i32)
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (ref.as_non_null
+  ;; CHECK-NEXT:      (local.get $func)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i32.const 0)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $unneeded_is_bad_kinds
+    (param $func (ref null func))
+    (param $data (ref null data))
+    (param $i31 (ref null i31))
+    (drop
+      (ref.is_func (local.get $data))
+    )
+    (drop
+      (ref.is_data (local.get $i31))
+    )
+    (drop
+      (ref.is_i31 (local.get $func))
+    )
+    ;; also check non-nullable types as inputs
+    (drop
+      (ref.is_func (ref.as_non_null (local.get $data)))
+    )
+    (drop
+      (ref.is_data (ref.as_non_null (local.get $i31)))
+    )
+    (drop
+      (ref.is_i31 (ref.as_non_null (local.get $func)))
+    )
+  )
+
+  ;; ref.as_non_null is not needed on a non-nullable value, and if something is
+  ;; a func we don't need that either etc., and can just return the value.
+  ;; CHECK:      (func $unneeded_as (param $struct (ref $struct)) (param $func (ref func)) (param $data dataref) (param $i31 i31ref)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (local.get $struct)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (local.get $func)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (local.get $data)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (local.get $i31)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $unneeded_as
+    (param $struct (ref $struct))
+    (param $func (ref func))
+    (param $data (ref data))
+    (param $i31 (ref i31))
+    (drop
+      (ref.as_non_null (local.get $struct))
+    )
+    (drop
+      (ref.as_func (local.get $func))
+    )
+    (drop
+      (ref.as_data (local.get $data))
+    )
+    (drop
+      (ref.as_i31 (local.get $i31))
+    )
+  )
+
+  ;; similar to $unneeded_as, but the values are nullable. we can turn the
+  ;; more specific things into ref.as_non_null.
+  ;; CHECK:      (func $unneeded_as_null (param $struct (ref null $struct)) (param $func funcref) (param $data (ref null data)) (param $i31 (ref null i31))
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (ref.as_non_null
+  ;; CHECK-NEXT:    (local.get $struct)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (ref.as_non_null
+  ;; CHECK-NEXT:    (local.get $func)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (ref.as_non_null
+  ;; CHECK-NEXT:    (local.get $data)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (ref.as_non_null
+  ;; CHECK-NEXT:    (local.get $i31)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $unneeded_as_null
+    (param $struct (ref null $struct))
+    (param $func (ref null func))
+    (param $data (ref null data))
+    (param $i31 (ref null i31))
+    (drop
+      (ref.as_non_null (local.get $struct))
+    )
+    (drop
+      (ref.as_func (local.get $func))
+    )
+    (drop
+      (ref.as_data (local.get $data))
+    )
+    (drop
+      (ref.as_i31 (local.get $i31))
+    )
+  )
+
+  ;; similar to $unneeded_as, but the values are of mixed kind (as_func of
+  ;; data, etc.), so we know we will trap
+  ;; CHECK:      (func $unneeded_as_bad_kinds (param $func funcref) (param $data (ref null data)) (param $i31 (ref null i31))
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (local.get $data)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (unreachable)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (local.get $i31)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (unreachable)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (local.get $func)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (unreachable)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (ref.as_non_null
+  ;; CHECK-NEXT:      (local.get $data)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (unreachable)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (ref.as_non_null
+  ;; CHECK-NEXT:      (local.get $i31)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (unreachable)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (ref.as_non_null
+  ;; CHECK-NEXT:      (local.get $func)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (unreachable)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $unneeded_as_bad_kinds
+    (param $func (ref null func))
+    (param $data (ref null data))
+    (param $i31 (ref null i31))
+    (drop
+      (ref.as_func (local.get $data))
+    )
+    (drop
+      (ref.as_data (local.get $i31))
+    )
+    (drop
+      (ref.as_i31 (local.get $func))
+    )
+    ;; also check non-nullable types as inputs
+    (drop
+      (ref.as_func (ref.as_non_null (local.get $data)))
+    )
+    (drop
+      (ref.as_data (ref.as_non_null (local.get $i31)))
+    )
+    (drop
+      (ref.as_i31 (ref.as_non_null (local.get $func)))
+    )
+  )
+
+  ;; CHECK:      (func $unneeded_unreachability
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (ref.is_func
+  ;; CHECK-NEXT:    (unreachable)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (ref.as_func
+  ;; CHECK-NEXT:    (unreachable)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $unneeded_unreachability
+    ;; unreachable instructions can simply be ignored
+    (drop
+      (ref.is_func (unreachable))
+    )
+    (drop
+      (ref.as_func (unreachable))
+    )
+  )
 )