Skip to content

Commit 85abdee

Browse files
jnthntatumcopybara-github
jnthntatum
authored and
copybara-github
committed
Allow constant folding to recognize identifiers replaced with inlined constants.
PiperOrigin-RevId: 1005954928
1 parent c06497e commit 85abdee

5 files changed

Lines changed: 194 additions & 0 deletions

File tree

‎eval/compiler/constant_folding.cc‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -190,6 +190,14 @@ absl::Status ConstantFoldingExtension::OnPostVisit(PlannerContext& context,
190190
is_const_.pop_back();
191191

192192
if (is_const == IsConst::kNonConst) {
193+
if (context.options().enable_subexpression_inlining &&
194+
node.has_ident_expr()) {
195+
ExecutionPathView subplan = context.GetSubplan(node);
196+
if (subplan.size() == 1 &&
197+
google::api::expr::runtime::IsConstant(subplan[0])) {
198+
return absl::OkStatus();
199+
}
200+
}
193201
// update parent
194202
if (!is_const_.empty()) {
195203
is_const_.back() = IsConst::kNonConst;

‎eval/compiler/constant_folding_test.cc‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -564,6 +564,55 @@ TEST_F(UpdatedConstantFoldingTest, ErrorsOnUnexpectedOrder) {
564564
StatusIs(absl::StatusCode::kInternal));
565565
}
566566

567+
TEST_F(UpdatedConstantFoldingTest, FoldsListWithInlinedConstantIdent) {
568+
ASSERT_OK_AND_ASSIGN(std::unique_ptr<cel::Ast> list_ast,
569+
ParseFromCel("[a, 2]"));
570+
const Expr& create_list = list_ast->root_expr();
571+
const Expr& elem_a = create_list.list_expr().elements()[0].expr();
572+
const Expr& elem_two = create_list.list_expr().elements()[1].expr();
573+
574+
ProgramBuilder program_builder;
575+
program_builder.EnterSubexpression(&create_list);
576+
// Simulate `a` having been inlined as a constant step by FlatExprVisitor.
577+
program_builder.EnterSubexpression(&elem_a);
578+
program_builder.AddStep(
579+
ExpressionStep::MakeConstant(cel::IntValue(10), elem_a.id()));
580+
program_builder.ExitSubexpression(&elem_a);
581+
582+
program_builder.EnterSubexpression(&elem_two);
583+
program_builder.AddStep(
584+
ExpressionStep::MakeConstant(cel::IntValue(2), elem_two.id()));
585+
program_builder.ExitSubexpression(&elem_two);
586+
587+
program_builder.AddStep(CreateCreateListStep(
588+
create_list.list_expr().elements().size(), {}, create_list.id()));
589+
program_builder.ExitSubexpression(&create_list);
590+
591+
options_.enable_subexpression_inlining = true;
592+
std::shared_ptr<google::protobuf::Arena> arena;
593+
PlannerContext context(env_, resolver_, options_,
594+
type_registry_.GetComposedTypeProvider(),
595+
issue_collector_, program_builder, arena);
596+
597+
ProgramOptimizerFactory constant_folder_factory =
598+
CreateConstantFoldingOptimizer();
599+
ASSERT_OK_AND_ASSIGN(std::unique_ptr<ProgramOptimizer> constant_folder,
600+
constant_folder_factory(context, *list_ast));
601+
602+
ASSERT_THAT(constant_folder->OnPreVisit(context, create_list), IsOk());
603+
ASSERT_THAT(constant_folder->OnPreVisit(context, elem_a), IsOk());
604+
ASSERT_THAT(constant_folder->OnPostVisit(context, elem_a), IsOk());
605+
ASSERT_THAT(constant_folder->OnPreVisit(context, elem_two), IsOk());
606+
ASSERT_THAT(constant_folder->OnPostVisit(context, elem_two), IsOk());
607+
ASSERT_THAT(constant_folder->OnPostVisit(context, create_list), IsOk());
608+
609+
ExecutionPath path = std::move(program_builder).FlattenMain();
610+
EXPECT_THAT(path, SizeIs(1));
611+
cel::Value folded_list;
612+
ASSERT_TRUE(google::api::expr::runtime::GetIfConstant(path[0], folded_list));
613+
EXPECT_TRUE(folded_list.IsList());
614+
}
615+
567616
} // namespace
568617

569618
} // namespace cel::runtime_internal

