Track iterator stack limit separately from comprehension slots.

PiperOrigin-RevId: 1000446433

Author: jnthntatum

eval/compiler/BUILD

@@ -204,6 +204,11 @@ cc_test(
         ":cel_expression_builder_flat_impl",
         ":comprehension_vulnerability_check",
         ":flat_expr_builder",
+        "//base:ast",
+        "//common:source",
+        "//common:value",
+        "//common:value_testing",
+        "//eval/eval:evaluator_core",
         "//eval/public:activation",
         "//eval/public:builtin_func_registrar",
         "//eval/public:cel_attribute",
@@ -213,10 +218,15 @@ cc_test(
         "//eval/public/containers:container_backed_list_impl",
         "//eval/public/testing:matchers",
         "//eval/testutil:test_message_cc_proto",
+        "//extensions:bindings_ext",
         "//internal:testing",
         "//parser",
+        "//parser:parser_interface",
+        "//runtime:activation",
         "//runtime:runtime_options",
+        "//runtime/internal:runtime_env",
         "//runtime/internal:runtime_env_testing",
+        "//testutil:test_macros",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
         "@com_google_cel_spec//proto/cel/expr:syntax_cc_proto",

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;

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.
@@ -1145,21 +1169,36 @@ 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.
     }
 
+    // Iterator frames are releasable if not in a bind (no iterator push), and
+    // not in a lazy subexpression (which may evaluate in different implicit
+    // scopes).
+    const bool is_stack_releasable =
+        !is_bind && !in_bind_init && (!block_.has_value() || !block_->in);
+
+    if (is_stack_releasable) {
+      index_manager_.ReserveIterStack();
+    } else if (!is_bind) {
+      index_manager_.AddLazyIterStack();
+    }
+
     comprehension_stack_.push_back(
         {&expr, &comprehension, iter_slot, iter2_slot, accu_slot, slot_count,
          /*subexpression=*/-1,
@@ -1172,6 +1211,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,
@@ -1198,6 +1238,10 @@ class FlatExprVisitor : public cel::AstVisitor {
     record.visitor->PostVisit(&expr);
 
     index_manager_.ReleaseSlots(record.slot_count);
+    if (record.is_stack_releasable) {
+      ABSL_DCHECK(!record.is_optimizable_bind);
+      index_manager_.ReleaseIterStack();
+    }
     comprehension_stack_.pop_back();
   }
 
@@ -1499,6 +1543,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<ProgramOptimizer> optimizer) {
     program_optimizers_.push_back(std::move(optimizer));
   }
@@ -1530,6 +1578,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;
@@ -2298,8 +2347,8 @@ absl::StatusOr<FlatExpression> 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_

eval/compiler/flat_expr_builder_comprehensions_test.cc

@@ -14,15 +14,22 @@
  * limitations under the License.
  */
 
+#include <memory>
+#include <utility>
 #include <vector>
 
 #include "cel/expr/syntax.pb.h"
 #include "google/protobuf/field_mask.pb.h"
 #include "absl/status/status.h"
 #include "absl/strings/string_view.h"
+#include "base/ast.h"
+#include "common/source.h"
+#include "common/value.h"
+#include "common/value_testing.h"
 #include "eval/compiler/cel_expression_builder_flat_impl.h"
 #include "eval/compiler/comprehension_vulnerability_check.h"
 #include "eval/compiler/flat_expr_builder.h"
+#include "eval/eval/evaluator_core.h"
 #include "eval/public/activation.h"
 #include "eval/public/builtin_func_registrar.h"
 #include "eval/public/cel_attribute.h"
@@ -32,10 +39,15 @@
 #include "eval/public/containers/container_backed_list_impl.h"
 #include "eval/public/testing/matchers.h"
 #include "eval/testutil/test_message.pb.h"
+#include "extensions/bindings_ext.h"
 #include "internal/testing.h"
 #include "parser/parser.h"
+#include "parser/parser_interface.h"
+#include "runtime/activation.h"
+#include "runtime/internal/runtime_env.h"
 #include "runtime/internal/runtime_env_testing.h"
 #include "runtime/runtime_options.h"
+#include "testutil/test_macros.h"
 #include "google/protobuf/arena.h"
 #include "google/protobuf/text_format.h"
 
@@ -623,6 +635,119 @@ TEST_F(CelExpressionBuilderFlatImplComprehensionsTest,
           HasSubstr("Unexpected iter_var access in trivial comprehension")));
 }
 
