Rename @cel.attribute to cel.@Attribute to be consistent with other c…

…el-namespaced internal functions.

PiperOrigin-RevId: 641295530

eval/public/portable_cel_expr_builder_factory.cc

@@ -46,6 +46,8 @@ namespace {
 using ::cel::MemoryManagerRef;
 using ::cel::ast_internal::AstImpl;
 using ::cel::extensions::CreateSelectOptimizationProgramOptimizer;
+using ::cel::extensions::kCelAttribute;
+using ::cel::extensions::kCelHasField;
 using ::cel::extensions::ProtoMemoryManagerRef;
 using ::cel::extensions::SelectOptimizationAstUpdater;
 using ::cel::runtime_internal::CreateConstantFoldingOptimizer;
@@ -108,17 +110,18 @@ std::unique_ptr<CelExpressionBuilder> CreatePortableExprBuilder(
     // Add overloads for select optimization signature.
     // These are never bound, only used to prevent the builder from failing on
     // the overloads check.
-    absl::Status status = builder->GetRegistry()->RegisterLazyFunction(
-        CelFunctionDescriptor(cel::extensions::kCelAttribute, false,
-                              {cel::Kind::kAny, cel::Kind::kList}));
+    absl::Status status =
+        builder->GetRegistry()->RegisterLazyFunction(CelFunctionDescriptor(
+            kCelAttribute, false, {cel::Kind::kAny, cel::Kind::kList}));
     if (!status.ok()) {
-      ABSL_LOG(ERROR) << "Failed to register @cel.attribute: " << status;
+      ABSL_LOG(ERROR) << "Failed to register " << kCelAttribute << ": "
+                      << status;
     }
-    status = builder->GetRegistry()->RegisterLazyFunction(
-        CelFunctionDescriptor(cel::extensions::kFieldsHas, false,
-                              {cel::Kind::kAny, cel::Kind::kList}));
+    status = builder->GetRegistry()->RegisterLazyFunction(CelFunctionDescriptor(
+        kCelHasField, false, {cel::Kind::kAny, cel::Kind::kList}));
     if (!status.ok()) {
-      ABSL_LOG(ERROR) << "Failed to register @cel.hasField: " << status;
+      ABSL_LOG(ERROR) << "Failed to register " << kCelHasField << ": "
+                      << status;
     }
     // Add runtime implementation.
     flat_expr_builder.AddProgramOptimizer(

extensions/BUILD

@@ -221,20 +221,28 @@ cc_library(
     deps = [
         "//base:attributes",
         "//base:builtins",
+        "//base:function_descriptor",
         "//base/ast_internal:ast_impl",
         "//base/ast_internal:expr",
         "//common:ast_rewrite",
         "//common:casting",
+        "//common:expr",
         "//common:kind",
+        "//common:native_type",
         "//common:type",
         "//common:value",
+        "//eval/compiler:flat_expr_builder",
         "//eval/compiler:flat_expr_builder_extensions",
         "//eval/eval:attribute_trail",
         "//eval/eval:direct_expression_step",
         "//eval/eval:evaluator_core",
         "//eval/eval:expression_step_base",
+        "//internal:casts",
         "//internal:status_macros",
+        "//runtime:runtime_builder",
         "//runtime/internal:errors",
+        "//runtime/internal:runtime_friend_access",
+        "//runtime/internal:runtime_impl",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/functional:overload",

extensions/select_optimization.cc

@@ -37,19 +37,27 @@
 #include "base/ast_internal/expr.h"
 #include "base/attribute.h"
 #include "base/builtins.h"
+#include "base/function_descriptor.h"
 #include "common/ast_rewrite.h"
 #include "common/casting.h"
+#include "common/expr.h"
 #include "common/kind.h"
+#include "common/native_type.h"
 #include "common/type.h"
 #include "common/value.h"
 #include "common/value_manager.h"
+#include "eval/compiler/flat_expr_builder.h"
 #include "eval/compiler/flat_expr_builder_extensions.h"
 #include "eval/eval/attribute_trail.h"
 #include "eval/eval/direct_expression_step.h"
 #include "eval/eval/evaluator_core.h"
 #include "eval/eval/expression_step_base.h"
+#include "internal/casts.h"
 #include "internal/status_macros.h"
 #include "runtime/internal/errors.h"
+#include "runtime/internal/runtime_friend_access.h"
+#include "runtime/internal/runtime_impl.h"
+#include "runtime/runtime_builder.h"
 
 namespace cel::extensions {
 namespace {
@@ -90,7 +98,7 @@ struct SelectPath {
 };
 
 // Generates the AST representation of the qualification path for the optimized
-// select branch. I.e., the list-typed second argument of the @cel.attribute
+// select branch. I.e., the list-typed second argument of the cel.@attribute
 // call.
 Expr MakeSelectPathExpr(
     const std::vector<QualifierInstruction>& select_instructions) {
@@ -439,7 +447,7 @@ class RewriterImpl : public AstRewriterBase {
     SelectPath path = GetSelectPath(&expr);
 
     // generate the new cel.attribute call.
-    absl::string_view fn = path.test_only ? kFieldsHas : kCelAttribute;
+    absl::string_view fn = path.test_only ? kCelHasField : kCelAttribute;
 
     Expr operand(std::move(*path.operand));
     Expr call;
@@ -779,7 +787,7 @@ absl::Status SelectOptimizer::OnPostVisit(PlannerContext& context,
   }
 
   absl::string_view fn = node.call_expr().function();
-  if (fn != kFieldsHas && fn != kCelAttribute) {
+  if (fn != kCelHasField && fn != kCelAttribute) {
     return absl::OkStatus();
   }
 
@@ -801,7 +809,7 @@ absl::Status SelectOptimizer::OnPostVisit(PlannerContext& context,
 
   bool presence_test = false;
 
-  if (fn == kFieldsHas) {
+  if (fn == kCelHasField) {
     presence_test = true;
   }
 
@@ -867,6 +875,19 @@ absl::Status SelectOptimizer::OnPostVisit(PlannerContext& context,
   return context.ReplaceSubplan(node, std::move(path));
 }
 
+google::api::expr::runtime::FlatExprBuilder* GetFlatExprBuilder(
+    RuntimeBuilder& builder) {
+  auto& runtime =
+      runtime_internal::RuntimeFriendAccess::GetMutableRuntime(builder);
+  if (runtime_internal::RuntimeFriendAccess::RuntimeTypeId(runtime) ==
+      NativeTypeId::For<runtime_internal::RuntimeImpl>()) {
+    auto& runtime_impl =
+        cel::internal::down_cast<runtime_internal::RuntimeImpl&>(runtime);
+    return &runtime_impl.expr_builder();
+  }
+  return nullptr;
+}
+
 }  // namespace
 
 absl::Status SelectOptimizationAstUpdater::UpdateAst(PlannerContext& context,
@@ -884,4 +905,28 @@ CreateSelectOptimizationProgramOptimizer(
   };
 }
 
+absl::Status EnableSelectOptimization(
+    cel::RuntimeBuilder& builder, const SelectOptimizationOptions& options) {
+  auto* flat_expr_builder = GetFlatExprBuilder(builder);
+  if (flat_expr_builder == nullptr) {
+    return absl::InvalidArgumentError(
+        "SelectOptimization requires default runtime implementation");
+  }
+
+  flat_expr_builder->AddAstTransform(
+      std::make_unique<SelectOptimizationAstUpdater>());
+  // Add overloads for select optimization signature.
+  // These are never bound, only used to prevent the builder from failing on
+  // the overloads check.
+  CEL_RETURN_IF_ERROR(builder.function_registry().RegisterLazyFunction(
+      FunctionDescriptor(kCelAttribute, false, {Kind::kAny, Kind::kList})));
+
+  CEL_RETURN_IF_ERROR(builder.function_registry().RegisterLazyFunction(
+      FunctionDescriptor(kCelHasField, false, {Kind::kAny, Kind::kList})));
+  // Add runtime implementation.
+  flat_expr_builder->AddProgramOptimizer(
+      CreateSelectOptimizationProgramOptimizer(options));
+  return absl::OkStatus();
+}
+
 }  // namespace cel::extensions

extensions/select_optimization.h

@@ -18,11 +18,12 @@
 #include "absl/status/status.h"
 #include "base/ast_internal/ast_impl.h"
 #include "eval/compiler/flat_expr_builder_extensions.h"
+#include "runtime/runtime_builder.h"
 
 namespace cel::extensions {
 
-constexpr char kCelAttribute[] = "@cel.attribute";
-constexpr char kFieldsHas[] = "@cel.hasField";
+constexpr char kCelAttribute[] = "cel.@attribute";
+constexpr char kCelHasField[] = "cel.@hasField";
 
 // Configuration options for the select optimization.
 struct SelectOptimizationOptions {
@@ -36,6 +37,35 @@ struct SelectOptimizationOptions {
   bool force_fallback_implementation = false;
 };
 
+// Enable select optimization on the given RuntimeBuilder, replacing long
+// select chains with a single operation.
+//
+// This assumes that the type information at check time agrees with the
+// configured types at runtime.
+//
+// Important: The select optimization follows spec behavior for traversals.
+//  - `enable_empty_wrapper_null_unboxing` is ignored and optimized traversals
+//    always operates as though it is `true`.
+//  - `enable_heterogeneous_equality` is ignored and optimized traversals
+//    always operate as though it is `true`.
+//
+// This should only be called *once* on a given runtime builder.
+//
+// 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.
+absl::Status EnableSelectOptimization(
+    cel::RuntimeBuilder& builder,
+    const SelectOptimizationOptions& options = {});
+
+// ===============================================================
+// Implementation details -- CEL users should not depend on these.
+// Exposed here for enabling on Legacy APIs. They expose internal details
+// which are not guaranteed to be stable.
+// ===============================================================
+
 // Scans ast for optimizable select branches.
 //
 // In general, this should be done by a type checker but may be deferred to