Add CastableBase::To() overload with predicate

This allows for a more optimal way to filter the result of To(). Updated
Type query functions to make use of it. Added tests.

Change-Id: If3a65259345fbe6b92c6d367ab01fa718bb7cfee
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/44463
Reviewed-by: Ben Clayton <bclayton@google.com>
Commit-Queue: Antonio Maiorano <amaiorano@google.com>

src/castable.h

@@ -94,7 +94,7 @@ struct IsAnyOf;
 /// class `TO`
 /// @param obj the object to test from
 template <typename TO, typename FROM>
-bool Is(FROM* obj) {
+inline bool Is(FROM* obj) {
   constexpr const bool downcast = std::is_base_of<FROM, TO>::value;
   constexpr const bool upcast = std::is_base_of<TO, FROM>::value;
   constexpr const bool nocast = std::is_same<FROM, TO>::value;
@@ -111,6 +111,28 @@ bool Is(FROM* obj) {
   return obj->TypeInfo().Is(TypeInfo::Of<std::remove_const_t<TO>>());
 }
 
+/// @returns true if `obj` is a valid pointer, and is of, or derives from the
+/// class `TO`, and pred(const TO*) returns true
+/// @param obj the object to test from
+/// @param pred predicate function with signature `bool(const TO*)` called iff
+/// object is of, or derives from the class `TO`.
+template <typename TO, typename FROM, typename Pred>
+inline bool Is(FROM* obj, Pred&& pred) {
+  constexpr const bool downcast = std::is_base_of<FROM, TO>::value;
+  constexpr const bool upcast = std::is_base_of<TO, FROM>::value;
+  constexpr const bool nocast = std::is_same<FROM, TO>::value;
+  static_assert(upcast || downcast || nocast, "impossible cast");
+
+  if (obj == nullptr) {
+    return false;
+  }
+
+  bool is_type = upcast || nocast ||
+                 obj->TypeInfo().Is(TypeInfo::Of<std::remove_const_t<TO>>());
+
+  return is_type && pred(static_cast<std::add_const_t<TO>*>(obj));
+}
+
 /// @returns true if `obj` is of, or derives from any of the `TO`
 /// classes.
 /// @param obj the object to cast from
@@ -150,6 +172,15 @@ class CastableBase {
     return tint::Is<TO>(this);
   }
 
+  /// @returns true if this object is of, or derives from the class `TO` and
+  /// pred(const TO*) returns true
+  /// @param pred predicate function with signature `bool(const TO*)` called iff
+  /// object is of, or derives from the class `TO`.
+  template <typename TO, typename Pred>
+  inline bool Is(Pred&& pred) const {
+    return tint::Is<TO>(this, std::forward<Pred>(pred));
+  }
+
   /// @returns true if this object is of, or derives from any of the `TO`
   /// classes.
   template <typename... TO>
@@ -219,6 +250,16 @@ class Castable : public BASE {
     return tint::Is<TO>(static_cast<const CLASS*>(this));
   }
 
+  /// @returns true if this object is of, or derives from the class `TO` and
+  /// pred(const TO*) returns true
+  /// @param pred predicate function with signature `bool(const TO*)` called iff
+  /// object is of, or derives from the class `TO`.
+  template <typename TO, typename Pred>
+  inline bool Is(Pred&& pred) const {
+    return tint::Is<TO>(static_cast<const CLASS*>(this),
+                        std::forward<Pred>(pred));
+  }
+
   /// @returns true if this object is of, or derives from any of the `TO`
   /// classes.
   template <typename... TO>

src/castable_test.cc

@@ -73,6 +73,25 @@ TEST(CastableBase, Is) {
   ASSERT_TRUE(gecko->Is<Reptile>());
 }
 
+TEST(CastableBase, IsWithPredicate) {
+  std::unique_ptr<CastableBase> frog = std::make_unique<Frog>();
+
+  frog->Is<Animal>([&frog](const Animal* a) {
+    EXPECT_EQ(a, frog.get());
+    return true;
+  });
+
+  ASSERT_TRUE((frog->Is<Animal>([](const Animal* a) { return true; })));
+  ASSERT_FALSE((frog->Is<Animal>([](const Animal* a) { return false; })));
+
+  // Predicate not called if cast is invalid
+  auto expect_not_called = [] { FAIL() << "Should not be called"; };
+  ASSERT_FALSE((frog->Is<Bear>([&](const Animal* a) {
+    expect_not_called();
+    return true;
+  })));
+}
+
 TEST(CastableBase, IsAnyOf) {
   std::unique_ptr<CastableBase> frog = std::make_unique<Frog>();
   std::unique_ptr<CastableBase> bear = std::make_unique<Bear>();
@@ -138,6 +157,25 @@ TEST(Castable, Is) {
   ASSERT_TRUE(gecko->Is<Reptile>());
 }
 
+TEST(Castable, IsWithPredicate) {
+  std::unique_ptr<Animal> frog = std::make_unique<Frog>();
+
+  frog->Is<Animal>([&frog](const Animal* a) {
+    EXPECT_EQ(a, frog.get());
+    return true;
+  });
+
+  ASSERT_TRUE((frog->Is<Animal>([](const Animal* a) { return true; })));
+  ASSERT_FALSE((frog->Is<Animal>([](const Animal* a) { return false; })));
+
+  // Predicate not called if cast is invalid
+  auto expect_not_called = [] { FAIL() << "Should not be called"; };
+  ASSERT_FALSE((frog->Is<Bear>([&](const Animal* a) {
+    expect_not_called();
+    return true;
+  })));
+}
+
 TEST(Castable, As) {
   std::unique_ptr<Animal> frog = std::make_unique<Frog>();
   std::unique_ptr<Animal> bear = std::make_unique<Bear>();

src/type/type.cc

@@ -79,11 +79,13 @@ bool Type::is_float_scalar() const {
 }
 
 bool Type::is_float_matrix() const {
-  return Is<Matrix>() && As<Matrix>()->type()->is_float_scalar();
+  return Is<Matrix>(
+      [](const Matrix* m) { return m->type()->is_float_scalar(); });
 }
 
 bool Type::is_float_vector() const {
-  return Is<Vector>() && As<Vector>()->type()->is_float_scalar();
+  return Is<Vector>(
+      [](const Vector* v) { return v->type()->is_float_scalar(); });
 }
 
 bool Type::is_float_scalar_or_vector() const {
@@ -95,19 +97,19 @@ bool Type::is_integer_scalar() const {
 }
 
 bool Type::is_unsigned_integer_vector() const {
-  return Is<Vector>() && As<Vector>()->type()->Is<U32>();
+  return Is<Vector>([](const Vector* v) { return v->type()->Is<U32>(); });
 }
 
 bool Type::is_signed_integer_vector() const {
-  return Is<Vector>() && As<Vector>()->type()->Is<I32>();
+  return Is<Vector>([](const Vector* v) { return v->type()->Is<I32>(); });
 }
 
 bool Type::is_unsigned_scalar_or_vector() const {
-  return Is<U32>() || (Is<Vector>() && As<Vector>()->type()->Is<U32>());
+  return Is<U32>() || is_unsigned_integer_vector();
 }
 
 bool Type::is_signed_scalar_or_vector() const {
-  return Is<I32>() || (Is<Vector>() && As<Vector>()->type()->Is<I32>());
+  return Is<I32>() || is_signed_integer_vector();
 }
 
 bool Type::is_integer_scalar_or_vector() const {
@@ -115,7 +117,7 @@ bool Type::is_integer_scalar_or_vector() const {
 }
 
 bool Type::is_bool_vector() const {
-  return Is<Vector>() && As<Vector>()->type()->Is<Bool>();
+  return Is<Vector>([](const Vector* v) { return v->type()->Is<Bool>(); });
 }
 
 bool Type::is_bool_scalar_or_vector() const {