‎eval/compiler/flat_expr_builder_test.cc‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3744,6 +3744,106 @@ TEST(FlatExprBuilderTest, BlockInlinesTrivialSubexpressionsDisabled) {
37443744
// Here trace reports the init step on first access.
37453745
EXPECT_THAT(traced_ids, testing::ElementsAre(3, 4, 6, 9, 5, 1));
37463746
}
3747+
3748+
TEST(FlatExprBuilderTest, BindInlinesConstantAndFoldsMainExpression) {
3749+
ParsedExpr parsed_expr;
3750+
// Equivalent to: cel.bind(a, 1 + 2, (a + 3) + x)
3751+
ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(
3752+
R"pb(
3753+
expr {
3754+
id: 1
3755+
comprehension_expr {
3756+
iter_var: "#unused"
3757+
iter_range {
3758+
id: 2
3759+
list_expr {}
3760+
}
3761+
accu_var: "a"
3762+
accu_init {
3763+
id: 3
3764+
call_expr {
3765+
function: "_+_"
3766+
args {
3767+
id: 4
3768+
const_expr { int64_value: 1 }
3769+
}
3770+
args {
3771+
id: 5
3772+
const_expr { int64_value: 2 }
3773+
}
3774+
}
3775+
}
3776+
loop_condition {
3777+
id: 6
3778+
const_expr { bool_value: false }
3779+
}
3780+
loop_step {
3781+
id: 7
3782+
ident_expr { name: "a" }
3783+
}
3784+
result {
3785+
id: 8
3786+
call_expr {
3787+
function: "_+_"
3788+
args {
3789+
id: 9
3790+
call_expr {
3791+
function: "_+_"
3792+
args {
3793+
id: 10
3794+
ident_expr { name: "a" }
3795+
}
3796+
args {
3797+
id: 11
3798+
const_expr { int64_value: 3 }
3799+
}
3800+
}
3801+
}
3802+
args {
3803+
id: 12
3804+
ident_expr { name: "x" }
3805+
}
3806+
}
3807+
}
3808+
}
3809+
}
3810+
)pb",
3811+
&parsed_expr));
3812+
3813+
InterpreterOptions options;
3814+
google::protobuf::Arena arena;
3815+
options.constant_folding = true;
3816+
options.constant_arena = &arena;
3817+
options.enable_subexpression_inlining = true;
3818+
3819+
ASSERT_OK_AND_ASSIGN(
3820+
std::unique_ptr<CelExpressionBuilder> builder,
3821+
CreateConstantFoldingConformanceTestExprBuilder(options));
3822+
ASSERT_OK_AND_ASSIGN(std::unique_ptr<CelExpression> cel_expr,
3823+
builder->CreateExpression(&parsed_expr.expr(),
3824+
&parsed_expr.source_info()));
3825+
3826+
const FlatExpression& flat_expr =
3827+
static_cast<const CelExpressionFlatImpl*>(cel_expr.get())
3828+
->flat_expression();
3829+
ASSERT_FALSE(flat_expr.subexpressions().empty());
3830+
ExecutionPathView main_path = flat_expr.subexpressions()[0];
3831+
// Main expression consists of:
3832+
// 0: Constant(6) (folded from `a + 3` after `a` was inlined as Constant(3))
3833+
// 1: Identifier(x)
3834+
// 2: _+_
3835+
// 3: ClearSlot
3836+
ASSERT_THAT(main_path, SizeIs(4));
3837+
cel::Value folded_val;
3838+
ASSERT_TRUE(GetIfConstant(main_path[0], folded_val));
3839+
EXPECT_EQ(folded_val.GetInt().NativeValue(), 6);
3840+
EXPECT_EQ(main_path[0].id(), 9);
3841+
3842+
Activation activation;
3843+
activation.InsertValue("x", CelValue::CreateInt64(10));
3844+
ASSERT_OK_AND_ASSIGN(CelValue result, cel_expr->Evaluate(activation, &arena));
3845+
EXPECT_THAT(result, test::IsCelInt64(16));
3846+
}
37473847
} // namespace
37483848

