From 147f1f953a8bb225f89725f900ed6e51b52362c5 Mon Sep 17 00:00:00 2001 From: Jonathan Tatum Date: Tue, 27 Oct 2026 16:32:24 -0700 Subject: [PATCH] Allow constant folding to recognize identifiers replaced with inlined constants. PiperOrigin-RevId: 1008459534 --- eval/compiler/constant_folding.cc | 8 ++ eval/compiler/constant_folding_test.cc | 49 +++++++++ eval/compiler/flat_expr_builder_test.cc | 100 ++++++++++++++++++ extensions/BUILD | 1 + .../set_membership_optimization_test.cc | 36 +++++++ 5 files changed, 194 insertions(+) diff --git a/eval/compiler/constant_folding.cc b/eval/compiler/constant_folding.cc index 8e854a867..3d308b81f 100644 --- a/eval/compiler/constant_folding.cc +++ b/eval/compiler/constant_folding.cc @@ -190,6 +190,14 @@ absl::Status ConstantFoldingExtension::OnPostVisit(PlannerContext& context, is_const_.pop_back(); if (is_const == IsConst::kNonConst) { + if (context.options().enable_subexpression_inlining && + node.has_ident_expr()) { + ExecutionPathView subplan = context.GetSubplan(node); + if (subplan.size() == 1 && + google::api::expr::runtime::IsConstant(subplan[0])) { + return absl::OkStatus(); + } + } // update parent if (!is_const_.empty()) { is_const_.back() = IsConst::kNonConst; diff --git a/eval/compiler/constant_folding_test.cc b/eval/compiler/constant_folding_test.cc index d9ba91bfa..c612e7c1f 100644 --- a/eval/compiler/constant_folding_test.cc +++ b/eval/compiler/constant_folding_test.cc @@ -564,6 +564,55 @@ TEST_F(UpdatedConstantFoldingTest, ErrorsOnUnexpectedOrder) { StatusIs(absl::StatusCode::kInternal)); } +TEST_F(UpdatedConstantFoldingTest, FoldsListWithInlinedConstantIdent) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr list_ast, + ParseFromCel("[a, 2]")); + const Expr& create_list = list_ast->root_expr(); + const Expr& elem_a = create_list.list_expr().elements()[0].expr(); + const Expr& elem_two = create_list.list_expr().elements()[1].expr(); + + ProgramBuilder program_builder; + program_builder.EnterSubexpression(&create_list); + // Simulate `a` having been inlined as a constant step by FlatExprVisitor. + program_builder.EnterSubexpression(&elem_a); + program_builder.AddStep( + ExpressionStep::MakeConstant(cel::IntValue(10), elem_a.id())); + program_builder.ExitSubexpression(&elem_a); + + program_builder.EnterSubexpression(&elem_two); + program_builder.AddStep( + ExpressionStep::MakeConstant(cel::IntValue(2), elem_two.id())); + program_builder.ExitSubexpression(&elem_two); + + program_builder.AddStep(CreateCreateListStep( + create_list.list_expr().elements().size(), {}, create_list.id())); + program_builder.ExitSubexpression(&create_list); + + options_.enable_subexpression_inlining = true; + std::shared_ptr arena; + PlannerContext context(env_, resolver_, options_, + type_registry_.GetComposedTypeProvider(), + issue_collector_, program_builder, arena); + + ProgramOptimizerFactory constant_folder_factory = + CreateConstantFoldingOptimizer(); + ASSERT_OK_AND_ASSIGN(std::unique_ptr constant_folder, + constant_folder_factory(context, *list_ast)); + + ASSERT_THAT(constant_folder->OnPreVisit(context, create_list), IsOk()); + ASSERT_THAT(constant_folder->OnPreVisit(context, elem_a), IsOk()); + ASSERT_THAT(constant_folder->OnPostVisit(context, elem_a), IsOk()); + ASSERT_THAT(constant_folder->OnPreVisit(context, elem_two), IsOk()); + ASSERT_THAT(constant_folder->OnPostVisit(context, elem_two), IsOk()); + ASSERT_THAT(constant_folder->OnPostVisit(context, create_list), IsOk()); + + ExecutionPath path = std::move(program_builder).FlattenMain(); + EXPECT_THAT(path, SizeIs(1)); + cel::Value folded_list; + ASSERT_TRUE(google::api::expr::runtime::GetIfConstant(path[0], folded_list)); + EXPECT_TRUE(folded_list.IsList()); +} + } // namespace } // namespace cel::runtime_internal diff --git a/eval/compiler/flat_expr_builder_test.cc b/eval/compiler/flat_expr_builder_test.cc index 85938d068..a15a2e4f1 100644 --- a/eval/compiler/flat_expr_builder_test.cc +++ b/eval/compiler/flat_expr_builder_test.cc @@ -3744,6 +3744,106 @@ TEST(FlatExprBuilderTest, BlockInlinesTrivialSubexpressionsDisabled) { // Here trace reports the init step on first access. EXPECT_THAT(traced_ids, testing::ElementsAre(3, 4, 6, 9, 5, 1)); } + +TEST(FlatExprBuilderTest, BindInlinesConstantAndFoldsMainExpression) { + ParsedExpr parsed_expr; + // Equivalent to: cel.bind(a, 1 + 2, (a + 3) + x) + ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString( + R"pb( + expr { + id: 1 + comprehension_expr { + iter_var: "#unused" + iter_range { + id: 2 + list_expr {} + } + accu_var: "a" + accu_init { + id: 3 + call_expr { + function: "_+_" + args { + id: 4 + const_expr { int64_value: 1 } + } + args { + id: 5 + const_expr { int64_value: 2 } + } + } + } + loop_condition { + id: 6 + const_expr { bool_value: false } + } + loop_step { + id: 7 + ident_expr { name: "a" } + } + result { + id: 8 + call_expr { + function: "_+_" + args { + id: 9 + call_expr { + function: "_+_" + args { + id: 10 + ident_expr { name: "a" } + } + args { + id: 11 + const_expr { int64_value: 3 } + } + } + } + args { + id: 12 + ident_expr { name: "x" } + } + } + } + } + } + )pb", + &parsed_expr)); + + InterpreterOptions options; + google::protobuf::Arena arena; + options.constant_folding = true; + options.constant_arena = &arena; + options.enable_subexpression_inlining = true; + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr builder, + CreateConstantFoldingConformanceTestExprBuilder(options)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr cel_expr, + builder->CreateExpression(&parsed_expr.expr(), + &parsed_expr.source_info())); + + const FlatExpression& flat_expr = + static_cast(cel_expr.get()) + ->flat_expression(); + ASSERT_FALSE(flat_expr.subexpressions().empty()); + ExecutionPathView main_path = flat_expr.subexpressions()[0]; + // Main expression consists of: + // 0: Constant(6) (folded from `a + 3` after `a` was inlined as Constant(3)) + // 1: Identifier(x) + // 2: _+_ + // 3: ClearSlot + ASSERT_THAT(main_path, SizeIs(4)); + cel::Value folded_val; + ASSERT_TRUE(GetIfConstant(main_path[0], folded_val)); + EXPECT_EQ(folded_val.GetInt().NativeValue(), 6); + EXPECT_EQ(main_path[0].id(), 9); + + Activation activation; + activation.InsertValue("x", CelValue::CreateInt64(10)); + ASSERT_OK_AND_ASSIGN(CelValue result, cel_expr->Evaluate(activation, &arena)); + EXPECT_THAT(result, test::IsCelInt64(16)); +} } // namespace } // namespace google::api::expr::runtime diff --git a/extensions/BUILD b/extensions/BUILD index f6e5cab66..4b4c91a31 100644 --- a/extensions/BUILD +++ b/extensions/BUILD @@ -910,6 +910,7 @@ cc_test( name = "set_membership_optimization_test", srcs = ["set_membership_optimization_test.cc"], deps = [ + ":bindings_ext", ":set_membership_optimization", ":sets_functions", "//base:attributes", diff --git a/extensions/set_membership_optimization_test.cc b/extensions/set_membership_optimization_test.cc index 827108319..aff9f6dc5 100644 --- a/extensions/set_membership_optimization_test.cc +++ b/extensions/set_membership_optimization_test.cc @@ -42,6 +42,7 @@ #include "eval/compiler/flat_expr_builder.h" #include "eval/compiler/qualified_reference_resolver.h" #include "eval/eval/evaluator_core.h" +#include "extensions/bindings_ext.h" #include "extensions/sets_functions.h" #include "internal/status_macros.h" #include "internal/testing.h" @@ -89,6 +90,7 @@ absl::StatusOr> CreateTestCompiler() { NewCompilerBuilder(cel::internal::GetTestingDescriptorPool())); CEL_RETURN_IF_ERROR(builder->AddLibrary(StandardCompilerLibrary())); CEL_RETURN_IF_ERROR(builder->AddLibrary(SetsCompilerLibrary())); + CEL_RETURN_IF_ERROR(builder->AddLibrary(BindingsCompilerLibrary())); TypeCheckerBuilder& checker = builder->GetCheckerBuilder(); checker.set_container("cel.expr.conformance.proto2"); @@ -395,6 +397,40 @@ TEST(SetMembershipOptimizationPlanTest, } } +TEST(SetMembershipOptimizationPlanTest, + EmitsCompactPlanWithInlinedBoundConstantSet) { + RuntimeOptions options; + options.container = "cel.expr.conformance.proto2"; + options.enable_heterogeneous_equality = true; + options.enable_subexpression_inlining = true; + std::shared_ptr env = + NewTestingRuntimeEnv(); + ASSERT_THAT(RegisterStandardFunctions(env->function_registry, options), + IsOk()); + ASSERT_THAT(RegisterSetsFunctions(env->function_registry, options), IsOk()); + FlatExprBuilder builder(std::move(env), options); + builder.set_container(options.container); + builder.AddAstTransform( + NewReferenceResolverExtension(ReferenceResolverOption::kAlways)); + builder.AddProgramOptimizer( + cel::runtime_internal::CreateConstantFoldingOptimizer()); + builder.AddProgramOptimizer( + CreateSetMembershipOptimizationProgramOptimizer()); + + for (absl::string_view expr_str : { + "cel.bind(set_var, [1, 2, 3, 4], x in set_var)", + "cel.bind(set_var, ['a', 'b', 'c', 'd'], s in set_var)", + }) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr ast, CompileTestExpr(expr_str)); + ASSERT_OK_AND_ASSIGN(FlatExpression plan, + builder.CreateExpressionImpl(std::move(ast), nullptr)); + ASSERT_FALSE(plan.subexpressions().empty()); + // Main expression has 3 steps: Identifier(x/s), Bitset/StringSet step, and + // ClearSlot. + EXPECT_THAT(plan.subexpressions()[0], SizeIs(3)) << expr_str; + } +} + TEST(SetMembershipOptimizationPlanTest, EmitsCompactBitsetPlanForUncheckedAst) { RuntimeOptions options; options.container = "cel.expr.conformance.proto2";