Add specialized implementations for ==/!=/@in.

PiperOrigin-RevId: 716429211

Author: jnthntatum

eval/compiler/BUILD

@@ -101,6 +101,7 @@ cc_library(
         "//common:ast",
         "//common:ast_traverse",
         "//common:ast_visitor",
+        "//common:kind",
         "//common:memory",
         "//common:type",
         "//common:value",
@@ -111,6 +112,7 @@ cc_library(
         "//eval/eval:create_map_step",
         "//eval/eval:create_struct_step",
         "//eval/eval:direct_expression_step",
+        "//eval/eval:equality_steps",
         "//eval/eval:evaluator_core",
         "//eval/eval:function_step",
         "//eval/eval:ident_step",
@@ -159,6 +161,7 @@ cc_test(
         ":constant_folding",
         ":flat_expr_builder",
         ":qualified_reference_resolver",
+        "//base:builtins",
         "//base:function",
         "//base:function_descriptor",
         "//common:value",
@@ -178,18 +181,18 @@ cc_test(
         "//eval/public/containers:container_backed_map_impl",
         "//eval/public/structs:cel_proto_descriptor_pool_builder",
         "//eval/public/structs:cel_proto_wrapper",
-        "//eval/public/structs:protobuf_descriptor_type_provider",
         "//eval/public/testing:matchers",
         "//eval/testutil:test_message_cc_proto",
-        "//internal:proto_file_util",
         "//internal:proto_matchers",
         "//internal:status_macros",
         "//internal:testing",
         "//parser",
+        "//runtime:function_adapter",
         "//runtime:runtime_options",
         "//runtime/internal:runtime_env_testing",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:span",
         "@com_google_cel_spec//proto/cel/expr:checked_cc_proto",

eval/compiler/flat_expr_builder.cc

@@ -54,6 +54,7 @@
 #include "common/ast.h"
 #include "common/ast_traverse.h"
 #include "common/ast_visitor.h"
+#include "common/kind.h"
 #include "common/memory.h"
 #include "common/type.h"
 #include "common/value.h"
@@ -68,6 +69,7 @@
 #include "eval/eval/create_map_step.h"
 #include "eval/eval/create_struct_step.h"
 #include "eval/eval/direct_expression_step.h"
+#include "eval/eval/equality_steps.h"
 #include "eval/eval/evaluator_core.h"
 #include "eval/eval/function_step.h"
 #include "eval/eval/ident_step.h"
@@ -527,6 +529,36 @@ class FlatExprVisitor : public cel::AstVisitor {
                  const cel::ast_internal::Call& call) {
             return HandleNot(expr, call);
           };
+      if (options_.enable_heterogeneous_equality) {
+        for (const auto& in_op :
+             {cel::builtin::kIn, cel::builtin::kInDeprecated,
+              cel::builtin::kInFunction}) {
+          call_handlers_[in_op] = [this](const cel::ast_internal::Expr& expr,
+                                         const cel::ast_internal::Call& call) {
+            return HandleHeterogeneousEqualityIn(expr, call);
+          };
+        }
+        // Try to detect if the environment is setup with a custom equality
+        // implementation.
+        if (resolver_
+                .FindOverloads(cel::builtin::kEqual,
+                               /*receiver_style=*/false,
+                               {cel::Kind::kAny, cel::Kind::kAny})
+                .empty()) {
+          call_handlers_[cel::builtin::kEqual] =
+              [this](const cel::ast_internal::Expr& expr,
+                     const cel::ast_internal::Call& call) {
+                return HandleHeterogeneousEquality(expr, call,
+                                                   /*inequality=*/false);
+              };
+          call_handlers_[cel::builtin::kInequal] =
+              [this](const cel::ast_internal::Expr& expr,
+                     const cel::ast_internal::Call& call) {
+                return HandleHeterogeneousEquality(expr, call,
+                                                   /*inequality=*/true);
+              };
+        }
+      }
     }
   }
 
@@ -1874,6 +1906,13 @@ class FlatExprVisitor : public cel::AstVisitor {
   CallHandlerResult HandleNotStrictlyFalse(const cel::ast_internal::Expr& expr,
                                            const cel::ast_internal::Call& call);
 
+  CallHandlerResult HandleHeterogeneousEquality(
+      const cel::ast_internal::Expr& expr, const cel::ast_internal::Call& call,
+      bool inequality);
+
+  CallHandlerResult HandleHeterogeneousEqualityIn(
+      const cel::ast_internal::Expr& expr, const cel::ast_internal::Call& call);
+
   const Resolver& resolver_;
   ValueManager& value_factory_;
   absl::Status progress_status_;
@@ -2026,6 +2065,59 @@ FlatExprVisitor::CallHandlerResult FlatExprVisitor::HandleListAppend(
   return CallHandlerResult::kNotIntercepted;
 }
 
+FlatExprVisitor::CallHandlerResult FlatExprVisitor::HandleHeterogeneousEquality(
+    const cel::ast_internal::Expr& expr, const cel::ast_internal::Call& call,
+    bool inequality) {
+  if (!ValidateOrError(
+          call.args().size() == 2,
+          "unexpected number of args for builtin equality operator")) {
+    return CallHandlerResult::kIntercepted;
+  }
+  auto depth = RecursionEligible();
+
+  if (depth.has_value()) {
+    auto args = ExtractRecursiveDependencies();
+    if (args.size() != 2) {
+      SetProgressStatusError(absl::InvalidArgumentError(
+          "unexpected number of args for builtin equality operator"));
+      return CallHandlerResult::kIntercepted;
+    }
+    SetRecursiveStep(
+        CreateDirectEqualityStep(std::move(args[0]), std::move(args[1]),
+                                 inequality, expr.id()),
+        *depth + 1);
+    return CallHandlerResult::kIntercepted;
+  }
+  AddStep(CreateEqualityStep(inequality, expr.id()));
+  return CallHandlerResult::kIntercepted;
+}
+
+FlatExprVisitor::CallHandlerResult
+FlatExprVisitor::HandleHeterogeneousEqualityIn(
+    const cel::ast_internal::Expr& expr, const cel::ast_internal::Call& call) {
+  if (!ValidateOrError(call.args().size() == 2,
+                       "unexpected number of args for builtin 'in' operator")) {
+    return CallHandlerResult::kIntercepted;
+  }
+
+  auto depth = RecursionEligible();
+  if (depth.has_value()) {
+    auto args = ExtractRecursiveDependencies();
+    if (args.size() != 2) {
+      SetProgressStatusError(absl::InvalidArgumentError(
+          "unexpected number of args for builtin 'in' operator"));
+      return CallHandlerResult::kIntercepted;
+    }
+    SetRecursiveStep(
+        CreateDirectInStep(std::move(args[0]), std::move(args[1]), expr.id()),
+        *depth + 1);
+    return CallHandlerResult::kIntercepted;
+  }
+
+  AddStep(CreateInStep(expr.id()));
+  return CallHandlerResult::kIntercepted;
+}
+
 void BinaryCondVisitor::PreVisit(const cel::ast_internal::Expr* expr) {
   switch (cond_) {
     case BinaryCond::kAnd:

eval/compiler/flat_expr_builder_test.cc

@@ -16,6 +16,7 @@
 
 #include "eval/compiler/flat_expr_builder.h"
 
+#include <cstdint>
 #include <functional>
 #include <memory>
 #include <string>
@@ -28,9 +29,11 @@
 #include "google/protobuf/descriptor.pb.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
 #include "absl/strings/str_split.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
+#include "base/builtins.h"
 #include "base/function.h"
 #include "base/function_descriptor.h"
 #include "common/value.h"
@@ -51,16 +54,15 @@
 #include "eval/public/portable_cel_function_adapter.h"
 #include "eval/public/structs/cel_proto_descriptor_pool_builder.h"
 #include "eval/public/structs/cel_proto_wrapper.h"
-#include "eval/public/structs/protobuf_descriptor_type_provider.h"
 #include "eval/public/testing/matchers.h"
 #include "eval/public/unknown_attribute_set.h"
 #include "eval/public/unknown_set.h"
 #include "eval/testutil/test_message.pb.h"
-#include "internal/proto_file_util.h"
 #include "internal/proto_matchers.h"
 #include "internal/status_macros.h"
 #include "internal/testing.h"
 #include "parser/parser.h"
+#include "runtime/function_adapter.h"
 #include "runtime/internal/runtime_env_testing.h"
 #include "runtime/runtime_options.h"
 #include "cel/expr/conformance/proto3/test_all_types.pb.h"
@@ -73,7 +75,9 @@ namespace google::api::expr::runtime {
 
 namespace {
 
+using ::absl_testing::IsOk;
 using ::absl_testing::StatusIs;
+using ::cel::BytesValue;
 using ::cel::Value;
 using ::cel::expr::conformance::proto3::TestAllTypes;
 using ::cel::internal::test::EqualsProto;
@@ -1842,6 +1846,64 @@ TEST(FlatExprBuilderTest, TypeResolve) {
   EXPECT_TRUE(result.BoolOrDie());
 }
 
+TEST(FlatExprBuilderTest, FastEquality) {
+  TestMessage message;
+  ASSERT_OK_AND_ASSIGN(ParsedExpr parsed_expr, parser::Parse("'foo' == 'bar'"));
+  cel::RuntimeOptions options;
+  options.enable_fast_builtins = true;
+  InterpreterOptions legacy_options;
+  legacy_options.enable_fast_builtins = true;
+  CelExpressionBuilderFlatImpl builder(NewTestingRuntimeEnv(), options);
+  ASSERT_THAT(RegisterBuiltinFunctions(builder.GetRegistry(), legacy_options),
+              IsOk());
+  ASSERT_OK_AND_ASSIGN(auto expression,
+                       builder.CreateExpression(&parsed_expr.expr(),
+                                                &parsed_expr.source_info()));
+
+  Activation activation;
+  google::protobuf::Arena arena;
+  ASSERT_OK_AND_ASSIGN(CelValue result,
+                       expression->Evaluate(activation, &arena));
+
+  ASSERT_TRUE(result.IsBool()) << result.DebugString();
+  EXPECT_FALSE(result.BoolOrDie());
+}
+
+TEST(FlatExprBuilderTest, FastEqualityDisabledWithCustomEquality) {
+  TestMessage message;
+  ASSERT_OK_AND_ASSIGN(ParsedExpr parsed_expr, parser::Parse("1 == b'\001'"));
+  cel::RuntimeOptions options;
+  options.enable_fast_builtins = true;
+  InterpreterOptions legacy_options;
+  legacy_options.enable_fast_builtins = true;
+  CelExpressionBuilderFlatImpl builder(NewTestingRuntimeEnv(), options);
+  ASSERT_THAT(RegisterBuiltinFunctions(builder.GetRegistry(), legacy_options),
+              IsOk());
+
+  auto& registry = builder.GetRegistry()->InternalGetRegistry();
+
+  auto status = cel::BinaryFunctionAdapter<bool, int64_t, const BytesValue&>::
+      RegisterGlobalOverload(
+          "_==_",
+          [](auto&, int64_t lhs, const cel::BytesValue& rhs) -> bool {
+            return true;
+          },
+          registry);
+  ASSERT_THAT(status, IsOk());
+
+  ASSERT_OK_AND_ASSIGN(auto expression,
+                       builder.CreateExpression(&parsed_expr.expr(),
+                                                &parsed_expr.source_info()));
+
+  Activation activation;
+  google::protobuf::Arena arena;
+  ASSERT_OK_AND_ASSIGN(CelValue result,
+                       expression->Evaluate(activation, &arena));
+
+  ASSERT_TRUE(result.IsBool()) << result.DebugString();
+  EXPECT_TRUE(result.BoolOrDie());
+}
+
 TEST(FlatExprBuilderTest, AnyPackingList) {
   google::protobuf::LinkMessageReflection<TestAllTypes>();
   ASSERT_OK_AND_ASSIGN(ParsedExpr parsed_expr,

eval/compiler/qualified_reference_resolver.cc

@@ -59,6 +59,14 @@ bool IsSpecialFunction(absl::string_view function_name) {
          function_name == cel::builtin::kIndex ||
          function_name == cel::builtin::kTernary ||
          function_name == kOptionalOr || function_name == kOptionalOrValue ||
+         function_name == cel::builtin::kEqual ||
+         function_name == cel::builtin::kInequal ||
+         function_name == cel::builtin::kNot ||
+         function_name == cel::builtin::kNotStrictlyFalse ||
+         function_name == cel::builtin::kNotStrictlyFalseDeprecated ||
+         function_name == cel::builtin::kIn ||
+         function_name == cel::builtin::kInDeprecated ||
+         function_name == cel::builtin::kInFunction ||
          function_name == "cel.@block";
 }
 

eval/eval/BUILD

@@ -457,6 +457,57 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "equality_steps",
+    srcs = [
+        "equality_steps.cc",
+    ],
+    hdrs = [
+        "equality_steps.h",
+    ],
+    deps = [
+        ":attribute_trail",
+        ":direct_expression_step",
+        ":evaluator_core",
+        ":expression_step_base",
+        "//base:builtins",
+        "//common:value",
+        "//common:value_kind",
+        "//internal:number",
+        "//internal:status_macros",
+        "//runtime/internal:errors",
+        "//runtime/standard:equality_functions",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+    ],
+)
+
+cc_test(
+    name = "equality_steps_test",
+    srcs = [
+        "equality_steps_test.cc",
+    ],
+    deps = [
+        ":attribute_trail",
+        ":direct_expression_step",
+        ":equality_steps",
+        ":evaluator_core",
+        "//base:attributes",
+        "//common:memory",
+        "//common:type",
+        "//common:value",
+        "//common:value_kind",
+        "//common:value_testing",
+        "//internal:testing",
+        "//runtime:activation",
+        "//runtime:runtime_options",
+        "@com_google_absl//absl/log:absl_check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_protobuf//:protobuf",
+    ],
+)
+
 cc_library(
     name = "comprehension_step",
     srcs = [

eval/eval/attribute_utility.cc

@@ -202,8 +202,16 @@ void Accumulator::Add(const UnknownValue& value) {
 void Accumulator::Add(const AttributeTrail& attr) { parent_.Add(*this, attr); }
 
 void Accumulator::MaybeAdd(const Value& v) {
-  if (InstanceOf<UnknownValue>(v)) {
-    Add(Cast<UnknownValue>(v));
+  if (v.IsUnknown()) {
+    Add(v.GetUnknown());
+  }
+}
+
+void Accumulator::MaybeAdd(const Value& v, const AttributeTrail& attr) {
+  if (v.IsUnknown()) {
+    Add(v.GetUnknown());
+  } else if (parent_.CheckForUnknown(attr, /*use_partial=*/true)) {
+    Add(attr);
   }
 }
 

eval/eval/attribute_utility.h

@@ -34,6 +34,15 @@ class AttributeUtility {
     // Add to the accumulated set of unknowns if value is UnknownValue.
     void MaybeAdd(const cel::Value& v);
 
+    // Add to the accumulated set of unknowns if value is UnknownValue or
+    // the attribute trail is (partially) unknown. This version prefers
+    // preserving an already present unknown value over a new one matching the
+    // attribute trail.
+    //
+    // Uses partial matching (a pattern matches the attribute or any
+    // sub-attribute).
+    void MaybeAdd(const cel::Value& v, const AttributeTrail& attr);
+
     bool IsEmpty() const;
 
     cel::UnknownValue Build() &&;