Add conformance tests with select optimization enabled.

Update the optimizer to permit invalid constant map keys (by not rewriting the select chain when encountered). These should not be allowed per spec, but the CEL implementations generally permit it to tolerate `dyn` keys.

PiperOrigin-RevId: 839852761

Author: jnthntatum

conformance/BUILD

@@ -53,6 +53,7 @@ cc_library(
         "//extensions:math_ext_decls",
         "//extensions:math_ext_macros",
         "//extensions:proto_ext",
+        "//extensions:select_optimization",
         "//extensions:strings",
         "//extensions/protobuf:enum_adapter",
         "//internal:status_macros",
@@ -297,6 +298,25 @@ gen_conformance_tests(
     skip_tests = _TESTS_TO_SKIP_LEGACY + _TESTS_TO_SKIP_CHECKED,
 )
 
+# select optimization is only supported for checked expressions.
+gen_conformance_tests(
+    name = "conformance_legacy_select_opt",
+    checked = True,
+    data = _ALL_TESTS,
+    modern = False,
+    select_opt = True,
+    skip_tests = _TESTS_TO_SKIP_LEGACY + _TESTS_TO_SKIP_CHECKED,
+)
+
+gen_conformance_tests(
+    name = "conformance_select_opt",
+    checked = True,
+    data = _ALL_TESTS,
+    modern = True,
+    select_opt = True,
+    skip_tests = _TESTS_TO_SKIP_MODERN + _TESTS_TO_SKIP_CHECKED,
+)
+
 # Generates a bunch of `cc_test` whose names follow the pattern
 # `conformance_dashboard_..._{arena|refcount}_{optimized|unoptimized}_{recursive|iterative}`.
 gen_conformance_tests(

conformance/run.bzl

@@ -46,12 +46,14 @@ def _conformance_test_name(name, optimize, recursive):
         ],
     )
 
-def _conformance_test_args(modern, optimize, recursive, skip_check, skip_tests, dashboard):
+def _conformance_test_args(modern, optimize, recursive, select_opt, skip_check, skip_tests, dashboard):
     args = []
     if modern:
         args.append("--modern")
     if optimize:
         args.append("--opt")
+    if select_opt:
+        args.append("--select_optimization")
     if recursive:
         args.append("--recursive")
     if skip_check:
