From 9e4b7a2c8d21ff70c1d59d575210beea8418c119 Mon Sep 17 00:00:00 2001 From: Jonathan Tatum Date: Mon, 19 Oct 2026 16:26:08 -0700 Subject: [PATCH] Add tighter bound on value stack estimate. PiperOrigin-RevId: 1002674949 --- eval/compiler/BUILD | 6 +- eval/compiler/flat_expr_builder.cc | 116 ++++++++-- eval/compiler/flat_expr_builder_extensions.h | 6 +- eval/compiler/flat_expr_builder_test.cc | 220 ++++++++++++++++++- eval/eval/BUILD | 4 + eval/eval/evaluator_core.cc | 73 +++++- eval/eval/evaluator_core.h | 28 ++- eval/eval/evaluator_core_test.cc | 96 ++++++++ eval/eval/function_step.h | 3 + extensions/select_optimization.cc | 12 +- 10 files changed, 518 insertions(+), 46 deletions(-) diff --git a/eval/compiler/BUILD b/eval/compiler/BUILD index 746b6af8f..f57a59ebf 100644 --- a/eval/compiler/BUILD +++ b/eval/compiler/BUILD @@ -122,6 +122,7 @@ cc_library( "//runtime/internal:runtime_env", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base:nullability", + "@com_google_absl//absl/cleanup", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/container:node_hash_map", @@ -146,6 +147,7 @@ cc_test( ":cel_expression_builder_flat_impl", ":constant_folding", ":flat_expr_builder", + ":flat_expr_builder_extensions", ":qualified_reference_resolver", "//base:builtins", "//checker:validation_result", @@ -159,8 +161,8 @@ cc_test( "//compiler", "//compiler:compiler_factory", "//compiler:standard_library", - "//eval/compiler:flat_expr_builder_extensions", "//eval/eval:evaluator_core", + "//eval/eval:expression_step_logic", "//eval/public:activation", "//eval/public:ast_rewrite", "//eval/public:builtin_func_registrar", @@ -181,6 +183,7 @@ cc_test( "//eval/public/structs:cel_proto_wrapper", "//eval/public/testing:matchers", "//eval/testutil:test_message_cc_proto", + "//extensions/protobuf:ast_converters", "//internal:proto_matchers", "//internal:status_macros", "//internal:testing", @@ -202,7 +205,6 @@ cc_test( "@com_google_cel_spec//proto/cel/expr:checked_cc_proto", "@com_google_cel_spec//proto/cel/expr:syntax_cc_proto", "@com_google_cel_spec//proto/cel/expr/conformance/proto3:test_all_types_cc_proto", - "@com_google_protobuf//:field_mask_cc_proto", "@com_google_protobuf//:protobuf", ], ) diff --git a/eval/compiler/flat_expr_builder.cc b/eval/compiler/flat_expr_builder.cc index 7270ee42e..ec15a759d 100644 --- a/eval/compiler/flat_expr_builder.cc +++ b/eval/compiler/flat_expr_builder.cc @@ -32,6 +32,7 @@ #include "absl/algorithm/container.h" #include "absl/base/nullability.h" +#include "absl/cleanup/cleanup.h" #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" #include "absl/container/node_hash_map.h" @@ -971,13 +972,13 @@ class FlatExprVisitor : public cel::AstVisitor { *std::move(field_type), select_expr.test_only(), options_.enable_empty_wrapper_null_unboxing, enable_optional_types_), - expr.id()); + expr.id(), /*stack_delta=*/0); return; } AddStep(CreateSelectStep(std::move(field), select_expr.test_only(), options_.enable_empty_wrapper_null_unboxing, enable_optional_types_), - expr.id()); + expr.id(), /*stack_delta=*/0); } // Call node handler group. @@ -1383,9 +1384,14 @@ class FlatExprVisitor : public cel::AstVisitor { std::vector fields = std::move(status_or_resolved_fields.value().second); + size_t num_fields = fields.size(); + int64_t stack_delta = + num_fields <= static_cast(std::numeric_limits::max()) + ? 1 - static_cast(num_fields) + : std::numeric_limits::max(); AddStep(CreateCreateStructStep(std::move(resolved_name), std::move(fields), MakeOptionalIndicesSet(struct_expr)), - expr.id()); + expr.id(), stack_delta); } void PostVisitMap(const cel::Expr& expr, @@ -1406,9 +1412,15 @@ class FlatExprVisitor : public cel::AstVisitor { } } - AddStep(CreateCreateStructStepForMap(map_expr.entries().size(), + size_t num_entries = map_expr.entries().size(); + int64_t stack_delta = + num_entries <= + static_cast(std::numeric_limits::max() / 2) + ? 1 - 2 * static_cast(num_entries) + : std::numeric_limits::max(); + AddStep(CreateCreateStructStepForMap(num_entries, MakeOptionalIndicesSet(map_expr)), - expr.id()); + expr.id(), stack_delta); } absl::Status progress_status() const { return progress_status_; } @@ -1471,11 +1483,11 @@ class FlatExprVisitor : public cel::AstVisitor { // may free the step at that point. template std::enable_if_t, T*> AddStep( - std::unique_ptr step, int64_t expr_id = -1) { + std::unique_ptr step, int64_t expr_id = -1, int64_t stack_delta = 1) { if (progress_status_.ok() && !PlanningSuppressed()) { T* ptr = step.get(); - program_builder_.AddStep( - ExpressionStep::MakeGenericStep(std::move(step), expr_id)); + program_builder_.AddStep(ExpressionStep::MakeGenericStep( + std::move(step), expr_id, stack_delta)); return ptr; } return nullptr; @@ -1483,9 +1495,10 @@ class FlatExprVisitor : public cel::AstVisitor { template std::enable_if_t, T*> AddStep( - absl::StatusOr> step, int64_t expr_id = -1) { + absl::StatusOr> step, int64_t expr_id = -1, + int64_t stack_delta = 1) { if (step.ok()) { - return AddStep(*std::move(step), expr_id); + return AddStep(*std::move(step), expr_id, stack_delta); } else { SetProgressStatusIfError(step.status()); } @@ -1728,7 +1741,7 @@ FlatExprVisitor::CallHandlerResult FlatExprVisitor::HandleIndex( } AddStep(CreateContainerAccessStep(call_expr, enable_optional_types_), - expr.id()); + expr.id(), /*stack_delta=*/-1); return CallHandlerResult::kIntercepted; } @@ -2041,7 +2054,7 @@ void ExhaustiveTernaryCondVisitor::PreVisit(const cel::Expr* expr) { } void ExhaustiveTernaryCondVisitor::PostVisit(const cel::Expr* expr) { - visitor_->AddStep(CreateTernaryStep(), expr->id()); + visitor_->AddStep(CreateTernaryStep(), expr->id(), /*stack_delta=*/-2); } void ComprehensionVisitor::PreVisit(const cel::Expr* expr) { @@ -2232,6 +2245,52 @@ std::vector FlattenExpressionTable( return subexpression_indexes; } +std::optional CheckedDeltaAdd(int64_t current, + std::optional delta) { + if (!delta.has_value()) { + return std::nullopt; + } + if (*delta > 0 && current > std::numeric_limits::max() - *delta) { + return std::nullopt; + } + if (*delta < 0 && current < std::numeric_limits::min() - *delta) { + return std::nullopt; + } + current += *delta; + if (current < 0) { + return std::nullopt; + } + return current; +} + +// Conservative estimate of the maximum value stack size needed for the given +// subexpressions. +// +// If overflow occurs, returns fallback_size, which is the total number of +// steps in the program. +size_t EstimateMaxStackSize(absl::Span subexpressions, + size_t fallback_size) { + size_t total_max_stack = 0; + for (ExecutionPathView path : subexpressions) { + int64_t current = 0; + int64_t max_depth = 0; + for (const ExpressionStep& step : path) { + std::optional next = CheckedDeltaAdd(current, step.StackDelta()); + if (!next.has_value()) { + return fallback_size; + } + current = *next; + max_depth = std::max(max_depth, current); + } + if (static_cast(max_depth) > + std::numeric_limits::max() - total_max_stack) { + return fallback_size; + } + total_max_stack += static_cast(max_depth); + } + return total_max_stack; +} + absl::Status CheckAstExtensions( const std::vector& extensions, bool enable_select_optimization) { @@ -2266,11 +2325,27 @@ absl::StatusOr FlatExprBuilder::CreateExpressionImpl( ? RuntimeIssue::Severity::kWarning : RuntimeIssue::Severity::kError; IssueCollector issue_collector(max_severity); + auto record_issues = absl::MakeCleanup([issues, &issue_collector] { + if (issues == nullptr) { + return; + } + *issues = issue_collector.ExtractIssues(); + }); - CEL_ASSIGN_OR_RETURN(std::vector runtime_extensions, - ExtractAndValidateRuntimeExtensions(*ast)); - CEL_RETURN_IF_ERROR( - CheckAstExtensions(runtime_extensions, enable_select_optimization_)); + absl::StatusOr> runtime_extensions = + ExtractAndValidateRuntimeExtensions(*ast); + + if (!runtime_extensions.ok()) { + CEL_RETURN_IF_ERROR(issue_collector.AddIssue( + RuntimeIssue::CreateError(runtime_extensions.status()))); + } else { + absl::Status status = + CheckAstExtensions(*runtime_extensions, enable_select_optimization_); + if (!status.ok()) { + CEL_RETURN_IF_ERROR( + issue_collector.AddIssue(RuntimeIssue::CreateError(status))); + } + } Resolver resolver(container_, function_registry_, type_registry_, GetTypeProvider(), @@ -2318,17 +2393,16 @@ absl::StatusOr FlatExprBuilder::CreateExpressionImpl( "Expression too large to be executed, exceeds uint32_t limits."); } - if (issues != nullptr) { - (*issues) = issue_collector.ExtractIssues(); - } - ExecutionPath execution_path; std::vector subexpressions = FlattenExpressionTable(program_builder, execution_path); + size_t value_stack_size = + EstimateMaxStackSize(subexpressions, execution_path.size()); return FlatExpression(std::move(execution_path), std::move(subexpressions), visitor.slot_count(), visitor.iter_stack_count(), - GetTypeProvider(), options_, std::move(arena)); + GetTypeProvider(), options_, std::move(arena), + value_stack_size); } const cel::TypeProvider& FlatExprBuilder::GetTypeProvider() const { return use_legacy_type_provider_ diff --git a/eval/compiler/flat_expr_builder_extensions.h b/eval/compiler/flat_expr_builder_extensions.h index 90916f5b4..d6d3abb17 100644 --- a/eval/compiler/flat_expr_builder_extensions.h +++ b/eval/compiler/flat_expr_builder_extensions.h @@ -334,9 +334,9 @@ class PlannerContext { absl::Status AddSubplanStep(const cel::Expr& node, ExpressionStep step); absl::Status AddSubplanStep(const cel::Expr& node, std::unique_ptr step, - int64_t expr_id = -1) { - return AddSubplanStep( - node, ExpressionStep::MakeGenericStep(std::move(step), expr_id)); + int64_t expr_id = -1, int64_t stack_delta = 1) { + return AddSubplanStep(node, ExpressionStep::MakeGenericStep( + std::move(step), expr_id, stack_delta)); } const Resolver& resolver() const { return resolver_; } diff --git a/eval/compiler/flat_expr_builder_test.cc b/eval/compiler/flat_expr_builder_test.cc index 24539f4b7..da83692d5 100644 --- a/eval/compiler/flat_expr_builder_test.cc +++ b/eval/compiler/flat_expr_builder_test.cc @@ -16,6 +16,7 @@ #include #include +#include #include #include #include @@ -24,7 +25,6 @@ #include "cel/expr/checked.pb.h" #include "cel/expr/syntax.pb.h" -#include "google/protobuf/field_mask.pb.h" #include "google/protobuf/descriptor.pb.h" #include "absl/container/flat_hash_map.h" #include "absl/log/absl_check.h" @@ -52,6 +52,7 @@ #include "eval/compiler/flat_expr_builder_extensions.h" #include "eval/compiler/qualified_reference_resolver.h" #include "eval/eval/evaluator_core.h" +#include "eval/eval/expression_step_logic.h" #include "eval/public/activation.h" #include "eval/public/ast_rewrite.h" #include "eval/public/builtin_func_registrar.h" @@ -72,6 +73,7 @@ #include "eval/public/unknown_attribute_set.h" #include "eval/public/unknown_set.h" #include "eval/testutil/test_message.pb.h" +#include "extensions/protobuf/ast_converters.h" #include "internal/proto_matchers.h" #include "internal/status_macros.h" #include "internal/testing.h" @@ -203,7 +205,7 @@ TEST(FlatExprBuilderTest, RuntimeExtensionsError) { auto* ext = source_info.add_extensions(); ext->set_id("ext1"); ext->add_affected_components( - cel::expr::SourceInfo_Extension_Component_COMPONENT_RUNTIME); + cel::expr::SourceInfo::Extension::COMPONENT_RUNTIME); CelExpressionBuilderFlatImpl builder(NewTestingRuntimeEnv()); EXPECT_THAT(builder.CreateExpression(&expr, &source_info).status(), @@ -1207,8 +1209,8 @@ TEST(FlatExprBuilderTest, CheckedExprActivationMissesReferences) { std::vector> map_pairs{ {CelValue::CreateStringView("var2"), CelValue::CreateBool(false)}}; - std::unique_ptr map_value = - *CreateContainerBackedMap(absl::MakeSpan(map_pairs)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr map_value, + CreateContainerBackedMap(absl::MakeSpan(map_pairs))); activation.InsertValue("bar", CelValue::CreateMap(map_value.get())); ASSERT_OK_AND_ASSIGN(result, cel_expr->Evaluate(activation, &arena)); ASSERT_TRUE(result.IsBool()); @@ -3160,6 +3162,216 @@ TEST(FlatExprBuilderTest, EXPECT_THAT(result, test::IsCelBool(true)); } +TEST(FlatExprBuilderTest, StackSizeEstimationBuiltinsAndGenericSteps) { + cel::RuntimeOptions options; + options.short_circuiting = false; + CelExpressionBuilderFlatImpl builder(NewTestingRuntimeEnv(), options); + ASSERT_THAT(RegisterBuiltinFunctions(builder.GetRegistry()), IsOk()); + + { + ASSERT_OK_AND_ASSIGN(ParsedExpr parsed, parser::Parse("1 + 2 + 3 + 4")); + ASSERT_OK_AND_ASSIGN(std::unique_ptr ast, + cel::extensions::CreateAstFromParsedExpr(parsed)); + ASSERT_OK_AND_ASSIGN(FlatExpression expr, + builder.flat_expr_builder().CreateExpressionImpl( + std::move(ast), nullptr)); + EXPECT_EQ(expr.path().size(), 7); + EXPECT_EQ(expr.value_stack_size(), 2); + } + + { + ASSERT_OK_AND_ASSIGN(ParsedExpr parsed, + parser::Parse("{'a': 1, 'b': 2}['a']")); + ASSERT_OK_AND_ASSIGN(std::unique_ptr ast, + cel::extensions::CreateAstFromParsedExpr(parsed)); + ASSERT_OK_AND_ASSIGN(FlatExpression expr, + builder.flat_expr_builder().CreateExpressionImpl( + std::move(ast), nullptr)); + // 'a', 1, 'b', 2, CreateMap(2), 'a', ContainerAccess + EXPECT_EQ(expr.path().size(), 7); + EXPECT_EQ(expr.value_stack_size(), 4); + } + + { + ASSERT_OK_AND_ASSIGN( + ParsedExpr parsed, + parser::Parse("google.api.expr.runtime.TestMessage{int64_value: 1, " + "string_value: 'a'}.int64_value")); + ASSERT_OK_AND_ASSIGN(std::unique_ptr ast, + cel::extensions::CreateAstFromParsedExpr(parsed)); + ASSERT_OK_AND_ASSIGN(FlatExpression expr, + builder.flat_expr_builder().CreateExpressionImpl( + std::move(ast), nullptr)); + // 1, 'a', CreateStruct(2), Select + EXPECT_EQ(expr.path().size(), 4); + EXPECT_EQ(expr.value_stack_size(), 2); + } + + { + ASSERT_OK_AND_ASSIGN(ParsedExpr parsed, parser::Parse("true ? 1 : 2")); + ASSERT_OK_AND_ASSIGN(std::unique_ptr ast, + cel::extensions::CreateAstFromParsedExpr(parsed)); + ASSERT_OK_AND_ASSIGN(FlatExpression expr, + builder.flat_expr_builder().CreateExpressionImpl( + std::move(ast), nullptr)); + // true, 1, 2, TernaryStep + EXPECT_EQ(expr.path().size(), 4); + EXPECT_EQ(expr.value_stack_size(), 3); + } +} + +TEST(FlatExprBuilderTest, StackSizeEstimationShortCircuiting) { + cel::RuntimeOptions options; + options.short_circuiting = true; + CelExpressionBuilderFlatImpl builder(NewTestingRuntimeEnv(), options); + ASSERT_THAT(RegisterBuiltinFunctions(builder.GetRegistry()), IsOk()); + + { + // Conditionals over-count in linear stack estimation because TernaryJump + // and FixedJump have a worst-case delta of 0 and both branches are walked + // sequentially without a merge step. + ASSERT_OK_AND_ASSIGN(ParsedExpr parsed, + parser::Parse("(c1 ? 1 : 2) + (c2 ? 3 : 4)")); + ASSERT_OK_AND_ASSIGN(std::unique_ptr ast, + cel::extensions::CreateAstFromParsedExpr(parsed)); + ASSERT_OK_AND_ASSIGN(FlatExpression expr, + builder.flat_expr_builder().CreateExpressionImpl( + std::move(ast), nullptr)); + // c1, TernaryJump, 1, FixedJump, 2, c2, TernaryJump, 3, FixedJump, 4, + + EXPECT_EQ(expr.path().size(), 11); + EXPECT_EQ(expr.value_stack_size(), 6); + } + + { + // Comprehension: [1, 2, 3].all(x, x > 0) + ASSERT_OK_AND_ASSIGN(ParsedExpr parsed, + parser::Parse("[1, 2, 3].all(x, x > 0)")); + ASSERT_OK_AND_ASSIGN(std::unique_ptr ast, + cel::extensions::CreateAstFromParsedExpr(parsed)); + ASSERT_OK_AND_ASSIGN(FlatExpression expr, + builder.flat_expr_builder().CreateExpressionImpl( + std::move(ast), nullptr)); + EXPECT_EQ(expr.path().size(), 19); + EXPECT_EQ(expr.value_stack_size(), 5); + } + + { + // Mixed logical operators: (a && b) || (c && d) + ASSERT_OK_AND_ASSIGN(ParsedExpr parsed, + parser::Parse("(a && b) || (c && d)")); + ASSERT_OK_AND_ASSIGN(std::unique_ptr ast, + cel::extensions::CreateAstFromParsedExpr(parsed)); + ASSERT_OK_AND_ASSIGN(FlatExpression expr, + builder.flat_expr_builder().CreateExpressionImpl( + std::move(ast), nullptr)); + // a, AndJump, b, And(2), OrJump, c, AndJump, d, And(2), Or(2) + EXPECT_EQ(expr.path().size(), 10); + EXPECT_EQ(expr.value_stack_size(), 3); + } +} + +TEST(FlatExprBuilderTest, StackSizeEstimationWithLazySubexpressions) { + ParsedExpr parsed_expr; + ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString( + R"pb( + expr: { + call_expr: { + function: "cel.@block" + args { + list_expr: { + elements { + call_expr: { + function: "_+_" + args { const_expr: { int64_value: 1 } } + args { const_expr: { int64_value: 2 } } + } + } + elements { + call_expr: { + function: "_+_" + args { const_expr: { int64_value: 3 } } + args { const_expr: { int64_value: 4 } } + } + } + } + } + args { + call_expr: { + function: "_+_" + args { ident_expr: { name: "@index0" } } + args { ident_expr: { name: "@index1" } } + } + } + } + } + )pb", + &parsed_expr)); + + CelExpressionBuilderFlatImpl builder(NewTestingRuntimeEnv()); + ASSERT_THAT(RegisterBuiltinFunctions(builder.GetRegistry()), IsOk()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr ast, + cel::extensions::CreateAstFromParsedExpr(parsed_expr)); + ASSERT_OK_AND_ASSIGN(FlatExpression expr, + builder.flat_expr_builder().CreateExpressionImpl( + std::move(ast), nullptr)); + // Main program: LazyInit(@index0), LazyInit(@index1), +, ClearSlots -> max 2 + // Subexpr 1: 1, 2, + -> max 2 + // Subexpr 2: 3, 4, + -> max 2 + // Total max stack: 2 + 2 + 2 = 6 (while total steps = 10) + EXPECT_EQ(expr.path().size(), 10); + EXPECT_EQ(expr.value_stack_size(), 6); +} + +class NoOpStepWithCustomDelta : public ExpressionStepLogic { + public: + void Evaluate(ExecutionFrame* frame) const override {} +}; + +class CustomDeltaInjectorOptimizer : public ProgramOptimizer { + public: + explicit CustomDeltaInjectorOptimizer(int64_t stack_delta) + : stack_delta_(stack_delta) {} + + absl::Status OnPreVisit(PlannerContext& context, + const cel::Expr& node) override { + return absl::OkStatus(); + } + + absl::Status OnPostVisit(PlannerContext& context, + const cel::Expr& node) override { + if (node.id() == 1) { + return context.AddSubplanStep( + node, std::make_unique(), -1, stack_delta_); + } + return absl::OkStatus(); + } + + private: + int64_t stack_delta_; +}; + +TEST(FlatExprBuilderTest, StackSizeEstimationOverflowAndUnderflowFallback) { + for (int64_t bad_delta : + {static_cast(std::numeric_limits::max()), + static_cast(-10)}) { + CelExpressionBuilderFlatImpl builder(NewTestingRuntimeEnv()); + ASSERT_THAT(RegisterBuiltinFunctions(builder.GetRegistry()), IsOk()); + builder.flat_expr_builder().AddProgramOptimizer( + [bad_delta](PlannerContext&, const cel::Ast&) + -> absl::StatusOr> { + return std::make_unique(bad_delta); + }); + + ASSERT_OK_AND_ASSIGN(ParsedExpr parsed, parser::Parse("1 + 2 + 3")); + ASSERT_OK_AND_ASSIGN(std::unique_ptr ast, + cel::extensions::CreateAstFromParsedExpr(parsed)); + ASSERT_OK_AND_ASSIGN(FlatExpression expr, + builder.flat_expr_builder().CreateExpressionImpl( + std::move(ast), nullptr)); + EXPECT_EQ(expr.value_stack_size(), expr.path().size()); + } +} + TEST(FlatExprBuilderTest, ListAppendComprehensionWithTargetFailsPlanning) { ParsedExpr parsed_expr; ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString( diff --git a/eval/eval/BUILD b/eval/eval/BUILD index b477b35f7..26d0561d1 100644 --- a/eval/eval/BUILD +++ b/eval/eval/BUILD @@ -418,6 +418,7 @@ cc_test( deps = [ ":cel_expression_flat_impl", ":evaluator_core", + ":expression_step_logic", "//base:data", "//common:value", "//eval/compiler:cel_expression_builder_flat_impl", @@ -429,9 +430,12 @@ cc_test( "//internal:testing_descriptor_pool", "//internal:testing_message_factory", "//runtime:activation", + "//runtime:function_overload_reference", + "//runtime:function_registry", "//runtime:runtime_options", "//runtime/internal:runtime_env_testing", "//runtime/internal:runtime_type_provider", + "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/status", "@com_google_cel_spec//proto/cel/expr:syntax_cc_proto", "@com_google_protobuf//:protobuf", diff --git a/eval/eval/evaluator_core.cc b/eval/eval/evaluator_core.cc index 2fc6eda93..1d8e13b3f 100644 --- a/eval/eval/evaluator_core.cc +++ b/eval/eval/evaluator_core.cc @@ -118,6 +118,12 @@ class EvaluationStatus final { alignas(absl::Status) char status_[sizeof(absl::Status)]; }; +int64_t OneMinusArgCount(size_t count) { + ABSL_DCHECK(count <= + static_cast(std::numeric_limits::max())); + return 1 - static_cast(count); +} + } // namespace void ExpressionStep::EvaluateMutableListAppendStep(ExecutionFrame& frame) { @@ -221,17 +227,17 @@ FlatExpressionEvaluatorState FlatExpression::MakeEvaluatorState( const google::protobuf::DescriptorPool* absl_nonnull descriptor_pool, google::protobuf::MessageFactory* absl_nonnull message_factory, google::protobuf::Arena* absl_nonnull arena) const { - return FlatExpressionEvaluatorState(path_.size(), comprehension_slots_size_, - iterator_stack_size_, type_provider_, - descriptor_pool, message_factory, arena); + return FlatExpressionEvaluatorState( + value_stack_size_, comprehension_slots_size_, iterator_stack_size_, + type_provider_, descriptor_pool, message_factory, arena); } FlatExpressionEvaluatorState FlatExpression::MakeEvaluatorState( const google::protobuf::DescriptorPool* absl_nonnull descriptor_pool, google::protobuf::MessageFactory* absl_nullable message_factory) const { - return FlatExpressionEvaluatorState(path_.size(), comprehension_slots_size_, - iterator_stack_size_, type_provider_, - descriptor_pool, message_factory); + return FlatExpressionEvaluatorState( + value_stack_size_, comprehension_slots_size_, iterator_stack_size_, + type_provider_, descriptor_pool, message_factory); } absl::StatusOr FlatExpression::EvaluateWithCallback( @@ -360,4 +366,59 @@ FixedJumpStepInfo* GetIfFixedJumpStep(ExpressionStep& step) { return nullptr; } +std::optional ExpressionStep::StackDelta() const { + switch (header_.kind) { + case ExpressionStepKind::kMovedFrom: + return 0; + case ExpressionStepKind::kGenericLogic: + if (header_.stack_delta == std::numeric_limits::max()) { + // Impractical, but technically allowed. + return std::nullopt; + } + return header_.stack_delta; + case ExpressionStepKind::kIntConstant: + case ExpressionStepKind::kBoolConstant: + case ExpressionStepKind::kDoubleConstant: + case ExpressionStepKind::kNullConstant: + case ExpressionStepKind::kUintConstant: + case ExpressionStepKind::kOtherConstant: + case ExpressionStepKind::kLazyInit: + case ExpressionStepKind::kReadSlot: + case ExpressionStepKind::kIdentifier: + case ExpressionStepKind::kNewMutableList: + return 1; + case ExpressionStepKind::kAssignSlotAndPop: + case ExpressionStepKind::kComprehensionFinish: + case ExpressionStepKind::kComprehensionNext: + case ExpressionStepKind::kComprehensionNext2: + case ExpressionStepKind::kComprehensionCond: + case ExpressionStepKind::kComprehensionCond2: + case ExpressionStepKind::kFastIn: + case ExpressionStepKind::kFastEqual: + case ExpressionStepKind::kFastNotEqual: + case ExpressionStepKind::kMutableListAppend: + return -1; + case ExpressionStepKind::kClearSlots: + case ExpressionStepKind::kBooleanNot: + case ExpressionStepKind::kNotStrictlyFalse: + case ExpressionStepKind::kBooleanOrJump: + case ExpressionStepKind::kBooleanAndJump: + case ExpressionStepKind::kTernaryJump: + case ExpressionStepKind::kFixedJump: + return 0; + case ExpressionStepKind::kBooleanOr: + case ExpressionStepKind::kBooleanAnd: + return OneMinusArgCount(u_.arg_count); + case ExpressionStepKind::kEagerFunction: + return OneMinusArgCount(u_.eager_function_step->num_arguments()); + case ExpressionStepKind::kLazyFunction: + return OneMinusArgCount(u_.lazy_function_step->num_arguments()); + case ExpressionStepKind::kCreateList: + return OneMinusArgCount(u_.create_list_step->num_elements()); + case ExpressionStepKind::kCreateSmallList: + return OneMinusArgCount(u_.create_small_list_step.list_size); + } + return std::nullopt; +} + } // namespace google::api::expr::runtime diff --git a/eval/eval/evaluator_core.h b/eval/eval/evaluator_core.h index b389c38a6..efb6e4461 100644 --- a/eval/eval/evaluator_core.h +++ b/eval/eval/evaluator_core.h @@ -19,6 +19,7 @@ #include #include #include +#include #include #include #include @@ -166,9 +167,20 @@ class ExpressionStep { const ExpressionStepLogic* GetGenericStep() const; bool IsGenericStep() const; + // Returns the worst-case change in value stack size when evaluating this + // step, or std::nullopt if the delta overflows. + std::optional StackDelta() const; + static ExpressionStep MakeGenericStep( - std::unique_ptr logic, int64_t id = -1) { + std::unique_ptr logic, int64_t id = -1, + int64_t stack_delta = 1) { ExpressionStep step(ExpressionStepKind::kGenericLogic, id); + if (stack_delta < std::numeric_limits::min() || + stack_delta >= std::numeric_limits::max()) { + step.header_.stack_delta = std::numeric_limits::max(); + } else { + step.header_.stack_delta = static_cast(stack_delta); + } step.u_.logic = logic.release(); return step; } @@ -365,7 +377,8 @@ class ExpressionStep { struct Header { ExpressionStepKind kind; - uint16_t reserved; + // Note: This can be moved if steps are all migrated. + int16_t stack_delta; int32_t id; }; @@ -869,11 +882,13 @@ class FlatExpression { FlatExpression(ExecutionPath path, size_t comprehension_slots_size, const cel::TypeProvider& type_provider, const cel::RuntimeOptions& options, - absl_nullable std::shared_ptr arena = nullptr) + absl_nullable std::shared_ptr arena = nullptr, + std::optional value_stack_size = std::nullopt) : path_(std::move(path)), subexpressions_({path_}), comprehension_slots_size_(comprehension_slots_size), iterator_stack_size_(comprehension_slots_size), + value_stack_size_(value_stack_size.value_or(path_.size())), type_provider_(type_provider), options_(options), arena_(std::move(arena)) {} @@ -883,11 +898,13 @@ class FlatExpression { size_t comprehension_slots_size, size_t iterator_stack_size, const cel::TypeProvider& type_provider, const cel::RuntimeOptions& options, - absl_nullable std::shared_ptr arena = nullptr) + absl_nullable std::shared_ptr arena = nullptr, + std::optional value_stack_size = std::nullopt) : path_(std::move(path)), subexpressions_(std::move(subexpressions)), comprehension_slots_size_(comprehension_slots_size), iterator_stack_size_(iterator_stack_size), + value_stack_size_(value_stack_size.value_or(path_.size())), type_provider_(type_provider), options_(options), arena_(std::move(arena)) {} @@ -932,6 +949,8 @@ class FlatExpression { size_t iterator_stack_size() const { return iterator_stack_size_; } + size_t value_stack_size() const { return value_stack_size_; } + const cel::TypeProvider& type_provider() const { return type_provider_; } private: @@ -939,6 +958,7 @@ class FlatExpression { std::vector subexpressions_; size_t comprehension_slots_size_; size_t iterator_stack_size_; + size_t value_stack_size_; const cel::TypeProvider& type_provider_; cel::RuntimeOptions options_; // Arena used during planning phase, may hold constant values so should be diff --git a/eval/eval/evaluator_core_test.cc b/eval/eval/evaluator_core_test.cc index 1f43f3f35..061db087d 100644 --- a/eval/eval/evaluator_core_test.cc +++ b/eval/eval/evaluator_core_test.cc @@ -1,17 +1,23 @@ #include "eval/eval/evaluator_core.h" +#include #include #include #include +#include #include #include #include "cel/expr/syntax.pb.h" +#include "absl/container/flat_hash_set.h" #include "absl/status/status.h" #include "base/type_provider.h" #include "common/value.h" #include "eval/compiler/cel_expression_builder_flat_impl.h" #include "eval/eval/cel_expression_flat_impl.h" +#include "eval/eval/create_list_step.h" +#include "eval/eval/expression_step_logic.h" +#include "eval/eval/function_step.h" #include "eval/internal/interop.h" #include "eval/public/activation.h" #include "eval/public/builtin_func_registrar.h" @@ -20,6 +26,8 @@ #include "internal/testing_descriptor_pool.h" #include "internal/testing_message_factory.h" #include "runtime/activation.h" +#include "runtime/function_overload_reference.h" +#include "runtime/function_registry.h" #include "runtime/internal/runtime_env_testing.h" #include "runtime/internal/runtime_type_provider.h" #include "runtime/runtime_options.h" @@ -390,4 +398,92 @@ TEST(EvaluatorCoreTest, TraceTest) { ASSERT_THAT(eval_status, IsOk()); } +TEST(EvaluatorCoreTest, StepStackDelta) { + EXPECT_EQ(ExpressionStep::MakeConstant(cel::IntValue(1)).StackDelta(), 1); + EXPECT_EQ(ExpressionStep::MakeConstant(cel::BoolValue(true)).StackDelta(), 1); + EXPECT_EQ(ExpressionStep::MakeConstant(cel::DoubleValue(1.0)).StackDelta(), + 1); + EXPECT_EQ(ExpressionStep::MakeConstant(cel::NullValue()).StackDelta(), 1); + EXPECT_EQ(ExpressionStep::MakeConstant(cel::UintValue(1)).StackDelta(), 1); + EXPECT_EQ( + ExpressionStep::MakeConstant(cel::StringValue::Literal("a")).StackDelta(), + 1); + + ExpressionStep moved_step = ExpressionStep::MakeConstant(cel::IntValue(1)); + ExpressionStep dest_step = std::move(moved_step); + EXPECT_EQ(moved_step.StackDelta(), 0); // NOLINT(bugprone-use-after-move) + EXPECT_EQ(dest_step.StackDelta(), 1); + + EXPECT_EQ(ExpressionStep::MakeLazyInitStep(0, 1).StackDelta(), 1); + EXPECT_EQ(ExpressionStep::MakeAssignSlotAndPopStep(0).StackDelta(), -1); + EXPECT_EQ(ExpressionStep::MakeClearSlotsStep(0, 2).StackDelta(), 0); + EXPECT_EQ(ExpressionStep::MakeReadSlotStep(0).StackDelta(), 1); + + EXPECT_EQ(ExpressionStep::MakeBooleanNotStep().StackDelta(), 0); + EXPECT_EQ(ExpressionStep::MakeNotStrictlyFalseStep().StackDelta(), 0); + EXPECT_EQ(ExpressionStep::MakeBooleanOrStep(3).StackDelta(), -2); + EXPECT_EQ(ExpressionStep::MakeBooleanAndStep(2).StackDelta(), -1); + + EXPECT_EQ(ExpressionStep::MakeComprehensionFinishStep(0).StackDelta(), -1); + EXPECT_EQ(ExpressionStep::MakeComprehensionNextStep().StackDelta(), -1); + EXPECT_EQ(ExpressionStep::MakeComprehensionNext2Step().StackDelta(), -1); + EXPECT_EQ(ExpressionStep::MakeComprehensionCondStep().StackDelta(), -1); + EXPECT_EQ(ExpressionStep::MakeComprehensionCond2Step().StackDelta(), -1); + + EXPECT_EQ(ExpressionStep::MakeBooleanOrJumpStep(2).StackDelta(), 0); + EXPECT_EQ(ExpressionStep::MakeBooleanAndJumpStep(2).StackDelta(), 0); + EXPECT_EQ(ExpressionStep::MakeTernaryJumpStep().StackDelta(), 0); + EXPECT_EQ(ExpressionStep::MakeFixedJumpStep().StackDelta(), 0); + + EXPECT_EQ(ExpressionStep::MakeIdentifierStep("x").StackDelta(), 1); + EXPECT_EQ(ExpressionStep::MakeFastInStep().StackDelta(), -1); + EXPECT_EQ(ExpressionStep::MakeFastEqualStep().StackDelta(), -1); + EXPECT_EQ(ExpressionStep::MakeFastNotEqualStep().StackDelta(), -1); + EXPECT_EQ(ExpressionStep::MakeNewMutableListStep().StackDelta(), 1); + EXPECT_EQ(ExpressionStep::MakeMutableListAppendStep().StackDelta(), -1); + + EXPECT_EQ(ExpressionStep::MakeCreateSmallListStep(SmallListStepInfo{0, 4}) + .StackDelta(), + -3); + EXPECT_EQ( + ExpressionStep::MakeCreateListStep( + std::make_unique(5, absl::flat_hash_set{})) + .StackDelta(), + -4); + + EXPECT_EQ(ExpressionStep::MakeEagerFunctionStep( + std::make_unique( + std::vector{}, "fn", + /*num_args=*/3, /*receiver_style=*/false, /*expr_id=*/1)) + .StackDelta(), + -2); + EXPECT_EQ(ExpressionStep::MakeLazyFunctionStep( + std::make_unique( + std::vector{}, "fn", + /*num_args=*/2, /*receiver_style=*/false, /*expr_id=*/1)) + .StackDelta(), + -1); + + EXPECT_EQ(ExpressionStep::MakeGenericStep( + std::make_unique()) + .StackDelta(), + 1); + EXPECT_EQ(ExpressionStep::MakeGenericStep( + std::make_unique(), /*id=*/-1, + /*stack_delta=*/-2) + .StackDelta(), + -2); + EXPECT_EQ(ExpressionStep::MakeGenericStep( + std::make_unique(), /*id=*/-1, + /*stack_delta=*/std::numeric_limits::max()) + .StackDelta(), + std::nullopt); + EXPECT_EQ(ExpressionStep::MakeGenericStep( + std::make_unique(), /*id=*/-1, + /*stack_delta=*/ + static_cast(std::numeric_limits::min()) - 1) + .StackDelta(), + std::nullopt); +} + } // namespace google::api::expr::runtime diff --git a/eval/eval/function_step.h b/eval/eval/function_step.h index c5d4893da..4587da2c5 100644 --- a/eval/eval/function_step.h +++ b/eval/eval/function_step.h @@ -40,6 +40,9 @@ std::unique_ptr CreateFunctionStep( // Common base class for EagerFunctionStep and LazyFunctionStep. class FunctionStepBase { + public: + size_t num_arguments() const { return num_arguments_; } + private: friend class EagerFunctionStep; friend class LazyFunctionStep; diff --git a/extensions/select_optimization.cc b/extensions/select_optimization.cc index 019812438..7747eeafe 100644 --- a/extensions/select_optimization.cc +++ b/extensions/select_optimization.cc @@ -235,7 +235,7 @@ absl::StatusOr ListIndexFromQualifier(const AttributeQualifier& qual) { int64_t value = -1; switch (qual.kind()) { case Kind::kInt: - value = *qual.GetInt64Key(); + value = *qual.AsInt(); break; default: // TODO(uncreated-issue/51): type-checker will reject an unsigned literal, but @@ -255,11 +255,11 @@ absl::StatusOr MapKeyFromQualifier(const AttributeQualifier& qual, google::protobuf::Arena* absl_nonnull arena) { switch (qual.kind()) { case Kind::kInt: - return cel::IntValue(*qual.GetInt64Key()); + return cel::IntValue(*qual.AsInt()); case Kind::kUint: - return cel::UintValue(*qual.GetUint64Key()); + return cel::UintValue(*qual.AsUint()); case Kind::kBool: - return cel::BoolValue(*qual.GetBoolKey()); + return cel::BoolValue(*qual.AsBool()); case Kind::kString: // Safe as `qual` is valid for the lifetime of the expression/program. return StringValue::WrapUnsafe(*qual.GetStringKey()); @@ -841,8 +841,8 @@ absl::Status SelectOptimizer::OnPostVisit(PlannerContext& context, absl::c_move(operand_subplan, std::back_inserter(path)); path.push_back(ExpressionStep::MakeGenericStep( - std::make_unique(node.id(), std::move(impl)), - node.id())); + std::make_unique(node.id(), std::move(impl)), node.id(), + /*stack_delta=*/0)); return context.ReplaceSubplan(node, std::move(path)); }