Add set membership optimization using precomputed bitmasks for small integer sets.

PiperOrigin-RevId: 1003462035

Author: jnthntatum

MODULE.bazel

@@ -59,12 +59,12 @@ bazel_dep(
 )
 bazel_dep(
     name = "cel-spec",
-    version = "0.25.2",
+    version = "0.25.3",
     repo_name = "com_google_cel_spec",
 )
 git_override(
     module_name = "cel-spec",
-    commit = "ba58ae5007845f3a1279b488cdeb79645ce958bb",
+    commit = "6dccde40b346f89761ddb890abc6ca503a453134",
     remote = "https://github.com/cel-expr/cel-spec",
 )
 

conformance/BUILD

@@ -53,6 +53,8 @@ cc_library(
         "//extensions:math_ext_macros",
         "//extensions:proto_ext",
         "//extensions:select_optimization",
+        "//extensions:set_membership_optimization",
+        "//extensions:sets_functions",
         "//extensions:strings",
         "//extensions/protobuf:enum_adapter",
         "//internal:status_macros",
@@ -162,6 +164,7 @@ _ALL_TESTS = [
     "@com_google_cel_spec//tests/simple:testdata/unknowns.textproto",
     "@com_google_cel_spec//tests/simple:testdata/wrappers.textproto",
     "@com_google_cel_spec//tests/simple:testdata/block_ext.textproto",
+    "@com_google_cel_spec//tests/simple:testdata/sets_ext.textproto",
     "@com_google_cel_spec//tests/simple:testdata/type_deduction.textproto",
 ]
 
@@ -286,17 +289,17 @@ gen_conformance_tests(
     name = "conformance_legacy_select_opt",
     checked = True,
     data = _ALL_TESTS,
+    extended_optimizations = True,
     modern = False,
-    select_opt = True,
     skip_tests = _TESTS_TO_SKIP_LEGACY,
 )
 
 gen_conformance_tests(
     name = "conformance_select_opt",
     checked = True,
     data = _ALL_TESTS,
+    extended_optimizations = True,
     modern = True,
-    select_opt = True,
     skip_tests = _TESTS_TO_SKIP_MODERN,
 )
 

conformance/run.bzl

@@ -56,14 +56,14 @@ def _conformance_test_name(name, pratt, optimize):
         ],
     )
 
-def _conformance_test_args(modern, optimize, select_opt, skip_check, dashboard, enable_variadic_logical_operators, pratt):
+def _conformance_test_args(modern, optimize, extended_optimizations, skip_check, dashboard, enable_variadic_logical_operators, pratt):
     args = []
     if modern:
         args.append("--modern")
     if optimize:
         args.append("--opt")
-    if select_opt:
-        args.append("--select_optimization")
+    if extended_optimizations:
+        args.append("--extended_optimizations")
     if skip_check:
         args.append("--skip_check")
     else:
