Track iterator stack limit separately from comprehension slots.

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;

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<ProgramOptimizer> 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<FlatExpression> 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_

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",
     ],
 )

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);
 

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<cel::Value> FlatExpression::EvaluateWithCallback(

eval/eval/evaluator_core.h

@@ -450,24 +450,30 @@ 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),
         message_factory_(message_factory) {}
 
   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<ExecutionPathView> 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<google::protobuf::Arena> 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<ExecutionPathView> 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

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);

eval/eval/iterator_stack.h

@@ -20,7 +20,6 @@
 #include <vector>
 
 #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<Entry, 2> entries_;
+  std::vector<Entry> entries_;
   size_t max_size_;
 };
 

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_) {}