+TEST_F(CelExpressionBuilderFlatImplComprehensionsTest,
+       BindInsideComprehensionWithComprehensionInInitIterStackCount) {
+  std::unique_ptr<cel::ParserBuilder> parser_builder = cel::NewParserBuilder();
+  ASSERT_THAT(
+      parser_builder->AddLibrary(cel::extensions::BindingsParserLibrary()),
+      IsOk());
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<cel::Parser> parser,
+                       parser_builder->Build());
+
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<cel::Source> source,
+      parser->PrepareSource(
+          "[[1, 2], [3, 4]].map(l, cel.bind(filtered, l.filter(x, x > 1), "
+          "size(filtered) > 0 ? filtered[0] : 0)) == [2, 3]"));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<cel::Ast> ast, parser->Parse(*source));
+
+  cel::RuntimeOptions options = GetRuntimeOptions();
+  std::shared_ptr<cel::runtime_internal::RuntimeEnv> env =
+      NewTestingRuntimeEnv();
+  CelExpressionBuilderFlatImpl builder(env, options);
+  ASSERT_THAT(RegisterBuiltinFunctions(builder.GetRegistry()), IsOk());
+
+  ASSERT_OK_AND_ASSIGN(
+      FlatExpression flat_expr,
+      builder.flat_expr_builder().CreateExpressionImpl(std::move(ast),
+                                                       /*issues=*/nullptr));
+  EXPECT_EQ(flat_expr.iterator_stack_size(), 2);
+
+  cel::Activation activation;
+  google::protobuf::Arena arena;
+  FlatExpressionEvaluatorState state = flat_expr.MakeEvaluatorState(
+      env->descriptor_pool.get(), env->MutableMessageFactory(), &arena);
+  ASSERT_OK_AND_ASSIGN(
+      cel::Value result,
+      flat_expr.EvaluateWithCallback(activation, /*embedder_context=*/nullptr,
+                                     EvaluationListener(), state));
+  EXPECT_THAT(result, cel::test::BoolValueIs(true));
+}
+
+TEST_F(CelExpressionBuilderFlatImplComprehensionsTest,
+       BlockWithComprehensionInBindingAndBodyIterStackCount) {
+  std::unique_ptr<cel::ParserBuilder> parser_builder = cel::NewParserBuilder();
+  ASSERT_THAT(parser_builder->AddLibrary(cel::test::TestMacrosParserLibrary()),
+              IsOk());
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<cel::Parser> parser,
+                       parser_builder->Build());
+
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<cel::Source> source,
+      parser->PrepareSource(
+          "cel.block([[1, 2, 3].filter(x, x > 1), [1, 2, 3].filter(x, x < 3)], "
+          "[1, 2].map(y, y in cel.index(0) || y in cel.index(1))) == [true, "
+          "true]"));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<cel::Ast> ast, parser->Parse(*source));
+
+  cel::RuntimeOptions options = GetRuntimeOptions();
+  std::shared_ptr<cel::runtime_internal::RuntimeEnv> env =
+      NewTestingRuntimeEnv();
+  CelExpressionBuilderFlatImpl builder(env, options);
+  ASSERT_THAT(RegisterBuiltinFunctions(builder.GetRegistry()), IsOk());
+
+  ASSERT_OK_AND_ASSIGN(
+      FlatExpression flat_expr,
+      builder.flat_expr_builder().CreateExpressionImpl(std::move(ast),
+                                                       /*issues=*/nullptr));
+  EXPECT_EQ(flat_expr.iterator_stack_size(), 3);
+
+  cel::Activation activation;
+  google::protobuf::Arena arena;
+  FlatExpressionEvaluatorState state = flat_expr.MakeEvaluatorState(
+      env->descriptor_pool.get(), env->MutableMessageFactory(), &arena);
+  ASSERT_OK_AND_ASSIGN(
+      cel::Value result,
+      flat_expr.EvaluateWithCallback(activation, /*embedder_context=*/nullptr,
+                                     EvaluationListener(), state));
+  EXPECT_THAT(result, cel::test::BoolValueIs(true));
+}
+
+TEST_F(CelExpressionBuilderFlatImplComprehensionsTest,
+       NonNestedComprehensionsIterStackCount) {
+  std::unique_ptr<cel::ParserBuilder> parser_builder = cel::NewParserBuilder();
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<cel::Parser> parser,
+                       parser_builder->Build());
+
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<cel::Source> source,
+      parser->PrepareSource(
+          "[1, 2].map(x, x + 1) + [3, 4].filter(y, y > 3) == [2, 3, 4]"));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<cel::Ast> ast, parser->Parse(*source));
+
+  cel::RuntimeOptions options = GetRuntimeOptions();
+  std::shared_ptr<cel::runtime_internal::RuntimeEnv> env =
+      NewTestingRuntimeEnv();
+  CelExpressionBuilderFlatImpl builder(env, options);
+  ASSERT_THAT(RegisterBuiltinFunctions(builder.GetRegistry()), IsOk());
+
+  ASSERT_OK_AND_ASSIGN(
+      FlatExpression flat_expr,
+      builder.flat_expr_builder().CreateExpressionImpl(std::move(ast),
+                                                       /*issues=*/nullptr));
+  EXPECT_EQ(flat_expr.iterator_stack_size(), 1);
+
+  cel::Activation activation;
+  google::protobuf::Arena arena;
+  FlatExpressionEvaluatorState state = flat_expr.MakeEvaluatorState(
+      env->descriptor_pool.get(), env->MutableMessageFactory(), &arena);
+  ASSERT_OK_AND_ASSIGN(
+      cel::Value result,
+      flat_expr.EvaluateWithCallback(activation, /*embedder_context=*/nullptr,
+                                     EvaluationListener(), state));
+  EXPECT_THAT(result, cel::test::BoolValueIs(true));
+}
+
 }  // namespace
 
 }  // namespace google::api::expr::runtime

eval/eval/BUILD

@@ -927,6 +927,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",
     ],
 )

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

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