@@ -78,10 +78,10 @@ def _conformance_test_args(modern, optimize, select_opt, skip_check, dashboard,
         args.append("--noenable_pratt_parser")
     return args
 
-def _conformance_test(name, data, modern, optimize, select_opt, skip_check, skip_tests, tags, dashboard, enable_variadic_logical_operators, pratt):
+def _conformance_test(name, data, modern, optimize, extended_optimizations, skip_check, skip_tests, tags, dashboard, enable_variadic_logical_operators, pratt):
     cc_test(
         name = _conformance_test_name(name, pratt, optimize),
-        args = _conformance_test_args(modern, optimize, select_opt, skip_check, dashboard, enable_variadic_logical_operators, pratt) + ["$(rlocationpath {})".format(test) for test in data],
+        args = _conformance_test_args(modern, optimize, extended_optimizations, skip_check, dashboard, enable_variadic_logical_operators, pratt) + ["$(rlocationpath {})".format(test) for test in data],
         env = select(
             {
                 "@platforms//os:windows": {"CEL_SKIP_TESTS": ",".join(skip_tests + _TESTS_TO_SKIP_WINDOWS)},
@@ -93,15 +93,15 @@ def _conformance_test(name, data, modern, optimize, select_opt, skip_check, skip
         tags = tags,
     )
 
-def gen_conformance_tests(name, data, modern = False, checked = False, select_opt = False, dashboard = False, skip_tests = [], tags = [], enable_variadic_logical_operators = False):
+def gen_conformance_tests(name, data, modern = False, checked = False, extended_optimizations = False, dashboard = False, skip_tests = [], tags = [], enable_variadic_logical_operators = False):
     """Generates conformance tests.
 
     Args:
         name: prefix for all tests
         data: textproto targets describing conformance tests
         modern: run using modern APIs
         checked: whether to apply type checking
-        select_opt: enable select optimization
+        extended_optimizations: enable extended optimizations (select and set membership)
         dashboard: enable dashboard mode
         skip_tests: tests to skip in the format of the cel-spec test runner. See documentation
             in github.com/google/cel-spec/tests/simple/simple_test.go
@@ -119,7 +119,7 @@ def gen_conformance_tests(name, data, modern = False, checked = False, select_op
                 data,
                 modern = modern,
                 optimize = optimize,
-                select_opt = select_opt,
+                extended_optimizations = extended_optimizations,
                 skip_check = skip_check,
                 skip_tests = _expand_tests_to_skip(skip_tests),
                 tags = tags,

conformance/run.cc

@@ -62,7 +62,8 @@ ABSL_FLAG(
 ABSL_FLAG(std::vector<std::string>, skip_tests, {}, "Tests to skip");
 ABSL_FLAG(bool, dashboard, false, "Dashboard mode, ignore test failures");
 ABSL_FLAG(bool, skip_check, true, "Skip type checking the expressions");
-ABSL_FLAG(bool, select_optimization, false, "Enable select optimization.");
+ABSL_FLAG(bool, extended_optimizations, false,
+          "Enable extended optimizations.");
 ABSL_FLAG(bool, enable_variadic_logical_operators, false,
           "Enable parsing logical AND & OR operators as a single flat variadic "
           "call.");
@@ -261,7 +262,7 @@ NewConformanceServiceFromFlags() {
       cel_conformance::ConformanceServiceOptions{
           .optimize = absl::GetFlag(FLAGS_opt),
           .modern = absl::GetFlag(FLAGS_modern),
-          .select_optimization = absl::GetFlag(FLAGS_select_optimization),
+          .extended_optimizations = absl::GetFlag(FLAGS_extended_optimizations),
           .enable_variadic_logical_operators =
               absl::GetFlag(FLAGS_enable_variadic_logical_operators),
           .enable_pratt_parser = absl::GetFlag(FLAGS_enable_pratt_parser),

conformance/service.cc

@@ -67,6 +67,8 @@
 #include "extensions/proto_ext.h"
 #include "extensions/protobuf/enum_adapter.h"
 #include "extensions/select_optimization.h"
+#include "extensions/set_membership_optimization.h"
+#include "extensions/sets_functions.h"
 #include "extensions/strings.h"
 #include "internal/status_macros.h"
 #include "parser/macro_registry.h"
@@ -196,6 +198,8 @@ absl::Status CheckImpl(google::protobuf::Arena* arena,
         builder->AddLibrary(cel::extensions::ComprehensionsV2CheckerLibrary()));
     CEL_RETURN_IF_ERROR(
         builder->AddLibrary(cel::extensions::ListsCheckerLibrary()));
+    CEL_RETURN_IF_ERROR(
+        builder->AddLibrary(cel::extensions::SetsCheckerLibrary()));
   }
 
   for (const auto& decl : request.type_env()) {
@@ -245,7 +249,7 @@ absl::Status CheckImpl(google::protobuf::Arena* arena,
 class LegacyConformanceServiceImpl : public ConformanceServiceInterface {
  public:
   static absl::StatusOr<std::unique_ptr<LegacyConformanceServiceImpl>> Create(
-      bool optimize, bool recursive, bool select_optimization,
+      bool optimize, bool recursive, bool extended_optimizations,
       bool enable_variadic_logical_operators, bool enable_pratt_parser) {
     static auto* constant_arena = new Arena();
 
@@ -293,8 +297,8 @@ class LegacyConformanceServiceImpl : public ConformanceServiceInterface {
       options.enable_typed_field_access = true;
     }
 
-    if (select_optimization) {
-      std::cerr << "Enabling select optimizations" << std::endl;
+    if (extended_optimizations) {
+      std::cerr << "Enabling extended optimizations" << std::endl;
       options.enable_select_optimization = true;
     }
 
@@ -323,6 +327,8 @@ class LegacyConformanceServiceImpl : public ConformanceServiceInterface {
         builder->GetRegistry(), options));
     CEL_RETURN_IF_ERROR(cel::extensions::RegisterMathExtensionFunctions(
         builder->GetRegistry(), options));
+    CEL_RETURN_IF_ERROR(cel::extensions::RegisterSetsFunctions(
+        builder->GetRegistry(), options));
 
     return absl::WrapUnique(new LegacyConformanceServiceImpl(
         std::move(builder), enable_variadic_logical_operators,
@@ -446,7 +452,7 @@ class LegacyConformanceServiceImpl : public ConformanceServiceInterface {
 class ModernConformanceServiceImpl : public ConformanceServiceInterface {
  public:
   static absl::StatusOr<std::unique_ptr<ModernConformanceServiceImpl>> Create(
-      bool optimize, bool recursive, bool select_optimization,
+      bool optimize, bool recursive, bool extended_optimizations,
       bool enable_variadic_logical_operators, bool enable_pratt_parser) {
     google::protobuf::LinkMessageReflection<
         cel::expr::conformance::proto3::TestAllTypes>();
@@ -490,7 +496,7 @@ class ModernConformanceServiceImpl : public ConformanceServiceInterface {
     }
 
     return absl::WrapUnique(new ModernConformanceServiceImpl(
-        options, optimize, select_optimization,
+        options, optimize, extended_optimizations,
         enable_variadic_logical_operators, enable_pratt_parser));
   }
 
@@ -512,8 +518,10 @@ class ModernConformanceServiceImpl : public ConformanceServiceInterface {
     }
     CEL_RETURN_IF_ERROR(cel::EnableReferenceResolver(
         builder, cel::ReferenceResolverEnabled::kAlways));
-    if (enable_select_optimization_) {
+    if (enable_extended_optimizations_) {
       CEL_RETURN_IF_ERROR(cel::extensions::EnableSelectOptimization(builder));
+      CEL_RETURN_IF_ERROR(
+          cel::extensions::EnableSetMembershipOptimization(builder));
     }
 
     auto& type_registry = builder.type_registry();
@@ -542,6 +550,8 @@ class ModernConformanceServiceImpl : public ConformanceServiceInterface {
         builder.function_registry(), options));
     CEL_RETURN_IF_ERROR(cel::extensions::RegisterListsFunctions(
         builder.function_registry(), options));
+    CEL_RETURN_IF_ERROR(cel::extensions::RegisterSetsFunctions(
+        builder.function_registry(), options));
 
     return std::move(builder).Build();
   }
@@ -641,12 +651,12 @@ class ModernConformanceServiceImpl : public ConformanceServiceInterface {
  private:
   ModernConformanceServiceImpl(const RuntimeOptions& options,
                                bool enable_optimizations,
-                               bool enable_select_optimization,
+                               bool enable_extended_optimizations,
                                bool enable_variadic_logical_operators,
                                bool enable_pratt_parser)
       : options_(options),
         enable_optimizations_(enable_optimizations),
-        enable_select_optimization_(enable_select_optimization),
+        enable_extended_optimizations_(enable_extended_optimizations),
         enable_variadic_logical_operators_(enable_variadic_logical_operators),
         enable_pratt_parser_(enable_pratt_parser) {}
 
@@ -678,7 +688,7 @@ class ModernConformanceServiceImpl : public ConformanceServiceInterface {
 
   RuntimeOptions options_;
   bool enable_optimizations_;
-  bool enable_select_optimization_;
+  bool enable_extended_optimizations_;
   bool enable_variadic_logical_operators_;
   bool enable_pratt_parser_;
 };
@@ -693,11 +703,11 @@ absl::StatusOr<std::unique_ptr<ConformanceServiceInterface>>
 NewConformanceService(const ConformanceServiceOptions& options) {
   if (options.modern) {
     return google::api::expr::runtime::ModernConformanceServiceImpl::Create(
-        options.optimize, options.recursive, options.select_optimization,
+        options.optimize, options.recursive, options.extended_optimizations,
         options.enable_variadic_logical_operators, options.enable_pratt_parser);
   } else {
     return google::api::expr::runtime::LegacyConformanceServiceImpl::Create(
-        options.optimize, options.recursive, options.select_optimization,
+        options.optimize, options.recursive, options.extended_optimizations,
         options.enable_variadic_logical_operators, options.enable_pratt_parser);
   }
 }

conformance/service.h

@@ -45,7 +45,7 @@ struct ConformanceServiceOptions {
   bool modern;
   bool arena;
   bool recursive;
-  bool select_optimization;
+  bool extended_optimizations;
   bool enable_variadic_logical_operators = false;
   bool enable_pratt_parser = true;
 };

eval/compiler/flat_expr_builder_extensions.cc

@@ -303,8 +303,15 @@ Subexpression* absl_nullable ProgramBuilder::MakeSubexpression(
   return it->second.get();
 }
 
-bool PlannerContext::IsSubplanInspectable(const cel::Expr& node) const {
-  return program_builder_.GetSubexpression(&node) != nullptr;
+bool PlannerContext::NodeHasInspectablePlan(const cel::Expr& node) const {
+  const Subexpression* subexpression = program_builder_.GetSubexpression(&node);
+  if (subexpression == nullptr) {
+    return false;
+  }
+  if (subexpression->IsFlattened()) {
+    return !subexpression->flattened_elements().empty();
+  }
+  return !subexpression->elements().empty();
 }
 
 ExecutionPathView PlannerContext::GetSubplan(const cel::Expr& node) {

eval/compiler/flat_expr_builder_extensions.h

@@ -308,7 +308,7 @@ class PlannerContext {
   // Returns true if the subplan is inspectable.
   //
   // If false, the node is not mapped to a subexpression in the program builder.
-  bool IsSubplanInspectable(const cel::Expr& node) const;
+  bool NodeHasInspectablePlan(const cel::Expr& node) const;
 
   // Return a view to the current subplan representing node.
   //

eval/compiler/flat_expr_builder_extensions_test.cc

@@ -144,7 +144,7 @@ TEST_F(PlannerContextTest, GetPlan) {
                                                  UniquePtrHolds(step_ptrs.a)));
 
   Expr d;
-  EXPECT_FALSE(context.IsSubplanInspectable(d));
+  EXPECT_FALSE(context.NodeHasInspectablePlan(d));
   EXPECT_THAT(context.GetSubplan(d), IsEmpty());
 }
 
@@ -193,8 +193,8 @@ TEST_F(PlannerContextTest, ExtractPlan) {
                          type_registry_.GetComposedTypeProvider(),
                          issue_collector_, program_builder, arena);
 
-  EXPECT_TRUE(context.IsSubplanInspectable(a));
-  EXPECT_TRUE(context.IsSubplanInspectable(b));
+  EXPECT_TRUE(context.NodeHasInspectablePlan(a));
+  EXPECT_TRUE(context.NodeHasInspectablePlan(b));
 
   ASSERT_OK_AND_ASSIGN(ExecutionPath extracted, context.ExtractSubplan(b));
 
@@ -233,7 +233,7 @@ TEST_F(PlannerContextTest, ReplacePlanUpdatesParent) {
                          type_registry_.GetComposedTypeProvider(),
                          issue_collector_, program_builder, arena);
 
-  EXPECT_TRUE(context.IsSubplanInspectable(a));
+  EXPECT_TRUE(context.NodeHasInspectablePlan(a));
 
   ASSERT_THAT(context.ReplaceSubplan(c, {}), IsOk());
 

eval/eval/BUILD

@@ -35,6 +35,7 @@ package_group(
 cc_library(
     name = "evaluator_core",
     srcs = [
+        "bitset_steps.cc",
         "comprehension_step.cc",
         "create_list_step.cc",
         "equality_steps.cc",
@@ -45,6 +46,7 @@ cc_library(
         "logic_step.cc",
     ],
     hdrs = [
+        "bitset_steps.h",
         "comprehension_step.h",
         "create_list_step.h",
         "equality_steps.h",
@@ -356,6 +358,40 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "bitset_steps",
+    deps = [
+        ":evaluator_core",
+    ],
+)
+
+cc_test(
+    name = "bitset_steps_test",
+    srcs = [
+        "bitset_steps_test.cc",
+    ],
+    deps = [
+        ":attribute_trail",
+        ":evaluator_core",
+        ":expression_step_logic",
+        "//base:attributes",
+        "//common:value",
+        "//common:value_kind",
+        "//common:value_testing",
+        "//internal:testing",
+        "//internal:testing_descriptor_pool",
+        "//internal:testing_message_factory",
+        "//runtime:activation",
+        "//runtime:runtime_options",
+        "//runtime/internal:runtime_type_provider",
+        "@com_google_absl//absl/base:nullability",
+        "@com_google_absl//absl/log:absl_check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_protobuf//:protobuf",
+    ],
+)
+
 cc_test(
     name = "equality_steps_test",
     srcs = [