From 921dde9a2def4a0ad0ef7031160dcda243964f8d Mon Sep 17 00:00:00 2001 From: Sean Huh Date: Fri, 16 Oct 2026 16:27:56 -0700 Subject: [PATCH] Allow planner to accept 3rd/4th args on select optimized ASTs from Java. Add a cross-stack test PiperOrigin-RevId: 1001173467 --- eval/compiler/flat_expr_builder.cc | 24 +++--- eval/compiler/flat_expr_builder.h | 5 ++ eval/compiler/flat_expr_builder_test.cc | 45 +++++++++++ eval/public/cel_expr_builder_factory.cc | 28 ++++--- extensions/select_optimization.cc | 24 +++--- extensions/select_optimization.h | 3 +- extensions/select_optimization_test.cc | 100 +++++++++++++++++------- 7 files changed, 160 insertions(+), 69 deletions(-) diff --git a/eval/compiler/flat_expr_builder.cc b/eval/compiler/flat_expr_builder.cc index 8e1c5606d..7270ee42e 100644 --- a/eval/compiler/flat_expr_builder.cc +++ b/eval/compiler/flat_expr_builder.cc @@ -2233,12 +2233,17 @@ std::vector FlattenExpressionTable( } absl::Status CheckAstExtensions( - const std::vector& extensions) { + const std::vector& extensions, + bool enable_select_optimization) { for (const cel::ExtensionSpec& extension : extensions) { if (extension.id() == "cel_block" && extension.version().major() == 1) { // cel_block v1 is always supported. continue; } + if (extension.id() == "select_optimization" && + extension.version().major() == 1 && enable_select_optimization) { + continue; + } // TODO(uncreated-issue/89): Add support for json field names. return absl::InvalidArgumentError(absl::StrCat( @@ -2262,19 +2267,10 @@ absl::StatusOr FlatExprBuilder::CreateExpressionImpl( : RuntimeIssue::Severity::kError; IssueCollector issue_collector(max_severity); - absl::StatusOr> runtime_extensions = - ExtractAndValidateRuntimeExtensions(*ast); - - if (!runtime_extensions.ok()) { - CEL_RETURN_IF_ERROR(issue_collector.AddIssue( - RuntimeIssue::CreateError(runtime_extensions.status()))); - } - - auto status = CheckAstExtensions(*runtime_extensions); - if (!status.ok()) { - CEL_RETURN_IF_ERROR( - issue_collector.AddIssue(RuntimeIssue::CreateError(status))); - } + CEL_ASSIGN_OR_RETURN(std::vector runtime_extensions, + ExtractAndValidateRuntimeExtensions(*ast)); + CEL_RETURN_IF_ERROR( + CheckAstExtensions(runtime_extensions, enable_select_optimization_)); Resolver resolver(container_, function_registry_, type_registry_, GetTypeProvider(), diff --git a/eval/compiler/flat_expr_builder.h b/eval/compiler/flat_expr_builder.h index aa4d0b4e5..fc6066883 100644 --- a/eval/compiler/flat_expr_builder.h +++ b/eval/compiler/flat_expr_builder.h @@ -81,6 +81,10 @@ class FlatExprBuilder { bool optional_types_enabled() const { return enable_optional_types_; } + // Called when select optimization is enabled to accept ASTs tagged with the + // `select_optimization` runtime extension. + void enable_select_optimization() { enable_select_optimization_ = true; } + private: const cel::TypeProvider& GetTypeProvider() const; @@ -90,6 +94,7 @@ class FlatExprBuilder { cel::RuntimeOptions options_; std::string container_; bool enable_optional_types_ = false; + bool enable_select_optimization_ = false; // TODO(uncreated-issue/45): evaluate whether we should use a shared_ptr here to // allow built expressions to keep the registries alive. const cel::FunctionRegistry& function_registry_; diff --git a/eval/compiler/flat_expr_builder_test.cc b/eval/compiler/flat_expr_builder_test.cc index 5d36dfb37..24539f4b7 100644 --- a/eval/compiler/flat_expr_builder_test.cc +++ b/eval/compiler/flat_expr_builder_test.cc @@ -211,6 +211,51 @@ TEST(FlatExprBuilderTest, RuntimeExtensionsError) { HasSubstr("unsupported CEL extension: ext1"))); } +TEST(FlatExprBuilderTest, SelectOptimizationRuntimeExtension) { + Expr expr; + expr.mutable_const_expr()->set_bool_value(true); + auto create = [&](CelExpressionBuilder& builder, absl::string_view id, + int major, int minor) { + SourceInfo source_info; + auto* ext = source_info.add_extensions(); + ext->set_id(id); + ext->mutable_version()->set_major(major); + ext->mutable_version()->set_minor(minor); + ext->add_affected_components(SourceInfo::Extension::COMPONENT_RUNTIME); + return builder.CreateExpression(&expr, &source_info).status(); + }; + + CelExpressionBuilderFlatImpl builder(NewTestingRuntimeEnv()); + // Other program optimizers don't imply select optimization support. + builder.flat_expr_builder().AddProgramOptimizer( + [](auto&, const auto&) { return nullptr; }); + EXPECT_THAT( + create(builder, "select_optimization", 1, 0), + StatusIs( + absl::StatusCode::kInvalidArgument, + HasSubstr("unsupported CEL extension: select_optimization@1.0"))); + + builder.flat_expr_builder().enable_select_optimization(); + EXPECT_THAT(create(builder, "select_optimization", 1, 0), IsOk()); + // Minor versions are informational (see SourceInfo.Extension.Version). + EXPECT_THAT(create(builder, "select_optimization", 1, 1), IsOk()); + EXPECT_THAT(create(builder, "select_optimization", 2, 0), + StatusIs(absl::StatusCode::kInvalidArgument)); + EXPECT_THAT(create(builder, "ext1", 1, 0), + StatusIs(absl::StatusCode::kInvalidArgument)); + + // The legacy API also plans the 3-arg form under default fail_on_warnings. + ASSERT_OK_AND_ASSIGN(ParsedExpr parsed, + parser::Parse("f(x, [[1, 'field', 5]], 0)")); + parsed.mutable_expr()->mutable_call_expr()->set_function("cel.@attribute"); + expr = parsed.expr(); + InterpreterOptions options; + options.enable_select_optimization = true; + EXPECT_THAT( + create(*CreateCelExpressionBuilder(options), "select_optimization", 1, 0), + IsOk()); +} + TEST(FlatExprBuilderTest, ConstValueUnset) { Expr expr; SourceInfo source_info; diff --git a/eval/public/cel_expr_builder_factory.cc b/eval/public/cel_expr_builder_factory.cc index a56c450b0..14f2b4111 100644 --- a/eval/public/cel_expr_builder_factory.cc +++ b/eval/public/cel_expr_builder_factory.cc @@ -122,19 +122,23 @@ std::unique_ptr CreateCelExpressionBuilder( // 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( - kCelAttribute, false, {cel::Kind::kAny, cel::Kind::kList})); - if (!status.ok()) { - ABSL_LOG(ERROR) << "Failed to register " << kCelAttribute << ": " - << status; - } - status = builder->GetRegistry()->RegisterLazyFunction(CelFunctionDescriptor( - kCelHasField, false, {cel::Kind::kAny, cel::Kind::kList})); - if (!status.ok()) { - ABSL_LOG(ERROR) << "Failed to register " << kCelHasField << ": " - << status; + for (const auto& desc : { + CelFunctionDescriptor(kCelAttribute, false, + {cel::Kind::kAny, cel::Kind::kList}), + CelFunctionDescriptor( + kCelAttribute, false, + {cel::Kind::kAny, cel::Kind::kList, cel::Kind::kAny}), + CelFunctionDescriptor(kCelHasField, false, + {cel::Kind::kAny, cel::Kind::kList}), + }) { + if (absl::Status status = + builder->GetRegistry()->RegisterLazyFunction(desc); + !status.ok()) { + ABSL_LOG(ERROR) << "Failed to register " << desc.name() << ": " + << status; + } } + flat_expr_builder.enable_select_optimization(); // Add runtime implementation. flat_expr_builder.AddProgramOptimizer( CreateSelectOptimizationProgramOptimizer()); diff --git a/extensions/select_optimization.cc b/extensions/select_optimization.cc index 02855dec1..5f0512011 100644 --- a/extensions/select_optimization.cc +++ b/extensions/select_optimization.cc @@ -166,7 +166,8 @@ std::optional GetSelectInstruction( } absl::StatusOr SelectQualifierFromList(const ListExpr& list) { - if (list.elements().size() != 2) { + // Type codes may follow [number, name]; they're unused here. + if (list.elements().size() < 2 || list.elements().size() > 4) { return absl::InvalidArgumentError("Invalid cel.attribute select list"); } @@ -774,16 +775,15 @@ absl::Status SelectOptimizer::OnPostVisit(PlannerContext& context, if (fn != kCelHasField && fn != kCelAttribute) { return absl::OkStatus(); } + const bool presence_test = fn == kCelHasField; - if (node.call_expr().args().size() < 2 || - node.call_expr().args().size() > 3) { + // cel.@attribute may take a 3rd arg that only types the result; it's unused. + const size_t max_args = presence_test ? 2 : 3; + const size_t num_args = node.call_expr().args().size(); + if (num_args < 2 || num_args > max_args) { return absl::InvalidArgumentError("Invalid cel.attribute call"); } - if (node.call_expr().args().size() == 3) { - return absl::UnimplementedError("Optionals not yet supported"); - } - CEL_ASSIGN_OR_RETURN(std::vector instructions, SelectInstructionsFromCall(node.call_expr())); @@ -791,12 +791,6 @@ absl::Status SelectOptimizer::OnPostVisit(PlannerContext& context, return absl::InvalidArgumentError("Invalid cel.attribute no select steps."); } - bool presence_test = false; - - if (fn == kCelHasField) { - presence_test = true; - } - const Expr& operand = node.call_expr().args()[0]; absl::string_view identifier; if (operand.has_ident_expr()) { @@ -892,9 +886,13 @@ absl::Status EnableSelectOptimization( // 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( + kCelAttribute, false, {Kind::kAny, Kind::kList, Kind::kAny}))); CEL_RETURN_IF_ERROR(builder.function_registry().RegisterLazyFunction( FunctionDescriptor(kCelHasField, false, {Kind::kAny, Kind::kList}))); + flat_expr_builder->enable_select_optimization(); // Add runtime implementation. flat_expr_builder->AddProgramOptimizer( CreateSelectOptimizationProgramOptimizer(options)); diff --git a/extensions/select_optimization.h b/extensions/select_optimization.h index 4de81b1b0..de412afe9 100644 --- a/extensions/select_optimization.h +++ b/extensions/select_optimization.h @@ -38,7 +38,8 @@ struct SelectOptimizationOptions { }; // Enable select optimization on the given RuntimeBuilder, replacing long -// select chains with a single operation. +// select chains with a single operation. Also enables planning ASTs tagged with +// the `select_optimization` runtime extension. // // This assumes that the type information at check time agrees with the // configured types at runtime. diff --git a/extensions/select_optimization_test.cc b/extensions/select_optimization_test.cc index 619db4a07..ea5512976 100644 --- a/extensions/select_optimization_test.cc +++ b/extensions/select_optimization_test.cc @@ -1753,6 +1753,39 @@ class SelectOptimizationUnexpectedAstTest : public SelectOptimizationTest { return element; } + Expr NextInt(int64_t value) { + Expr result = NextExpr(); + result.mutable_const_expr().set_int_value(value); + return result; + } + + Expr NextString(absl::string_view value) { + Expr result = NextExpr(); + result.mutable_const_expr().set_string_value(value); + return result; + } + + template + Expr NextList(Exprs... elements) { + Expr result = NextExpr(); + (result.mutable_list_expr().add_elements().set_expr(std::move(elements)), + ...); + return result; + } + + // Plans `fn(ident, args...)` with only the select program optimizer. + template + absl::Status PlanSelectCall(absl::string_view fn, Exprs... args) { + auto ast = std::make_unique(NextExpr(), SourceInfo()); + auto& call = ast->mutable_root_expr().mutable_call_expr(); + call.set_function(fn); + call.add_args().mutable_ident_expr().set_name("ident"); + (call.mutable_args().push_back(std::move(args)), ...); + FlatExprBuilder builder(env_, runtime_options_); + builder.AddProgramOptimizer(CreateSelectOptimizationProgramOptimizer()); + return builder.CreateExpressionImpl(std::move(ast), nullptr).status(); + } + protected: int64_t next_id_; }; @@ -1907,37 +1940,46 @@ TEST_F(SelectOptimizationUnexpectedAstTest, StatusIs(absl::StatusCode::kInvalidArgument)); } -TEST_F(SelectOptimizationUnexpectedAstTest, OptionalNotYetSupported) { - std::unique_ptr ast = std::make_unique(NextExpr(), SourceInfo()); - - ast->mutable_root_expr().mutable_call_expr().set_function(kCelAttribute); - auto& call_args = ast->mutable_root_expr().mutable_call_expr().mutable_args(); - call_args.emplace_back(NextExpr()).mutable_ident_expr().set_name("ident"); - - auto& list_expr = call_args.emplace_back(NextExpr()).mutable_list_expr(); - auto& fields = list_expr.mutable_elements() - .emplace_back(NextListExprElement()) - .mutable_expr() - .mutable_list_expr() - .mutable_elements(); - - fields.emplace_back(NextListExprElement()) - .mutable_expr() - .mutable_const_expr() - .set_int64_value(1); - fields.emplace_back(NextListExprElement()) - .mutable_expr() - .mutable_const_expr() - .set_string_value("field"); - - call_args.emplace_back(NextExpr()).mutable_const_expr().set_int64_value(0); - - FlatExprBuilder builder(env_, runtime_options_); +TEST_F(SelectOptimizationUnexpectedAstTest, TypedTupleAttributeAccepted) { + using FieldDescriptor = google::protobuf::FieldDescriptor; + // [field_num, field_name, type_code] hops, a map leaf + // [field_num, field_name, -1, [key_type, value_type]], and a 3rd argument + // that only binds the result type. + EXPECT_THAT( + PlanSelectCall( + kCelAttribute, + NextList(NextList(NextInt(1), NextString("field"), + NextInt(FieldDescriptor::TYPE_MESSAGE)), + NextList(NextInt(2), NextString("map_field"), NextInt(-1), + NextList(NextInt(FieldDescriptor::TYPE_STRING), + NextInt(FieldDescriptor::TYPE_INT32)))), + NextInt(0)), + IsOk()); +} - builder.AddProgramOptimizer(CreateSelectOptimizationProgramOptimizer()); +TEST_F(SelectOptimizationUnexpectedAstTest, TooManyArguments) { + EXPECT_THAT( + PlanSelectCall(kCelHasField, + NextList(NextList(NextInt(1), NextString("field"))), + NextInt(0)), + StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("Invalid cel.attribute call"))); + EXPECT_THAT( + PlanSelectCall(kCelAttribute, + NextList(NextList(NextInt(1), NextString("field"))), + NextInt(0), NextInt(0)), + StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("Invalid cel.attribute call"))); +} - EXPECT_THAT(builder.CreateExpressionImpl(std::move(ast), nullptr), - StatusIs(absl::StatusCode::kUnimplemented)); +TEST_F(SelectOptimizationUnexpectedAstTest, + MalformedSelectPathTooManyElements) { + EXPECT_THAT( + PlanSelectCall(kCelAttribute, + NextList(NextList(NextInt(1), NextString("field"), + NextInt(5), NextInt(0), NextInt(0)))), + StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("Invalid cel.attribute select list"))); } } // namespace