@@ -63,16 +65,16 @@ def _conformance_test_args(modern, optimize, recursive, skip_check, skip_tests,
         args.append("--dashboard")
     return args
 
-def _conformance_test(name, data, modern, optimize, recursive, skip_check, skip_tests, tags, dashboard):
+def _conformance_test(name, data, modern, optimize, recursive, select_opt, skip_check, skip_tests, tags, dashboard):
     cc_test(
         name = _conformance_test_name(name, optimize, recursive),
-        args = _conformance_test_args(modern, optimize, recursive, skip_check, skip_tests, dashboard) + ["$(location " + test + ")" for test in data],
+        args = _conformance_test_args(modern, optimize, recursive, select_opt, skip_check, skip_tests, dashboard) + ["$(location " + test + ")" for test in data],
         data = data,
         deps = ["//conformance:run"],
         tags = tags,
     )
 
-def gen_conformance_tests(name, data, modern = False, checked = False, dashboard = False, skip_tests = [], tags = []):
+def gen_conformance_tests(name, data, modern = False, checked = False, select_opt = False, dashboard = False, skip_tests = [], tags = []):
     """Generates conformance tests.
 
     Args:
@@ -97,6 +99,7 @@ def gen_conformance_tests(name, data, modern = False, checked = False, dashboard
                 modern = modern,
                 optimize = optimize,
                 recursive = recursive,
+                select_opt = select_opt,
                 skip_check = skip_check,
                 skip_tests = skip_tests,
                 tags = tags,

conformance/run.cc

@@ -63,6 +63,7 @@ ABSL_FLAG(bool, recursive, false,
 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.");
 
 namespace {
 
@@ -257,7 +258,9 @@ NewConformanceServiceFromFlags() {
       cel_conformance::ConformanceServiceOptions{
           .optimize = absl::GetFlag(FLAGS_opt),
           .modern = absl::GetFlag(FLAGS_modern),
-          .recursive = absl::GetFlag(FLAGS_recursive)});
+          .recursive = absl::GetFlag(FLAGS_recursive),
+          .select_optimization = absl::GetFlag(FLAGS_select_optimization),
+      });
   ABSL_CHECK_OK(status_or_service);
   return std::shared_ptr<cel_conformance::ConformanceServiceInterface>(
       std::move(*status_or_service));

conformance/service.cc

@@ -47,12 +47,10 @@
 #include "checker/type_checker_builder_factory.h"
 #include "common/ast.h"
 #include "common/ast_proto.h"
-#include "common/decl.h"
 #include "common/decl_proto_v1alpha1.h"
 #include "common/expr.h"
 #include "common/internal/value_conversion.h"
 #include "common/source.h"
-#include "common/type.h"
 #include "common/value.h"
 #include "eval/public/activation.h"
 #include "eval/public/builtin_func_registrar.h"
@@ -70,6 +68,7 @@
 #include "extensions/math_ext_macros.h"
 #include "extensions/proto_ext.h"
 #include "extensions/protobuf/enum_adapter.h"
+#include "extensions/select_optimization.h"
 #include "extensions/strings.h"
 #include "internal/status_macros.h"
 #include "parser/macro.h"
@@ -340,7 +339,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 optimize, bool recursive, bool select_optimization) {
     static auto* constant_arena = new Arena();
 
     google::protobuf::LinkMessageReflection<
@@ -385,6 +384,11 @@ class LegacyConformanceServiceImpl : public ConformanceServiceInterface {
       options.constant_arena = constant_arena;
     }
 
+    if (select_optimization) {
+      std::cerr << "Enabling select optimizations" << std::endl;
+      options.enable_select_optimization = true;
+    }
+
     if (recursive) {
       options.max_recursion_depth = 48;
     }
@@ -526,7 +530,7 @@ class LegacyConformanceServiceImpl : public ConformanceServiceInterface {
 class ModernConformanceServiceImpl : public ConformanceServiceInterface {
  public:
   static absl::StatusOr<std::unique_ptr<ModernConformanceServiceImpl>> Create(
-      bool optimize, bool recursive) {
+      bool optimize, bool recursive, bool select_optimization) {
     google::protobuf::LinkMessageReflection<
         cel::expr::conformance::proto3::TestAllTypes>();
     google::protobuf::LinkMessageReflection<
@@ -565,8 +569,8 @@ class ModernConformanceServiceImpl : public ConformanceServiceInterface {
       options.max_recursion_depth = 48;
     }
 
-    return absl::WrapUnique(
-        new ModernConformanceServiceImpl(options, optimize));
+    return absl::WrapUnique(new ModernConformanceServiceImpl(
+        options, optimize, select_optimization));
   }
 
   absl::StatusOr<std::unique_ptr<const cel::Runtime>> Setup(
@@ -583,6 +587,9 @@ class ModernConformanceServiceImpl : public ConformanceServiceInterface {
     }
     CEL_RETURN_IF_ERROR(cel::EnableReferenceResolver(
         builder, cel::ReferenceResolverEnabled::kAlways));
+    if (enable_select_optimization_) {
+      CEL_RETURN_IF_ERROR(cel::extensions::EnableSelectOptimization(builder));
+    }
 
     auto& type_registry = builder.type_registry();
     // Use linked pbs in the generated descriptor pool.
@@ -704,10 +711,12 @@ class ModernConformanceServiceImpl : public ConformanceServiceInterface {
   }
 
  private:
-  explicit ModernConformanceServiceImpl(const RuntimeOptions& options,
-                                        bool enable_optimizations)
-      : options_(options), enable_optimizations_(enable_optimizations) {}
-
+  ModernConformanceServiceImpl(const RuntimeOptions& options,
+                               bool enable_optimizations,
+                               bool enable_select_optimization)
+      : options_(options),
+        enable_optimizations_(enable_optimizations),
+        enable_select_optimization_(enable_select_optimization) {}
 
   static absl::StatusOr<std::unique_ptr<cel::TraceableProgram>> Plan(
       const cel::Runtime& runtime,
@@ -737,6 +746,7 @@ class ModernConformanceServiceImpl : public ConformanceServiceInterface {
 
   RuntimeOptions options_;
   bool enable_optimizations_;
+  bool enable_select_optimization_;
 };
 
 }  // namespace
@@ -749,10 +759,10 @@ 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.optimize, options.recursive, options.select_optimization);
   } else {
     return google::api::expr::runtime::LegacyConformanceServiceImpl::Create(
-        options.optimize, options.recursive);
+        options.optimize, options.recursive, options.select_optimization);
   }
 }
 

conformance/service.h

@@ -45,6 +45,7 @@ struct ConformanceServiceOptions {
   bool modern;
   bool arena;
   bool recursive;
+  bool select_optimization;
 };
 
 absl::StatusOr<std::unique_ptr<ConformanceServiceInterface>>

eval/public/cel_options.h

@@ -163,9 +163,6 @@ struct InterpreterOptions {
   //    always operates as though it is `true`.
   //  - `enable_heterogeneous_equality` is ignored and optimized traversals
   //    always operate as though it is `true`.
-  //
-  // Note: implementation in progress -- please consult the CEL team before
-  // enabling in an existing environment.
   bool enable_select_optimization = false;
 
   // Enable lazy cel.bind alias initialization.

extensions/BUILD

@@ -329,6 +329,7 @@ cc_library(
         "//eval/eval:evaluator_core",
         "//eval/eval:expression_step_base",
         "//internal:casts",
+        "//internal:number",
         "//internal:status_macros",
         "//runtime:runtime_builder",
         "//runtime/internal:errors",

extensions/select_optimization.cc

@@ -53,6 +53,7 @@
 #include "eval/eval/evaluator_core.h"
 #include "eval/eval/expression_step_base.h"
 #include "internal/casts.h"
+#include "internal/number.h"
 #include "internal/status_macros.h"
 #include "runtime/internal/errors.h"
 #include "runtime/internal/runtime_friend_access.h"
@@ -188,34 +189,45 @@ absl::StatusOr<SelectQualifier> SelectQualifierFromList(const ListExpr& list) {
                         field_name.const_expr().string_value()};
 }
 
+// Returns a qualifier instruction derived from a unoptimized ast.
 absl::StatusOr<QualifierInstruction> SelectInstructionFromConstant(
     const Constant& constant) {
-  if (constant.has_int64_value()) {
-    return QualifierInstruction(constant.int64_value());
-  } else if (constant.has_uint64_value()) {
-    return QualifierInstruction(constant.uint64_value());
+  if (constant.has_int_value()) {
+    return QualifierInstruction(constant.int_value());
+  } else if (constant.has_uint_value()) {
+    return QualifierInstruction(constant.uint_value());
   } else if (constant.has_bool_value()) {
     return QualifierInstruction(constant.bool_value());
   } else if (constant.has_string_value()) {
     return QualifierInstruction(constant.string_value());
+  } else if (constant.has_double_value()) {
+    cel::internal::Number number(constant.double_value());
+    if (number.LosslessConvertibleToInt()) {
+      return QualifierInstruction(number.AsInt());
+    } else if (number.LosslessConvertibleToUint()) {
+      return QualifierInstruction(number.AsUint());
+    }
   }
 
-  return absl::InvalidArgumentError("Invalid cel.attribute constant");
+  return absl::InvalidArgumentError("invalid index constant for cel.attribute");
 }
 
 absl::StatusOr<SelectQualifier> SelectQualifierFromConstant(
     const Constant& constant) {
-  if (constant.has_int64_value()) {
-    return AttributeQualifier::OfInt(constant.int64_value());
-  } else if (constant.has_uint64_value()) {
-    return AttributeQualifier::OfUint(constant.uint64_value());
+  if (constant.has_int_value()) {
+    return AttributeQualifier::OfInt(constant.int_value());
+  } else if (constant.has_uint_value()) {
+    return AttributeQualifier::OfUint(constant.uint_value());
   } else if (constant.has_bool_value()) {
     return AttributeQualifier::OfBool(constant.bool_value());
   } else if (constant.has_string_value()) {
     return AttributeQualifier::OfString(constant.string_value());
   }
+  // TODO(uncreated-issue/51): double keys could possibly be valid selectors, but
+  // the other stacks don't implement the optimization yet and we normalize the
+  // key to a uint or int if we do the late AST rewrite during planning.
 
-  return absl::InvalidArgumentError("Invalid cel.attribute constant");
+  return absl::InvalidArgumentError("invalid cel.attribute constant");
 }
 
 absl::StatusOr<size_t> ListIndexFromQualifier(const AttributeQualifier& qual) {
@@ -248,7 +260,7 @@ absl::StatusOr<Value> MapKeyFromQualifier(const AttributeQualifier& qual,
     case Kind::kBool:
       return cel::BoolValue(*qual.GetBoolKey());
     case Kind::kString:
-      return cel::StringValue(arena, *qual.GetStringKey());
+      return StringValue::From(*qual.GetStringKey(), arena);
     default:
       return runtime_internal::CreateNoMatchingOverloadError(
           cel::builtin::kIndex);
@@ -424,7 +436,8 @@ class RewriterImpl : public AstRewriterBase {
       auto qualifier_or =
           SelectInstructionFromConstant(qualifier_expr.const_expr());
       if (!qualifier_or.ok()) {
-        SetProgressStatus(qualifier_or.status());
+        // TODO(uncreated-issue/54): should warn, but by default warnings fail overall
+        // program planning.
         return;
       }
       candidates_[&expr] = std::move(qualifier_or).value();

extensions/select_optimization.h

@@ -54,8 +54,8 @@ struct SelectOptimizationOptions {
 // Assumes the default runtime implementation, an error with code
 // InvalidArgument is returned if it is not.
 //
-// Note: implementation in progress -- please consult the CEL team before
-// enabling in an existing environment.
+// Note: implementation does not support optional field traversal, and will
+// instead revert to the normal implementation instead of trying to optimize.
 absl::Status EnableSelectOptimization(
     cel::RuntimeBuilder& builder,
     const SelectOptimizationOptions& options = {});