diff --git a/eval/compiler/BUILD b/eval/compiler/BUILD index a8f40749d..291b61936 100644 --- a/eval/compiler/BUILD +++ b/eval/compiler/BUILD @@ -61,6 +61,7 @@ cc_test( ":resolver", "//common:expr", "//common:type", + "//common:value", "//eval/eval:evaluator_core", "//eval/eval:expression_step_logic", "//internal:testing", diff --git a/eval/compiler/flat_expr_builder.cc b/eval/compiler/flat_expr_builder.cc index 0d58c5864..394fd9a65 100644 --- a/eval/compiler/flat_expr_builder.cc +++ b/eval/compiler/flat_expr_builder.cc @@ -818,16 +818,24 @@ class FlatExprVisitor : public cel::AstVisitor { SlotLookupResult slot = LookupSlot(path); if (slot.subexpression >= 0) { - auto* subexpression = + const ProgramBuilder::Subexpression* subexpression = program_builder_.GetExtractedSubexpression(slot.subexpression); if (subexpression == nullptr) { SetProgressStatusIfError( absl::InternalError("bad subexpression reference")); return; } - // Off by one since mainline expression will be index 0. - AddStep(ExpressionStep::MakeLazyInitStep( - slot.slot, slot.subexpression + 1, expr.id())); + if (options_.enable_subexpression_inlining) { + if (std::optional inlined_step = + subexpression->MaybeInline(expr.id()); + inlined_step.has_value()) { + AddStep(*std::move(inlined_step)); + return; + } + } + // Off by one since mainline expression will be index 0. + AddStep(ExpressionStep::MakeLazyInitStep( + slot.slot, slot.subexpression + 1, expr.id())); return; } else if (slot.slot >= 0) { AddStep(ExpressionStep::MakeReadSlotStep(slot.slot, expr.id())); diff --git a/eval/compiler/flat_expr_builder_extensions.cc b/eval/compiler/flat_expr_builder_extensions.cc index f38b9c9e2..6850b798f 100644 --- a/eval/compiler/flat_expr_builder_extensions.cc +++ b/eval/compiler/flat_expr_builder_extensions.cc @@ -14,8 +14,10 @@ #include "eval/compiler/flat_expr_builder_extensions.h" #include +#include #include #include +#include #include #include @@ -28,7 +30,9 @@ #include "absl/status/statusor.h" #include "absl/types/variant.h" #include "common/expr.h" +#include "common/value.h" #include "eval/eval/evaluator_core.h" +#include "eval/eval/ident_step.h" namespace google::api::expr::runtime { @@ -36,6 +40,37 @@ namespace { using Subexpression = google::api::expr::runtime::ProgramBuilder::Subexpression; +const ExpressionStep* absl_nullable TryGetSingleStep( + const Subexpression* absl_nonnull subexpr) { + while (subexpr != nullptr) { + if (subexpr->IsFlattened()) { + const ExecutionPath& elements = subexpr->flattened_elements(); + if (elements.size() != 1) { + return nullptr; + } + return &elements[0]; + } + const auto& elements = subexpr->elements(); + if (elements.size() != 1) { + return nullptr; + } + const auto& element = elements[0]; + if (const ExpressionStep* step = absl::get_if(&element); + step != nullptr) { + return step; + } + // This might happen if an AST node is interpreted as a no-op, but we still + // create a subexpression for it. + if (Subexpression* const* child = absl::get_if(&element); + child != nullptr) { + subexpr = *child; + continue; + } + return nullptr; + } + return nullptr; +} + } // namespace Subexpression::Subexpression(const cel::Expr* self, ProgramBuilder* owner) @@ -185,7 +220,33 @@ bool Subexpression::ExtractTo(ExecutionPath& out) { return true; } -ExecutionPath ProgramBuilder::FlattenSubexpression(Subexpression* expr) { +bool Subexpression::IsInlinable() const { + const ExpressionStep* step = TryGetSingleStep(this); + if (step == nullptr) { + return false; + } + return IsConstant(*step) || GetIfIdentStep(*step) != nullptr; +} + +std::optional Subexpression::MaybeInline( + int64_t expr_id) const { + const ExpressionStep* step = TryGetSingleStep(this); + if (step == nullptr) { + return std::nullopt; + } + cel::Value value; + if (GetIfConstant(*step, value)) { + return ExpressionStep::MakeConstant(value, expr_id); + } + if (const IdentStepInfo* ident = GetIfIdentStep(*step); ident != nullptr) { + return ExpressionStep::MakeIdentifierStep(ident->ident_index, + ident->slot_index, expr_id); + } + return std::nullopt; +} + +ExecutionPath ProgramBuilder::FlattenSubexpression( + Subexpression* absl_nonnull expr) { ExecutionPath out; if (!expr) { diff --git a/eval/compiler/flat_expr_builder_extensions.h b/eval/compiler/flat_expr_builder_extensions.h index 2a888a8e2..c22ef2a74 100644 --- a/eval/compiler/flat_expr_builder_extensions.h +++ b/eval/compiler/flat_expr_builder_extensions.h @@ -179,6 +179,14 @@ class ProgramBuilder { // Returns false if the subexpression is not currently flattened. bool ExtractTo(ExecutionPath& out); + // Returns true if this subexpression is trivial (reduces to a single + // constant or identifier step) and can be inlined at reference sites. + bool IsInlinable() const; + + // Creates an inlined ExpressionStep for this subexpression with the + // given expr_id, or returns std::nullopt if !IsInlinable(). + std::optional MaybeInline(int64_t expr_id) const; + private: Subexpression(const cel::Expr* self, ProgramBuilder* owner); @@ -256,6 +264,15 @@ class ProgramBuilder { return extracted_subexpressions_[index]; } + const Subexpression* absl_nullable GetExtractedSubexpression( + size_t index) const { + if (index >= extracted_subexpressions_.size()) { + return nullptr; + } + + return extracted_subexpressions_[index]; + } + // Return index to the extracted subexpression. // // Returns -1 if the subexpression is not found. diff --git a/eval/compiler/flat_expr_builder_extensions_test.cc b/eval/compiler/flat_expr_builder_extensions_test.cc index 1add6b2b4..661bd8953 100644 --- a/eval/compiler/flat_expr_builder_extensions_test.cc +++ b/eval/compiler/flat_expr_builder_extensions_test.cc @@ -23,10 +23,12 @@ #include "absl/status/statusor.h" #include "common/expr.h" #include "common/type.h" +#include "common/value.h" #include "eval/compiler/resolver.h" #include "eval/eval/evaluator_core.h" #include "eval/eval/expression_step_logic.h" #include "eval/eval/function_step.h" +#include "eval/eval/ident_step.h" #include "internal/testing.h" #include "runtime/function_registry.h" #include "runtime/internal/issue_collector.h" @@ -539,5 +541,100 @@ TEST_F(ProgramBuilderTest, ExtractToRequiresFlatten) { UniquePtrHolds(step_ptrs.a))); } +TEST_F(ProgramBuilderTest, IsInlinableAndMaybeInlineConstant) { + Expr a; + Expr b; + ProgramBuilder program_builder; + + program_builder.EnterSubexpression(&a); + program_builder.EnterSubexpression(&b); + program_builder.AddStep(ExpressionStep::MakeConstant(cel::IntValue(42), 10)); + program_builder.ExitSubexpression(&b); + program_builder.ExitSubexpression(&a); + + Subexpression* subexpr_a = program_builder.GetSubexpression(&a); + Subexpression* subexpr_b = program_builder.GetSubexpression(&b); + ASSERT_TRUE(subexpr_a != nullptr); + ASSERT_TRUE(subexpr_b != nullptr); + + // Unflattened direct constant step and nested single-child subexpression. + EXPECT_TRUE(subexpr_b->IsInlinable()); + EXPECT_TRUE(subexpr_a->IsInlinable()); + + std::optional inlined_b = subexpr_b->MaybeInline(99); + ASSERT_TRUE(inlined_b.has_value()); + EXPECT_EQ(inlined_b->id(), 99); + cel::Value val; + ASSERT_TRUE(GetIfConstant(*inlined_b, val)); + EXPECT_EQ(val.GetInt().NativeValue(), 42); + + // Flattened constant step remains inlinable. + subexpr_a->Flatten(); + EXPECT_TRUE(subexpr_a->IsInlinable()); + std::optional inlined_a = subexpr_a->MaybeInline(100); + ASSERT_TRUE(inlined_a.has_value()); + EXPECT_EQ(inlined_a->id(), 100); + ASSERT_TRUE(GetIfConstant(*inlined_a, val)); + EXPECT_EQ(val.GetInt().NativeValue(), 42); +} + +TEST_F(ProgramBuilderTest, IsInlinableAndMaybeInlineIdentifier) { + Expr a; + ProgramBuilder program_builder; + + program_builder.EnterSubexpression(&a); + program_builder.AddStep(ExpressionStep::MakeIdentifierStep(3, 7, 10)); + program_builder.ExitSubexpression(&a); + + Subexpression* subexpr_a = program_builder.GetSubexpression(&a); + ASSERT_TRUE(subexpr_a != nullptr); + + EXPECT_TRUE(subexpr_a->IsInlinable()); + std::optional inlined = subexpr_a->MaybeInline(55); + ASSERT_TRUE(inlined.has_value()); + EXPECT_EQ(inlined->id(), 55); + const IdentStepInfo* ident_info = GetIfIdentStep(*inlined); + ASSERT_TRUE(ident_info != nullptr); + EXPECT_EQ(ident_info->ident_index, 3); + EXPECT_EQ(ident_info->slot_index, 7); + + subexpr_a->Flatten(); + EXPECT_TRUE(subexpr_a->IsInlinable()); + std::optional inlined_flat = subexpr_a->MaybeInline(56); + ASSERT_TRUE(inlined_flat.has_value()); + EXPECT_EQ(inlined_flat->id(), 56); + ident_info = GetIfIdentStep(*inlined_flat); + ASSERT_TRUE(ident_info != nullptr); + EXPECT_EQ(ident_info->ident_index, 3); + EXPECT_EQ(ident_info->slot_index, 7); +} + +TEST_F(ProgramBuilderTest, NotInlinableWhenNonTrivial) { + Expr a; + Expr b; + Expr c; + ProgramBuilder program_builder; + + ASSERT_OK_AND_ASSIGN(SimpleTreeSteps step_ptrs, + InitSimpleTree(a, b, c, program_builder)); + (void)step_ptrs; + + Subexpression* subexpr_a = program_builder.GetSubexpression(&a); + Subexpression* subexpr_b = program_builder.GetSubexpression(&b); + ASSERT_TRUE(subexpr_a != nullptr); + ASSERT_TRUE(subexpr_b != nullptr); + + // Single generic step is not inlinable. + EXPECT_FALSE(subexpr_b->IsInlinable()); + EXPECT_EQ(subexpr_b->MaybeInline(1), std::nullopt); + + // Multi-step subexpression is not inlinable (unflattened and flattened). + EXPECT_FALSE(subexpr_a->IsInlinable()); + EXPECT_EQ(subexpr_a->MaybeInline(1), std::nullopt); + subexpr_a->Flatten(); + EXPECT_FALSE(subexpr_a->IsInlinable()); + EXPECT_EQ(subexpr_a->MaybeInline(1), std::nullopt); +} + } // namespace } // namespace google::api::expr::runtime diff --git a/eval/compiler/flat_expr_builder_test.cc b/eval/compiler/flat_expr_builder_test.cc index 8d43b3191..85938d068 100644 --- a/eval/compiler/flat_expr_builder_test.cc +++ b/eval/compiler/flat_expr_builder_test.cc @@ -3596,6 +3596,154 @@ TEST(FlatExprBuilderTest, EXPECT_EQ(activation.LookupCount("single_var"), 1); EXPECT_EQ(activation.LookupCount("i"), 0); } + +// Returns a parsed expression containing a `cel.@block` call with trivial +// subexpressions. +// +// Rough policy equivalent: +// +// variables: +// - name: a +// expression: "10" +// - name: b +// expression: "x" +// rule: (a + b) + b +// +// Which lowers to: +// cel.@block([10, x], (@index0 + @index1) + @index1) +ParsedExpr MakeTrivialBlockParsedExpr() { + ParsedExpr parsed_expr; + ABSL_CHECK(google::protobuf::TextFormat::ParseFromString( + R"pb( + expr: { + id: 1 + call_expr: { + function: "cel.@block" + args { + id: 2 + list_expr: { + elements { + id: 3 + const_expr: { int64_value: 10 } + } + elements { + id: 4 + ident_expr: { name: "x" } + } + } + } + args { + id: 5 + call_expr: { + function: "_+_" + args { + id: 6 + call_expr: { + function: "_+_" + args { + id: 7 + ident_expr: { name: "@index0" } + } + args { + id: 8 + ident_expr: { name: "@index1" } + } + } + } + args { + id: 9 + ident_expr: { name: "@index1" } + } + } + } + } + } + )pb", + &parsed_expr)); + return parsed_expr; +} + +TEST(FlatExprBuilderTest, BlockInlinesTrivialSubexpressionsEnabled) { + ParsedExpr parsed_expr = MakeTrivialBlockParsedExpr(); + + cel::RuntimeOptions options; + options.enable_subexpression_inlining = true; + CelExpressionBuilderFlatImpl builder(NewTestingRuntimeEnv(), options); + ASSERT_THAT(RegisterBuiltinFunctions(builder.GetRegistry()), + absl_testing::IsOk()); + 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]; + ASSERT_GE(main_path.size(), 3); + + cel::Value const_val; + EXPECT_TRUE(GetIfConstant(main_path[0], const_val)); + EXPECT_EQ(main_path[0].id(), 7); + EXPECT_EQ(const_val.GetInt().NativeValue(), 10); + const IdentStepInfo* ident1 = GetIfIdentStep(main_path[1]); + ASSERT_NE(ident1, nullptr); + EXPECT_EQ(main_path[1].id(), 8); + + CountingActivation activation; + activation.InsertInt("x", 5); + google::protobuf::Arena arena; + std::vector traced_ids; + ASSERT_OK_AND_ASSIGN(CelValue result, + cel_expr->Trace(activation, &arena, + [&](int64_t id, const CelValue&, + google::protobuf::Arena*) -> absl::Status { + traced_ids.push_back(id); + return absl::OkStatus(); + })); + EXPECT_THAT(result, test::IsCelInt64(20)); + EXPECT_EQ(activation.LookupCount("x"), 1); + // We don't guarantee trace specifics here, just documenting that it + // is deterministic and changed from the non-inlined version. + EXPECT_THAT(traced_ids, testing::ElementsAre(7, 8, 6, 9, 5, 1)); +} + +TEST(FlatExprBuilderTest, BlockInlinesTrivialSubexpressionsDisabled) { + ParsedExpr parsed_expr = MakeTrivialBlockParsedExpr(); + + cel::RuntimeOptions options; + options.enable_subexpression_inlining = false; + CelExpressionBuilderFlatImpl builder(NewTestingRuntimeEnv(), options); + ASSERT_THAT(RegisterBuiltinFunctions(builder.GetRegistry()), + absl_testing::IsOk()); + 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]; + EXPECT_EQ(main_path.size(), 6); + EXPECT_EQ(main_path[0].kind(), ExpressionStepKind::kLazyInit); + + CountingActivation activation; + activation.InsertInt("x", 5); + google::protobuf::Arena arena; + std::vector traced_ids; + ASSERT_OK_AND_ASSIGN(CelValue result, + cel_expr->Trace(activation, &arena, + [&](int64_t id, const CelValue&, + google::protobuf::Arena*) -> absl::Status { + traced_ids.push_back(id); + return absl::OkStatus(); + })); + EXPECT_THAT(result, test::IsCelInt64(20)); + EXPECT_EQ(activation.LookupCount("x"), 1); + // Here trace reports the init step on first access. + EXPECT_THAT(traced_ids, testing::ElementsAre(3, 4, 6, 9, 5, 1)); +} } // namespace } // namespace google::api::expr::runtime diff --git a/eval/eval/evaluator_core.cc b/eval/eval/evaluator_core.cc index f8c6fe35d..a1d46382e 100644 --- a/eval/eval/evaluator_core.cc +++ b/eval/eval/evaluator_core.cc @@ -35,6 +35,7 @@ #include "eval/eval/comprehension_slots.h" #include "eval/eval/comprehension_step.h" #include "eval/eval/equality_steps.h" +#include "eval/eval/ident_step.h" #include "eval/eval/lazy_init_step.h" #include "eval/eval/logic_step.h" #include "runtime/activation_interface.h" @@ -376,6 +377,13 @@ IdentStepInfo* GetIfIdentStep(ExpressionStep& step) { return nullptr; } +const IdentStepInfo* GetIfIdentStep(const ExpressionStep& step) { + if (step.header_.kind == ExpressionStepKind::kIdentifier) { + return &step.u_.ident_step; + } + return nullptr; +} + std::optional ExpressionStep::StackDelta() const { switch (header_.kind) { case ExpressionStepKind::kMovedFrom: diff --git a/eval/eval/evaluator_core.h b/eval/eval/evaluator_core.h index 6f1ae2112..dd8f3d699 100644 --- a/eval/eval/evaluator_core.h +++ b/eval/eval/evaluator_core.h @@ -456,6 +456,7 @@ class ExpressionStep { friend TernaryJumpStepInfo* GetIfTernaryJumpStep(ExpressionStep& step); friend FixedJumpStepInfo* GetIfFixedJumpStep(ExpressionStep& step); friend IdentStepInfo* GetIfIdentStep(ExpressionStep& step); + friend const IdentStepInfo* GetIfIdentStep(const ExpressionStep& step); Header header_; // Note: ptr members are 'owned' by the step. @@ -1067,6 +1068,21 @@ ComprehensionCondStep* GetIfComprehensionCondStep(ExpressionStep& step); ComprehensionNextStep* GetIfComprehensionNextStep(ExpressionStep& step); BoolJumpStepInfo* GetIfBoolJumpStep(ExpressionStep& step); +BoolJumpStepInfo* GetIfBoolJumpStep(ExpressionStep* step); + +TernaryJumpStepInfo* GetIfTernaryJumpStep(ExpressionStep& step); +TernaryJumpStepInfo* GetIfTernaryJumpStep(ExpressionStep* step); + +FixedJumpStepInfo* GetIfFixedJumpStep(ExpressionStep& step); +FixedJumpStepInfo* GetIfFixedJumpStep(ExpressionStep* step); + +IdentStepInfo* GetIfIdentStep(ExpressionStep& step); +const IdentStepInfo* GetIfIdentStep(const ExpressionStep& step); +IdentStepInfo* GetIfIdentStep(ExpressionStep* step); +const IdentStepInfo* GetIfIdentStep(const ExpressionStep* step); + +// Implementation details. + inline BoolJumpStepInfo* GetIfBoolJumpStep(ExpressionStep* step) { if (step == nullptr) { return nullptr; @@ -1074,7 +1090,6 @@ inline BoolJumpStepInfo* GetIfBoolJumpStep(ExpressionStep* step) { return GetIfBoolJumpStep(*step); } -TernaryJumpStepInfo* GetIfTernaryJumpStep(ExpressionStep& step); inline TernaryJumpStepInfo* GetIfTernaryJumpStep(ExpressionStep* step) { if (step == nullptr) { return nullptr; @@ -1082,7 +1097,6 @@ inline TernaryJumpStepInfo* GetIfTernaryJumpStep(ExpressionStep* step) { return GetIfTernaryJumpStep(*step); } -FixedJumpStepInfo* GetIfFixedJumpStep(ExpressionStep& step); inline FixedJumpStepInfo* GetIfFixedJumpStep(ExpressionStep* step) { if (step == nullptr) { return nullptr; @@ -1090,7 +1104,6 @@ inline FixedJumpStepInfo* GetIfFixedJumpStep(ExpressionStep* step) { return GetIfFixedJumpStep(*step); } -IdentStepInfo* GetIfIdentStep(ExpressionStep& step); inline IdentStepInfo* GetIfIdentStep(ExpressionStep* step) { if (step == nullptr) { return nullptr; @@ -1098,7 +1111,12 @@ inline IdentStepInfo* GetIfIdentStep(ExpressionStep* step) { return GetIfIdentStep(*step); } -// Implementation details. +inline const IdentStepInfo* GetIfIdentStep(const ExpressionStep* step) { + if (step == nullptr) { + return nullptr; + } + return GetIfIdentStep(*step); +} inline ExpressionStep::~ExpressionStep() { switch (header_.kind) { diff --git a/eval/public/cel_options.cc b/eval/public/cel_options.cc index 93b67ad35..c3aef0167 100644 --- a/eval/public/cel_options.cc +++ b/eval/public/cel_options.cc @@ -45,6 +45,7 @@ cel::RuntimeOptions ConvertToRuntimeOptions(const InterpreterOptions& options) { options.enable_fast_builtins, options.enable_precision_preserving_double_format, options.enable_typed_field_access, + options.enable_subexpression_inlining, }; } diff --git a/eval/public/cel_options.h b/eval/public/cel_options.h index f28b37a3a..0bb6d81ed 100644 --- a/eval/public/cel_options.h +++ b/eval/public/cel_options.h @@ -225,6 +225,13 @@ struct InterpreterOptions { // path for field access when the type is known at plan time, instead of using // the generic field access implementation. bool enable_typed_field_access = true; + + // When enabled, trivial lazy subexpressions (such as constants or + // identifiers) bound via cel.bind or cel.@block will be inlined directly + // at reference sites during planning. + // + // Note: Enabling this option may change observed expression tracing behavior. + bool enable_subexpression_inlining = false; }; // LINT.ThenChange(//depot/google3/runtime/runtime_options.h) diff --git a/runtime/runtime_options.h b/runtime/runtime_options.h index eaf210366..1c52e9b4c 100644 --- a/runtime/runtime_options.h +++ b/runtime/runtime_options.h @@ -203,6 +203,13 @@ struct RuntimeOptions { // not what the planner expected (e.g. a map that was declared as a proto or // a different message with matching field names). bool enable_typed_field_access = true; + + // When enabled, trivial lazy subexpressions (such as constants or + // identifiers) bound via cel.bind or cel.@block will be inlined directly + // at reference sites during planning. + // + // Note: Enabling this option may change observed expression tracing behavior. + bool enable_subexpression_inlining = false; }; // LINT.ThenChange(//depot/google3/eval/public/cel_options.h)