Skip to content

Commit 089cf93

Browse files
jnthntatumcopybara-github
authored andcommitted
Inline trivial lazy subexpressions when planning CEL expressions.
PiperOrigin-RevId: 1005954927
1 parent ea48c57 commit 089cf93

11 files changed

Lines changed: 382 additions & 9 deletions

‎eval/compiler/BUILD‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ cc_test(
6161
":resolver",
6262
"//common:expr",
6363
"//common:type",
64+
"//common:value",
6465
"//eval/eval:evaluator_core",
6566
"//eval/eval:expression_step_logic",
6667
"//internal:testing",

‎eval/compiler/flat_expr_builder.cc‎

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -818,16 +818,24 @@ class FlatExprVisitor : public cel::AstVisitor {
818818
SlotLookupResult slot = LookupSlot(path);
819819

820820
if (slot.subexpression >= 0) {
821-
auto* subexpression =
821+
const ProgramBuilder::Subexpression* subexpression =
822822
program_builder_.GetExtractedSubexpression(slot.subexpression);
823823
if (subexpression == nullptr) {
824824
SetProgressStatusIfError(
825825
absl::InternalError("bad subexpression reference"));
826826
return;
827827
}
828-
// Off by one since mainline expression will be index 0.
829-
AddStep(ExpressionStep::MakeLazyInitStep(
830-
slot.slot, slot.subexpression + 1, expr.id()));
828+
if (options_.enable_subexpression_inlining) {
829+
if (std::optional<ExpressionStep> inlined_step =
830+
subexpression->MaybeInline(expr.id());
831+
inlined_step.has_value()) {
832+
AddStep(*std::move(inlined_step));
833+
return;
834+
}
835+
}
836+
// Off by one since mainline expression will be index 0.
837+
AddStep(ExpressionStep::MakeLazyInitStep(
838+
slot.slot, slot.subexpression + 1, expr.id()));
831839
return;
832840
} else if (slot.slot >= 0) {
833841
AddStep(ExpressionStep::MakeReadSlotStep(slot.slot, expr.id()));

‎eval/compiler/flat_expr_builder_extensions.cc‎

Lines changed: 62 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,10 @@
1414
#include "eval/compiler/flat_expr_builder_extensions.h"
1515

1616
#include <cstddef>
17+
#include <cstdint>
1718
#include <iterator>
1819
#include <memory>
20+
#include <optional>
1921
#include <utility>
2022
#include <vector>
2123

@@ -28,14 +30,47 @@
2830
#include "absl/status/statusor.h"
2931
#include "absl/types/variant.h"
3032
#include "common/expr.h"
33+
#include "common/value.h"
3134
#include "eval/eval/evaluator_core.h"
35+
#include "eval/eval/ident_step.h"
3236

3337
namespace google::api::expr::runtime {
3438

3539
namespace {
3640

3741
using Subexpression = google::api::expr::runtime::ProgramBuilder::Subexpression;
3842

43+
const ExpressionStep* absl_nullable TryGetSingleStep(
44+
const Subexpression* absl_nonnull subexpr) {
45+
while (subexpr != nullptr) {
46+
if (subexpr->IsFlattened()) {
47+
const ExecutionPath& elements = subexpr->flattened_elements();
48+
if (elements.size() != 1) {
49+
return nullptr;
50+
}
51+
return &elements[0];
52+
}
53+
const auto& elements = subexpr->elements();
54+
if (elements.size() != 1) {
55+
return nullptr;
56+
}
57+
const auto& element = elements[0];
58+
if (const ExpressionStep* step = absl::get_if<ExpressionStep>(&element);
59+
step != nullptr) {
60+
return step;
61+
}
62+
// This might happen if an AST node is interpreted as a no-op, but we still
63+
// create a subexpression for it.
64+
if (Subexpression* const* child = absl::get_if<Subexpression*>(&element);
65+
child != nullptr) {
66+
subexpr = *child;
67+
continue;
68+
}
69+
return nullptr;
70+
}
71+
return nullptr;
72+
}
73+
3974
} // namespace
4075

4176
Subexpression::Subexpression(const cel::Expr* self, ProgramBuilder* owner)
@@ -185,7 +220,33 @@ bool Subexpression::ExtractTo(ExecutionPath& out) {
185220
return true;
186221
}
187222

188-
ExecutionPath ProgramBuilder::FlattenSubexpression(Subexpression* expr) {
223+
bool Subexpression::IsInlinable() const {
224+
const ExpressionStep* step = TryGetSingleStep(this);
225+
if (step == nullptr) {
226+
return false;
227+
}
228+
return IsConstant(*step) || GetIfIdentStep(*step) != nullptr;
229+
}
230+
231+
std::optional<ExpressionStep> Subexpression::MaybeInline(
232+
int64_t expr_id) const {
233+
const ExpressionStep* step = TryGetSingleStep(this);
234+
if (step == nullptr) {
235+
return std::nullopt;
236+
}
237+
cel::Value value;
238+
if (GetIfConstant(*step, value)) {
239+
return ExpressionStep::MakeConstant(value, expr_id);
240+
}
241+
if (const IdentStepInfo* ident = GetIfIdentStep(*step); ident != nullptr) {
242+
return ExpressionStep::MakeIdentifierStep(ident->ident_index,
243+
ident->slot_index, expr_id);
244+
}
245+
return std::nullopt;
246+
}
247+
248+
ExecutionPath ProgramBuilder::FlattenSubexpression(
249+
Subexpression* absl_nonnull expr) {
189250
ExecutionPath out;
190251

191252
if (!expr) {

‎eval/compiler/flat_expr_builder_extensions.h‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,14 @@ class ProgramBuilder {
179179
// Returns false if the subexpression is not currently flattened.
180180
bool ExtractTo(ExecutionPath& out);
181181

182+
// Returns true if this subexpression is trivial (reduces to a single
183+
// constant or identifier step) and can be inlined at reference sites.
184+
bool IsInlinable() const;
185+
186+
// Creates an inlined ExpressionStep for this subexpression with the
187+
// given expr_id, or returns std::nullopt if !IsInlinable().
188+
std::optional<ExpressionStep> MaybeInline(int64_t expr_id) const;
189+
182190
private:
183191
Subexpression(const cel::Expr* self, ProgramBuilder* owner);
184192

@@ -256,6 +264,15 @@ class ProgramBuilder {
256264
return extracted_subexpressions_[index];
257265
}
258266

267+
const Subexpression* absl_nullable GetExtractedSubexpression(
268+
size_t index) const {
269+
if (index >= extracted_subexpressions_.size()) {
270+
return nullptr;
271+
}
272+
273+
return extracted_subexpressions_[index];
274+
}
275+
259276
// Return index to the extracted subexpression.
260277
//
261278
// Returns -1 if the subexpression is not found.

‎eval/compiler/flat_expr_builder_extensions_test.cc‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,10 +23,12 @@
2323
#include "absl/status/statusor.h"
2424
#include "common/expr.h"
2525
#include "common/type.h"
26+
#include "common/value.h"
2627
#include "eval/compiler/resolver.h"
2728
#include "eval/eval/evaluator_core.h"
2829
#include "eval/eval/expression_step_logic.h"
2930
#include "eval/eval/function_step.h"
31+
#include "eval/eval/ident_step.h"
3032
#include "internal/testing.h"
3133
#include "runtime/function_registry.h"
3234
#include "runtime/internal/issue_collector.h"
@@ -539,5 +541,100 @@ TEST_F(ProgramBuilderTest, ExtractToRequiresFlatten) {
539541
UniquePtrHolds(step_ptrs.a)));
540542
}
541543

544+
TEST_F(ProgramBuilderTest, IsInlinableAndMaybeInlineConstant) {
545+
Expr a;
546+
Expr b;
547+
ProgramBuilder program_builder;
548+
549+
program_builder.EnterSubexpression(&a);
550+
program_builder.EnterSubexpression(&b);
551+
program_builder.AddStep(ExpressionStep::MakeConstant(cel::IntValue(42), 10));
552+
program_builder.ExitSubexpression(&b);
553+
program_builder.ExitSubexpression(&a);
554+
555+
Subexpression* subexpr_a = program_builder.GetSubexpression(&a);
556+
Subexpression* subexpr_b = program_builder.GetSubexpression(&b);
557+
ASSERT_TRUE(subexpr_a != nullptr);
558+
ASSERT_TRUE(subexpr_b != nullptr);
559+
560+
// Unflattened direct constant step and nested single-child subexpression.
561+
EXPECT_TRUE(subexpr_b->IsInlinable());
562+
EXPECT_TRUE(subexpr_a->IsInlinable());
563+
564+
std::optional<ExpressionStep> inlined_b = subexpr_b->MaybeInline(99);
565+
ASSERT_TRUE(inlined_b.has_value());
566+
EXPECT_EQ(inlined_b->id(), 99);
567+
cel::Value val;
568+
ASSERT_TRUE(GetIfConstant(*inlined_b, val));
569+
EXPECT_EQ(val.GetInt().NativeValue(), 42);
570+
571+
// Flattened constant step remains inlinable.
572+
subexpr_a->Flatten();
573+
EXPECT_TRUE(subexpr_a->IsInlinable());
574+
std::optional<ExpressionStep> inlined_a = subexpr_a->MaybeInline(100);
575+
ASSERT_TRUE(inlined_a.has_value());
576+
EXPECT_EQ(inlined_a->id(), 100);
577+
ASSERT_TRUE(GetIfConstant(*inlined_a, val));
578+
EXPECT_EQ(val.GetInt().NativeValue(), 42);
579+
}
580+
581+
TEST_F(ProgramBuilderTest, IsInlinableAndMaybeInlineIdentifier) {
582+
Expr a;
583+
ProgramBuilder program_builder;
584+
585+
program_builder.EnterSubexpression(&a);
586+
program_builder.AddStep(ExpressionStep::MakeIdentifierStep(3, 7, 10));
587+
program_builder.ExitSubexpression(&a);
588+
589+
Subexpression* subexpr_a = program_builder.GetSubexpression(&a);
590+
ASSERT_TRUE(subexpr_a != nullptr);
591+
592+
EXPECT_TRUE(subexpr_a->IsInlinable());
593+
std::optional<ExpressionStep> inlined = subexpr_a->MaybeInline(55);
594+
ASSERT_TRUE(inlined.has_value());
595+
EXPECT_EQ(inlined->id(), 55);
596+
const IdentStepInfo* ident_info = GetIfIdentStep(*inlined);
597+
ASSERT_TRUE(ident_info != nullptr);
598+
EXPECT_EQ(ident_info->ident_index, 3);
599+
EXPECT_EQ(ident_info->slot_index, 7);
600+
601+
subexpr_a->Flatten();
602+
EXPECT_TRUE(subexpr_a->IsInlinable());
603+
std::optional<ExpressionStep> inlined_flat = subexpr_a->MaybeInline(56);
604+
ASSERT_TRUE(inlined_flat.has_value());
605+
EXPECT_EQ(inlined_flat->id(), 56);
606+
ident_info = GetIfIdentStep(*inlined_flat);
607+
ASSERT_TRUE(ident_info != nullptr);
608+
EXPECT_EQ(ident_info->ident_index, 3);
609+
EXPECT_EQ(ident_info->slot_index, 7);
610+
}
611+
612+
TEST_F(ProgramBuilderTest, NotInlinableWhenNonTrivial) {
613+
Expr a;
614+
Expr b;
615+
Expr c;
616+
ProgramBuilder program_builder;
617+
618+
ASSERT_OK_AND_ASSIGN(SimpleTreeSteps step_ptrs,
619+
InitSimpleTree(a, b, c, program_builder));
620+
(void)step_ptrs;
621+
622+
Subexpression* subexpr_a = program_builder.GetSubexpression(&a);
623+
Subexpression* subexpr_b = program_builder.GetSubexpression(&b);
624+
ASSERT_TRUE(subexpr_a != nullptr);
625+
ASSERT_TRUE(subexpr_b != nullptr);
626+
627+
// Single generic step is not inlinable.
628+
EXPECT_FALSE(subexpr_b->IsInlinable());
629+
EXPECT_EQ(subexpr_b->MaybeInline(1), std::nullopt);
630+
631+
// Multi-step subexpression is not inlinable (unflattened and flattened).
632+
EXPECT_FALSE(subexpr_a->IsInlinable());
633+
EXPECT_EQ(subexpr_a->MaybeInline(1), std::nullopt);
634+
subexpr_a->Flatten();
635+
EXPECT_FALSE(subexpr_a->IsInlinable());
636+
EXPECT_EQ(subexpr_a->MaybeInline(1), std::nullopt);
637+
}
638+
542639
} // namespace
543640
} // namespace google::api::expr::runtime

0 commit comments

Comments
 (0)