From 234f0b83bae708b64b1e70f3d5a9c2a2fd55282e Mon Sep 17 00:00:00 2001 From: Jonathan Tatum Date: Tue, 13 Oct 2026 16:48:03 -0700 Subject: [PATCH] Track iterator stack limit separately from comprehension slots. PiperOrigin-RevId: 998976105 --- eval/compiler/constant_folding.cc | 5 ++- eval/compiler/flat_expr_builder.cc | 52 ++---------------------------- eval/eval/BUILD | 1 - eval/eval/equality_steps_test.cc | 8 ++--- eval/eval/evaluator_core.cc | 8 ++--- eval/eval/evaluator_core.h | 21 ++++++------ eval/eval/evaluator_core_test.cc | 4 +-- eval/eval/iterator_stack.h | 3 +- eval/eval/lazy_init_step_test.cc | 3 +- 9 files changed, 28 insertions(+), 77 deletions(-) diff --git a/eval/compiler/constant_folding.cc b/eval/compiler/constant_folding.cc index 4a8d55c31..fb58d7d48 100644 --- a/eval/compiler/constant_folding.cc +++ b/eval/compiler/constant_folding.cc @@ -75,9 +75,8 @@ class ConstantFoldingExtension : public ProgramOptimizer { const TypeProvider& type_provider) : shared_arena_(std::move(shared_arena)), shared_message_factory_(std::move(shared_message_factory)), - state_(kDefaultStackLimit, kComprehensionSlotCount, - /*iterator_stack_size=*/0, type_provider, descriptor_pool, - message_factory, arena) {} + state_(kDefaultStackLimit, kComprehensionSlotCount, type_provider, + descriptor_pool, message_factory, arena) {} absl::Status OnPreVisit(google::api::expr::runtime::PlannerContext& context, const Expr& node) override; diff --git a/eval/compiler/flat_expr_builder.cc b/eval/compiler/flat_expr_builder.cc index 660e34182..f0920bf18 100644 --- a/eval/compiler/flat_expr_builder.cc +++ b/eval/compiler/flat_expr_builder.cc @@ -105,12 +105,7 @@ class FlatExprVisitor; // Helper for bookkeeping variables mapped to indexes. class IndexManager { public: - IndexManager() - : next_free_slot_(0), - max_slot_count_(0), - iter_stack_count_(0), - max_iter_stack_count_(0), - lazy_iter_stack_count_(0) {} + IndexManager() : next_free_slot_(0), max_slot_count_(0) {} size_t ReserveSlots(size_t n) { size_t result = next_free_slot_; @@ -126,30 +121,11 @@ class IndexManager { return next_free_slot_; } - void ReserveIterStack() { - ++iter_stack_count_; - max_iter_stack_count_ = std::max(max_iter_stack_count_, iter_stack_count_); - } - void ReleaseIterStack() { --iter_stack_count_; } - void AddLazyIterStack() { ++lazy_iter_stack_count_; } - size_t max_slot_count() const { return max_slot_count_; } - // Returns the maximum number of iterator stack slots required to evaluate the - // expression by tracking the maximum comprehension nesting in the main - // expression and dedicating a slot for any comprehensions evaluated lazily. - // This is not as tight as possible, but is a correct upper bound estimate and - // is tight for small expressions. - size_t IterStackRequirement() const { - return max_iter_stack_count_ + lazy_iter_stack_count_; - } - private: size_t next_free_slot_; size_t max_slot_count_; - size_t iter_stack_count_; - size_t max_iter_stack_count_; - size_t lazy_iter_stack_count_; }; // Helper for computing jump offsets. @@ -1169,34 +1145,21 @@ class FlatExprVisitor : public cel::AstVisitor { slot_count = 0; } } - // If this is in the scope of an optimized bind accu-init, account the slots // to the outermost bind-init scope. // // The init expression is effectively inlined at the first usage in the // critical path (which is unknown at plan time), so the used slots need to // be dedicated for the entire scope of that bind. - bool in_bind_init = false; for (ComprehensionStackRecord& record : comprehension_stack_) { if (record.in_accu_init && record.is_optimizable_bind) { record.slot_count += slot_count; slot_count = 0; - in_bind_init = true; break; } // If no bind init subexpression, account normally. } - const bool is_stack_releasable = - !in_bind_init && (!block_.has_value() || !block_->in); - if (!is_bind) { - if (is_stack_releasable) { - index_manager_.AddLazyIterStack(); - } else { - index_manager_.ReserveIterStack(); - } - } - comprehension_stack_.push_back( {&expr, &comprehension, iter_slot, iter2_slot, accu_slot, slot_count, /*subexpression=*/-1, @@ -1209,7 +1172,6 @@ class FlatExprVisitor : public cel::AstVisitor { options_.enable_comprehension_mutable_map, options_.short_circuiting), /*.is_optimizable_bind=*/is_bind, - /*.is_stack_releasable=*/is_stack_releasable, /*.iter_var_in_scope=*/false, /*.iter_var2_in_scope=*/false, /*.accu_var_in_scope=*/false, @@ -1236,9 +1198,6 @@ class FlatExprVisitor : public cel::AstVisitor { record.visitor->PostVisit(&expr); index_manager_.ReleaseSlots(record.slot_count); - if (record.is_stack_releasable) { - index_manager_.ReleaseIterStack(); - } comprehension_stack_.pop_back(); } @@ -1540,10 +1499,6 @@ class FlatExprVisitor : public cel::AstVisitor { size_t slot_count() const { return index_manager_.max_slot_count(); } - size_t iter_stack_count() const { - return index_manager_.IterStackRequirement(); - } - void AddOptimizer(std::unique_ptr optimizer) { program_optimizers_.push_back(std::move(optimizer)); } @@ -1575,7 +1530,6 @@ class FlatExprVisitor : public cel::AstVisitor { bool is_optimizable_list_append; bool is_optimizable_map_insert; bool is_optimizable_bind; - bool is_stack_releasable; bool iter_var_in_scope; bool iter_var2_in_scope; bool accu_var_in_scope; @@ -2344,8 +2298,8 @@ absl::StatusOr FlatExprBuilder::CreateExpressionImpl( FlattenExpressionTable(program_builder, execution_path); return FlatExpression(std::move(execution_path), std::move(subexpressions), - visitor.slot_count(), visitor.iter_stack_count(), - GetTypeProvider(), options_, std::move(arena)); + visitor.slot_count(), GetTypeProvider(), options_, + std::move(arena)); } const cel::TypeProvider& FlatExprBuilder::GetTypeProvider() const { return use_legacy_type_provider_ diff --git a/eval/eval/BUILD b/eval/eval/BUILD index 55b1e5e29..fa6c95a9b 100644 --- a/eval/eval/BUILD +++ b/eval/eval/BUILD @@ -929,7 +929,6 @@ cc_library( deps = [ "//common:value", "@com_google_absl//absl/base:nullability", - "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/log:absl_check", ], ) diff --git a/eval/eval/equality_steps_test.cc b/eval/eval/equality_steps_test.cc index 50970ca44..4ebfe10ef 100644 --- a/eval/eval/equality_steps_test.cc +++ b/eval/eval/equality_steps_test.cc @@ -81,7 +81,7 @@ TEST(IterativeTest, PartialAttrUnknown) { FlatExpressionEvaluatorState state( /*value_stack_size=*/5, - /*comprehension_slot_count=*/0, /*iterator_stack_size=*/0, type_provider, + /*comprehension_slot_count=*/0, type_provider, cel::internal::GetTestingDescriptorPool(), cel::internal::GetTestingMessageFactory(), &arena); @@ -113,7 +113,7 @@ TEST(IterativeTest, PartialAttrUnknownDisabled) { FlatExpressionEvaluatorState state( /*value_stack_size=*/5, - /*comprehension_slot_count=*/0, /*iterator_stack_size=*/0, type_provider, + /*comprehension_slot_count=*/0, type_provider, cel::internal::GetTestingDescriptorPool(), cel::internal::GetTestingMessageFactory(), &arena); @@ -183,7 +183,7 @@ TEST_P(EqualsTest, Iterative) { FlatExpressionEvaluatorState state( /*value_stack_size=*/5, - /*comprehension_slot_count=*/0, /*iterator_stack_size=*/0, type_provider, + /*comprehension_slot_count=*/0, type_provider, cel::internal::GetTestingDescriptorPool(), cel::internal::GetTestingMessageFactory(), &arena); @@ -328,7 +328,7 @@ TEST_P(InTest, Iterative) { FlatExpressionEvaluatorState state( /*value_stack_size=*/5, - /*comprehension_slot_count=*/0, /*iterator_stack_size=*/0, type_provider, + /*comprehension_slot_count=*/0, type_provider, cel::internal::GetTestingDescriptorPool(), cel::internal::GetTestingMessageFactory(), &arena); diff --git a/eval/eval/evaluator_core.cc b/eval/eval/evaluator_core.cc index 2fc6eda93..cbc47d136 100644 --- a/eval/eval/evaluator_core.cc +++ b/eval/eval/evaluator_core.cc @@ -222,16 +222,16 @@ FlatExpressionEvaluatorState FlatExpression::MakeEvaluatorState( 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); + 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); + type_provider_, descriptor_pool, + message_factory); } absl::StatusOr FlatExpression::EvaluateWithCallback( diff --git a/eval/eval/evaluator_core.h b/eval/eval/evaluator_core.h index b389c38a6..e415487c3 100644 --- a/eval/eval/evaluator_core.h +++ b/eval/eval/evaluator_core.h @@ -450,11 +450,14 @@ class FlatExpressionEvaluatorState { public: FlatExpressionEvaluatorState( size_t value_stack_size, size_t comprehension_slot_count, - size_t iterator_stack_size, const cel::TypeProvider& type_provider, + const cel::TypeProvider& type_provider, const google::protobuf::DescriptorPool* absl_nonnull descriptor_pool, google::protobuf::MessageFactory* absl_nullable message_factory = nullptr) : value_stack_(value_stack_size), - iterator_stack_(iterator_stack_size), + // We currently use comprehension_slot_count because it is less of an + // over estimate than value_stack_size. In future we should just + // calculate the correct capacity. + iterator_stack_(comprehension_slot_count), comprehension_slots_(comprehension_slot_count), type_provider_(type_provider), descriptor_pool_(descriptor_pool), @@ -462,12 +465,15 @@ class FlatExpressionEvaluatorState { FlatExpressionEvaluatorState( size_t value_stack_size, size_t comprehension_slot_count, - size_t iterator_stack_size, const cel::TypeProvider& type_provider, + const cel::TypeProvider& type_provider, const google::protobuf::DescriptorPool* absl_nonnull descriptor_pool, google::protobuf::MessageFactory* absl_nonnull message_factory, google::protobuf::Arena* absl_nonnull arena) : value_stack_(value_stack_size), - iterator_stack_(iterator_stack_size), + // We currently use comprehension_slot_count because it is less of an + // over estimate than value_stack_size. In future we should just + // calculate the correct capacity. + iterator_stack_(comprehension_slot_count), comprehension_slots_(comprehension_slot_count), type_provider_(type_provider), descriptor_pool_(descriptor_pool), @@ -873,21 +879,19 @@ class FlatExpression { : path_(std::move(path)), subexpressions_({path_}), comprehension_slots_size_(comprehension_slots_size), - iterator_stack_size_(comprehension_slots_size), type_provider_(type_provider), options_(options), arena_(std::move(arena)) {} FlatExpression(ExecutionPath path, std::vector subexpressions, - size_t comprehension_slots_size, size_t iterator_stack_size, + size_t comprehension_slots_size, const cel::TypeProvider& type_provider, const cel::RuntimeOptions& options, absl_nullable std::shared_ptr arena = nullptr) : path_(std::move(path)), subexpressions_(std::move(subexpressions)), comprehension_slots_size_(comprehension_slots_size), - iterator_stack_size_(iterator_stack_size), type_provider_(type_provider), options_(options), arena_(std::move(arena)) {} @@ -930,15 +934,12 @@ class FlatExpression { size_t comprehension_slots_size() const { return comprehension_slots_size_; } - size_t iterator_stack_size() const { return iterator_stack_size_; } - const cel::TypeProvider& type_provider() const { return type_provider_; } private: ExecutionPath path_; std::vector subexpressions_; size_t comprehension_slots_size_; - size_t iterator_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..c67dc8b58 100644 --- a/eval/eval/evaluator_core_test.cc +++ b/eval/eval/evaluator_core_test.cc @@ -84,7 +84,7 @@ TEST(EvaluatorCoreTest, ExecutionFrameNext) { cel::Activation activation; FlatExpressionEvaluatorState state( path.size(), - /*comprehension_slot_count=*/0, /*iterator_stack_size=*/0, type_provider, + /*comprehension_slots_size=*/0, type_provider, cel::internal::GetTestingDescriptorPool(), cel::internal::GetTestingMessageFactory(), &arena); ExecutionFrame frame(path, activation, options, state); @@ -111,7 +111,7 @@ TEST(EvaluatorCoreTest, JumpOutOfRangeAborts) { cel::Activation activation; FlatExpressionEvaluatorState state( path.size(), - /*comprehension_slot_count=*/0, /*iterator_stack_size=*/0, type_provider, + /*comprehension_slots_size=*/0, type_provider, cel::internal::GetTestingDescriptorPool(), cel::internal::GetTestingMessageFactory(), &arena); ExecutionFrame frame(path, activation, options, state); diff --git a/eval/eval/iterator_stack.h b/eval/eval/iterator_stack.h index ccb377dd3..0b0de14aa 100644 --- a/eval/eval/iterator_stack.h +++ b/eval/eval/iterator_stack.h @@ -20,7 +20,6 @@ #include #include "absl/base/nullability.h" -#include "absl/container/inlined_vector.h" #include "absl/log/absl_check.h" #include "common/value.h" @@ -93,7 +92,7 @@ class IteratorStack final { } private: - absl::InlinedVector entries_; + std::vector entries_; size_t max_size_; }; diff --git a/eval/eval/lazy_init_step_test.cc b/eval/eval/lazy_init_step_test.cc index d7211e597..3b5282d21 100644 --- a/eval/eval/lazy_init_step_test.cc +++ b/eval/eval/lazy_init_step_test.cc @@ -46,8 +46,7 @@ class LazyInitStepTest : public testing::Test { public: LazyInitStepTest() : type_provider_(cel::internal::GetTestingDescriptorPool()), - evaluator_state_(kValueStack, kComprehensionSlotCount, - /*iterator_stack_size=*/0, type_provider_, + evaluator_state_(kValueStack, kComprehensionSlotCount, type_provider_, cel::internal::GetTestingDescriptorPool(), cel::internal::GetTestingMessageFactory(), &arena_) {}