diff --git a/eval/compiler/constant_folding.cc b/eval/compiler/constant_folding.cc index fb58d7d48..4a8d55c31 100644 --- a/eval/compiler/constant_folding.cc +++ b/eval/compiler/constant_folding.cc @@ -75,8 +75,9 @@ 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, type_provider, - descriptor_pool, message_factory, arena) {} + state_(kDefaultStackLimit, kComprehensionSlotCount, + /*iterator_stack_size=*/0, 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 c764b22eb..fa0d556b6 100644 --- a/eval/compiler/flat_expr_builder.cc +++ b/eval/compiler/flat_expr_builder.cc @@ -105,7 +105,12 @@ class FlatExprVisitor; // Helper for bookkeeping variables mapped to indexes. class IndexManager { public: - IndexManager() : next_free_slot_(0), max_slot_count_(0) {} + IndexManager() + : next_free_slot_(0), + max_slot_count_(0), + iter_stack_count_(0), + max_iter_stack_count_(0), + lazy_iter_stack_count_(0) {} size_t ReserveSlots(size_t n) { size_t result = next_free_slot_; @@ -121,11 +126,30 @@ 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. @@ -1132,21 +1156,34 @@ 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, @@ -1159,6 +1196,7 @@ 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, @@ -1185,6 +1223,9 @@ 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(); } @@ -1476,6 +1517,10 @@ 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)); } @@ -1507,6 +1552,7 @@ 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; @@ -2275,8 +2321,8 @@ absl::StatusOr FlatExprBuilder::CreateExpressionImpl( FlattenExpressionTable(program_builder, execution_path); return FlatExpression(std::move(execution_path), std::move(subexpressions), - visitor.slot_count(), GetTypeProvider(), options_, - std::move(arena)); + visitor.slot_count(), visitor.iter_stack_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 08dacb9b2..f505954bf 100644 --- a/eval/eval/BUILD +++ b/eval/eval/BUILD @@ -949,6 +949,7 @@ 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 4ebfe10ef..50970ca44 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, type_provider, + /*comprehension_slot_count=*/0, /*iterator_stack_size=*/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, type_provider, + /*comprehension_slot_count=*/0, /*iterator_stack_size=*/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, type_provider, + /*comprehension_slot_count=*/0, /*iterator_stack_size=*/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, type_provider, + /*comprehension_slot_count=*/0, /*iterator_stack_size=*/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 cbc47d136..2fc6eda93 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_, - type_provider_, descriptor_pool, - message_factory, arena); + 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_, - type_provider_, descriptor_pool, - message_factory); + iterator_stack_size_, 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 d23dfd045..07d985804 100644 --- a/eval/eval/evaluator_core.h +++ b/eval/eval/evaluator_core.h @@ -430,14 +430,11 @@ class FlatExpressionEvaluatorState { public: FlatExpressionEvaluatorState( size_t value_stack_size, size_t comprehension_slot_count, - const cel::TypeProvider& type_provider, + size_t iterator_stack_size, 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), - // 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), + iterator_stack_(iterator_stack_size), comprehension_slots_(comprehension_slot_count), type_provider_(type_provider), descriptor_pool_(descriptor_pool), @@ -445,15 +442,12 @@ class FlatExpressionEvaluatorState { FlatExpressionEvaluatorState( size_t value_stack_size, size_t comprehension_slot_count, - const cel::TypeProvider& type_provider, + size_t iterator_stack_size, 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), - // 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), + iterator_stack_(iterator_stack_size), comprehension_slots_(comprehension_slot_count), type_provider_(type_provider), descriptor_pool_(descriptor_pool), @@ -859,19 +853,21 @@ 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 comprehension_slots_size, size_t iterator_stack_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)) {} @@ -914,12 +910,15 @@ 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 c67dc8b58..1f43f3f35 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_slots_size=*/0, type_provider, + /*comprehension_slot_count=*/0, /*iterator_stack_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_slots_size=*/0, type_provider, + /*comprehension_slot_count=*/0, /*iterator_stack_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 0b0de14aa..ccb377dd3 100644 --- a/eval/eval/iterator_stack.h +++ b/eval/eval/iterator_stack.h @@ -20,6 +20,7 @@ #include #include "absl/base/nullability.h" +#include "absl/container/inlined_vector.h" #include "absl/log/absl_check.h" #include "common/value.h" @@ -92,7 +93,7 @@ class IteratorStack final { } private: - std::vector entries_; + absl::InlinedVector entries_; size_t max_size_; }; diff --git a/eval/eval/lazy_init_step_test.cc b/eval/eval/lazy_init_step_test.cc index 3b5282d21..d7211e597 100644 --- a/eval/eval/lazy_init_step_test.cc +++ b/eval/eval/lazy_init_step_test.cc @@ -46,7 +46,8 @@ class LazyInitStepTest : public testing::Test { public: LazyInitStepTest() : type_provider_(cel::internal::GetTestingDescriptorPool()), - evaluator_state_(kValueStack, kComprehensionSlotCount, type_provider_, + evaluator_state_(kValueStack, kComprehensionSlotCount, + /*iterator_stack_size=*/0, type_provider_, cel::internal::GetTestingDescriptorPool(), cel::internal::GetTestingMessageFactory(), &arena_) {}