37493849
} // namespace google::api::expr::runtime

‎extensions/BUILD‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -910,6 +910,7 @@ cc_test(
910910
name = "set_membership_optimization_test",
911911
srcs = ["set_membership_optimization_test.cc"],
912912
deps = [
913+
":bindings_ext",
913914
":set_membership_optimization",
914915
":sets_functions",
915916
"//base:attributes",

‎extensions/set_membership_optimization_test.cc‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@
4242
#include "eval/compiler/flat_expr_builder.h"
4343
#include "eval/compiler/qualified_reference_resolver.h"
4444
#include "eval/eval/evaluator_core.h"
45+
#include "extensions/bindings_ext.h"
4546
#include "extensions/sets_functions.h"
4647
#include "internal/status_macros.h"
4748
#include "internal/testing.h"
@@ -89,6 +90,7 @@ absl::StatusOr<std::unique_ptr<Compiler>> CreateTestCompiler() {
8990
NewCompilerBuilder(cel::internal::GetTestingDescriptorPool()));
9091
CEL_RETURN_IF_ERROR(builder->AddLibrary(StandardCompilerLibrary()));
9192
CEL_RETURN_IF_ERROR(builder->AddLibrary(SetsCompilerLibrary()));
93+
CEL_RETURN_IF_ERROR(builder->AddLibrary(BindingsCompilerLibrary()));
9294

9395
TypeCheckerBuilder& checker = builder->GetCheckerBuilder();
9496
checker.set_container("cel.expr.conformance.proto2");
@@ -395,6 +397,40 @@ TEST(SetMembershipOptimizationPlanTest,
395397
}
396398
}
397399

400+
TEST(SetMembershipOptimizationPlanTest,
401+
EmitsCompactPlanWithInlinedBoundConstantSet) {
402+
RuntimeOptions options;
403+
options.container = "cel.expr.conformance.proto2";
404+
options.enable_heterogeneous_equality = true;
405+
options.enable_subexpression_inlining = true;
406+
std::shared_ptr<cel::runtime_internal::RuntimeEnv> env =
407+
NewTestingRuntimeEnv();
408+
ASSERT_THAT(RegisterStandardFunctions(env->function_registry, options),
409+
IsOk());
410+
ASSERT_THAT(RegisterSetsFunctions(env->function_registry, options), IsOk());
411+
FlatExprBuilder builder(std::move(env), options);
412+
builder.set_container(options.container);
413+
builder.AddAstTransform(
414+
NewReferenceResolverExtension(ReferenceResolverOption::kAlways));
415+
builder.AddProgramOptimizer(
416+
cel::runtime_internal::CreateConstantFoldingOptimizer());
417+
builder.AddProgramOptimizer(
418+
CreateSetMembershipOptimizationProgramOptimizer());
419+
420+
for (absl::string_view expr_str : {
421+
"cel.bind(set_var, [1, 2, 3, 4], x in set_var)",
422+
"cel.bind(set_var, ['a', 'b', 'c', 'd'], s in set_var)",
423+
}) {
424+
ASSERT_OK_AND_ASSIGN(std::unique_ptr<Ast> ast, CompileTestExpr(expr_str));
425+
ASSERT_OK_AND_ASSIGN(FlatExpression plan,
426+
builder.CreateExpressionImpl(std::move(ast), nullptr));
427+
ASSERT_FALSE(plan.subexpressions().empty());
428+
// Main expression has 3 steps: Identifier(x/s), Bitset/StringSet step, and
429+
// ClearSlot.
430+
EXPECT_THAT(plan.subexpressions()[0], SizeIs(3)) << expr_str;
431+
}
432+
}
433+
398434
TEST(SetMembershipOptimizationPlanTest, EmitsCompactBitsetPlanForUncheckedAst) {
399435
RuntimeOptions options;
400436
options.container = "cel.expr.conformance.proto2";

0 commit comments

Comments
 (0)