Update stack size estimate to use a tighter upper bound.

For many expressions, the previous method preallocated stacks that were much larger than can actually be used during evaluation.

PiperOrigin-RevId: 592953564

eval/compiler/flat_expr_builder.cc

@@ -16,6 +16,7 @@
 
 #include "eval/compiler/flat_expr_builder.h"
 
+#include <algorithm>
 #include <cstddef>
 #include <cstdint>
 #include <deque>
@@ -118,6 +119,28 @@ class IndexManager {
   size_t max_slot_count_;
 };
 
+// Estimate the maximum stack requirements for a subexpression.
+//
+// Maintain a running sum of the program Steps reported estimated stack change
+// upper bound.
+//
+// Assumptions: Jump steps and Call steps complicate this, but this should still
+// be valid for supported CEL expressions.
+// - While abstractly a Jump may invalidate this algorithm (e.g. deterministic
+//   jump over a stack shrinking step), the expression builder should not create
+//   programs where this is possible.
+// - Call steps are currently limited to lazy initializing cel.bind expressions,
+//   and report the estimated stack requirements for the called subexpression.
+size_t EstimateStackRequirement(const ExecutionPath& execution_path) {
+  int stack_estimate = 0;
+  int stack_limit = 1;
+  for (const auto& step : execution_path) {
+    stack_estimate += step->stack_delta();
+    stack_limit = std::max(stack_limit, stack_estimate);
+  }
+  return static_cast<size_t>(stack_limit);
+}
+
 // A convenience wrapper for offset-calculating logic.
 class Jump {
  public:
@@ -404,6 +427,7 @@ class FlatExprVisitor : public cel::ast_internal::AstVisitor {
   struct SlotLookupResult {
     int slot;
     int subexpression;
+    int stack_requirement;
   };
 
   // Helper to lookup a variable mapped to a slot.
@@ -428,10 +452,12 @@ class FlatExprVisitor : public cel::ast_internal::AstVisitor {
             record.comprehension->accu_var() == path) {
           int slot = record.accu_slot;
           int subexpression = -1;
+          int stack_requirement = -1;
           if (record.should_lazy_eval) {
             subexpression = record.subexpression;
+            stack_requirement = record.subexpression_stack_requirement;
           }
-          return {slot, subexpression};
+          return {slot, subexpression, stack_requirement};
         }
       }
     }
@@ -490,9 +516,9 @@ class FlatExprVisitor : public cel::ast_internal::AstVisitor {
     SlotLookupResult slot = LookupSlot(path);
 
     if (slot.subexpression >= 0) {
-      AddStep(
-          CreateCheckLazyInitStep(slot.slot, slot.subexpression, expr->id()));
-      AddStep(CreateAssignSlotStep(slot.slot));
+      AddStep(CreateCheckLazyInitStep(slot.slot, slot.subexpression,
+                                      slot.stack_requirement, expr->id()));
+      AddStep(CreateAssignSlotStep(slot.slot, (1 - slot.stack_requirement)));
       return;
     } else if (slot.slot >= 0) {
       AddStep(CreateIdentStepForSlot(*ident_expr, slot.slot, expr->id()));
@@ -728,6 +754,7 @@ class FlatExprVisitor : public cel::ast_internal::AstVisitor {
     comprehension_stack_.push_back(
         {expr, comprehension, iter_slot, accu_slot,
          /*subexpression=*/-1,
+         /*subexpression_stack_requirement=*/-1,
          IsOptimizableListAppend(comprehension,
                                  options_.enable_comprehension_list_append),
          is_bind,
@@ -980,6 +1007,7 @@ class FlatExprVisitor : public cel::ast_internal::AstVisitor {
     size_t accu_slot;
     // -1 indicates this shouldn't be used.
     int subexpression;
+    int subexpression_stack_requirement;
     bool is_optimizable_list_append;
     bool is_optimizable_bind;
     bool iter_var_in_scope;
@@ -1012,6 +1040,8 @@ class FlatExprVisitor : public cel::ast_internal::AstVisitor {
     // off by one since mainline expression is handled separately.
     size_t id = expression_table_.size();
     record.subexpression = id;
+    record.subexpression_stack_requirement =
+        static_cast<int>(EstimateStackRequirement(expression_table_.back()));
 
     record.visitor->MarkAccuInitExtracted();
 
@@ -1340,11 +1370,13 @@ absl::StatusOr<FlatExpression> FlatExprBuilder::CreateExpressionImpl(
     (*issues) = issue_collector.ExtractIssues();
   }
 
+  size_t stack_limit = EstimateStackRequirement(execution_path);
+
   std::vector<ExecutionPathView> subexpressions =
       FlattenExpressionTable(expression_table, execution_path);
 
   return FlatExpression(std::move(execution_path), std::move(subexpressions),
-                        visitor.slot_count(),
+                        stack_limit, visitor.slot_count(),
                         type_registry_.GetComposedTypeProvider(), options_);
 }
 

eval/eval/BUILD

@@ -438,6 +438,7 @@ cc_test(
     deps = [
         ":cel_expression_flat_impl",
         ":evaluator_core",
+        ":expression_step_base",
         "//base:data",
         "//eval/compiler:cel_expression_builder_flat_impl",
         "//eval/internal:interop",

eval/eval/compiler_constant_step.h

@@ -30,7 +30,8 @@ class CompilerConstantStep : public ExpressionStepBase {
  public:
   CompilerConstantStep(cel::Handle<cel::Value> value, int64_t expr_id,
                        bool comes_from_ast)
-      : ExpressionStepBase(expr_id, comes_from_ast), value_(std::move(value)) {}
+      : ExpressionStepBase(expr_id, comes_from_ast, /*stack_delta=*/1),
+        value_(std::move(value)) {}
 
   absl::Status Evaluate(ExecutionFrame* frame) const override;
 

eval/eval/comprehension_step.cc

@@ -44,7 +44,8 @@ class ComprehensionFinish : public ExpressionStepBase {
 };
 
 ComprehensionFinish::ComprehensionFinish(size_t accu_slot, int64_t expr_id)
-    : ExpressionStepBase(expr_id), accu_slot_(accu_slot) {}
+    : ExpressionStepBase(expr_id, true, /*stack_delta=*/-2),
+      accu_slot_(accu_slot) {}
 
 // Stack changes of ComprehensionFinish.
 //
@@ -73,7 +74,7 @@ absl::Status ComprehensionFinish::Evaluate(ExecutionFrame* frame) const {
 class ComprehensionInitStep : public ExpressionStepBase {
  public:
   explicit ComprehensionInitStep(int64_t expr_id)
-      : ExpressionStepBase(expr_id, false) {}
+      : ExpressionStepBase(expr_id, false, /*stack_delta=*/1) {}
   absl::Status Evaluate(ExecutionFrame* frame) const override;
 
  private:
@@ -148,7 +149,7 @@ absl::Status ComprehensionInitStep::Evaluate(ExecutionFrame* frame) const {
 
 ComprehensionNextStep::ComprehensionNextStep(size_t iter_slot, size_t accu_slot,
                                              int64_t expr_id)
-    : ExpressionStepBase(expr_id, false),
+    : ExpressionStepBase(expr_id, false, /*stack_delta=*/-1),
       iter_slot_(iter_slot),
       accu_slot_(accu_slot) {}
 
@@ -252,7 +253,7 @@ absl::Status ComprehensionNextStep::Evaluate(ExecutionFrame* frame) const {
 ComprehensionCondStep::ComprehensionCondStep(size_t iter_slot, size_t accu_slot,
                                              bool shortcircuiting,
                                              int64_t expr_id)
-    : ExpressionStepBase(expr_id, false),
+    : ExpressionStepBase(expr_id, false, /*stack_delta=*/-1),
       iter_slot_(iter_slot),
       accu_slot_(accu_slot),
       shortcircuiting_(shortcircuiting) {}

eval/eval/comprehension_step_test.cc

@@ -49,9 +49,9 @@ class ListKeysStepTest : public testing::Test {
       options.unknown_processing =
           cel::UnknownProcessingOptions::kAttributeAndFunction;
     }
-    return std::make_unique<CelExpressionFlatImpl>(
-        FlatExpression(std::move(path), /*comprehension_slot_count=*/0,
-                       TypeProvider::Builtin(), options));
+    return std::make_unique<CelExpressionFlatImpl>(FlatExpression(
+        std::move(path), /*value_stack_size=*/2, /*comprehension_slot_count=*/0,
+        TypeProvider::Builtin(), options));
   }
 
  private:

eval/eval/const_value_step_test.cc

@@ -44,9 +44,11 @@ absl::StatusOr<CelValue> RunConstantExpression(
   google::api::expr::runtime::ExecutionPath path;
   path.push_back(std::move(step));
 
-  CelExpressionFlatImpl impl(
-      FlatExpression(std::move(path), /*comprehension_slot_count=*/0,
-                     TypeProvider::Builtin(), cel::RuntimeOptions{}));
+  CelExpressionFlatImpl impl(FlatExpression(std::move(path),
+                                            /*value_stack_size=*/1,
+                                            /*comprehension_slot_count=*/0,
+                                            TypeProvider::Builtin(),
+                                            cel::RuntimeOptions{}));
 
   google::api::expr::runtime::Activation activation;
 

eval/eval/container_access_step.cc

@@ -48,7 +48,8 @@ inline constexpr int kNumContainerAccessArguments = 2;
 // message.
 class ContainerAccessStep : public ExpressionStepBase {
  public:
-  explicit ContainerAccessStep(int64_t expr_id) : ExpressionStepBase(expr_id) {}
+  explicit ContainerAccessStep(int64_t expr_id)
+      : ExpressionStepBase(expr_id, true, /*stack_delta=*/-1) {}
 
   absl::Status Evaluate(ExecutionFrame* frame) const override;
 

eval/eval/container_access_step_test.cc

@@ -72,9 +72,10 @@ CelValue EvaluateAttributeHelper(
   cel::RuntimeOptions options;
   options.unknown_processing = cel::UnknownProcessingOptions::kAttributeOnly;
   options.enable_heterogeneous_equality = false;
-  CelExpressionFlatImpl cel_expr(
-      FlatExpression(std::move(path), /*comprehension_slot_count=*/0,
-                     TypeProvider::Builtin(), options));
+  CelExpressionFlatImpl cel_expr(FlatExpression(
+      std::move(path),
+      /*value_stack_size=*/2,
+      /*comprehension_slot_count=*/0, TypeProvider::Builtin(), options));
   Activation activation;
 
   activation.InsertValue("container", container);

eval/eval/create_list_step.cc

@@ -31,7 +31,7 @@ using ::cel::runtime_internal::MutableListValue;
 class CreateListStep : public ExpressionStepBase {
  public:
   CreateListStep(int64_t expr_id, int list_size, bool immutable)
-      : ExpressionStepBase(expr_id),
+      : ExpressionStepBase(expr_id, true, /*stack_delta=*/1 - list_size),
         list_size_(list_size),
         immutable_(immutable) {}
 

eval/eval/create_list_step_test.cc

@@ -54,11 +54,11 @@ absl::StatusOr<CelValue> RunExpression(const std::vector<int64_t>& values,
   if (enable_unknowns) {
     options.unknown_processing = cel::UnknownProcessingOptions::kAttributeOnly;
   }
-  CelExpressionFlatImpl cel_expr(
-
-      FlatExpression(std::move(path),
-                     /*comprehension_slot_count=*/0, TypeProvider::Builtin(),
-                     options));
+  size_t stack_size = path.size();
+  CelExpressionFlatImpl cel_expr(FlatExpression(
+      std::move(path),
+      /*value_stack_size=*/stack_size,
+      /*comprehension_slot_count=*/0, TypeProvider::Builtin(), options));
   Activation activation;
 
   return cel_expr.Evaluate(activation, arena);
@@ -95,9 +95,10 @@ absl::StatusOr<CelValue> RunExpressionWithCelValues(
     options.unknown_processing = cel::UnknownProcessingOptions::kAttributeOnly;
   }
 
-  CelExpressionFlatImpl cel_expr(
-      FlatExpression(std::move(path), /*comprehension_slot_count=*/0,
-                     TypeProvider::Builtin(), options));
+  CelExpressionFlatImpl cel_expr(FlatExpression(
+      std::move(path),
+      /*value_stack_size=*/values.size(),
+      /*comprehension_slot_count=*/0, TypeProvider::Builtin(), options));
 
   return cel_expr.Evaluate(activation, arena);
 }
@@ -118,9 +119,11 @@ TEST(CreateListStepTest, TestCreateListStackUnderflow) {
                        CreateCreateListStep(create_list, dummy_expr.id()));
   path.push_back(std::move(step0));
 
-  CelExpressionFlatImpl cel_expr(
-      FlatExpression(std::move(path), /*comprehension_slot_count=*/0,
-                     TypeProvider::Builtin(), cel::RuntimeOptions{}));
+  CelExpressionFlatImpl cel_expr(FlatExpression(std::move(path),
+                                                /*value_stack_size=*/1,
+                                                /*comprehension_slot_count=*/0,
+                                                TypeProvider::Builtin(),
+                                                cel::RuntimeOptions{}));
   Activation activation;
 
   google::protobuf::Arena arena;

eval/eval/create_struct_step.cc

@@ -1,5 +1,6 @@
 #include "eval/eval/create_struct_step.h"
 
+#include <cstddef>
 #include <cstdint>
 #include <memory>
 #include <string>
@@ -815,7 +816,7 @@ class CreateStructStepForStruct final : public ExpressionStepBase {
  public:
   CreateStructStepForStruct(int64_t expr_id, Handle<StructType> type,
                             std::vector<StructFieldId> entries)
-      : ExpressionStepBase(expr_id),
+      : ExpressionStepBase(expr_id, true, 1 - entries.size()),
         type_(std::move(type)),
         entries_(std::move(entries)) {}
 
@@ -853,7 +854,8 @@ class CreateStructStepForWellKnownType final : public ExpressionStepBase {
 class CreateStructStepForMap final : public ExpressionStepBase {
  public:
   CreateStructStepForMap(int64_t expr_id, size_t entry_count)
-      : ExpressionStepBase(expr_id), entry_count_(entry_count) {}
+      : ExpressionStepBase(expr_id, true, 1 - (2 * entry_count)),
+        entry_count_(entry_count) {}
 
   absl::Status Evaluate(ExecutionFrame* frame) const override;
 

eval/eval/create_struct_step_test.cc

@@ -94,9 +94,10 @@ absl::StatusOr<CelValue> RunExpression(absl::string_view field,
   if (enable_unknowns) {
     options.unknown_processing = cel::UnknownProcessingOptions::kAttributeOnly;
   }
-  CelExpressionFlatImpl cel_expr(
-      FlatExpression(std::move(path), /*comprehension_slot_count=*/0,
-                     TypeProvider::Builtin(), options));
+  CelExpressionFlatImpl cel_expr(FlatExpression(
+      std::move(path),
+      /*value_stack_size=*/1,
+      /*comprehension_slot_count=*/0, TypeProvider::Builtin(), options));
   Activation activation;
   activation.InsertValue("message", value);
 
@@ -187,9 +188,10 @@ absl::StatusOr<CelValue> RunCreateMapExpression(
     options.unknown_processing = cel::UnknownProcessingOptions::kAttributeOnly;
   }
 
-  CelExpressionFlatImpl cel_expr(
-      FlatExpression(std::move(path), /*comprehension_slot_count=*/0,
-                     TypeProvider::Builtin(), options));
+  CelExpressionFlatImpl cel_expr(FlatExpression(
+      std::move(path),
+      /*value_stack_size=*/2 * values.size(),
+      /*comprehension_slot_count=*/0, TypeProvider::Builtin(), options));
   return cel_expr.Evaluate(activation, arena);
 }
 
@@ -226,9 +228,10 @@ TEST_P(CreateCreateStructStepTest, TestEmptyMessageCreation) {
   if (GetParam()) {
     options.unknown_processing = cel::UnknownProcessingOptions::kAttributeOnly;
   }
-  CelExpressionFlatImpl cel_expr(
-      FlatExpression(std::move(path), /*comprehension_slot_count=*/0,
-                     TypeProvider::Builtin(), options));
+  CelExpressionFlatImpl cel_expr(FlatExpression(
+      std::move(path),
+      /*value_stack_size=*/1,
+      /*comprehension_slot_count=*/0, TypeProvider::Builtin(), options));
   Activation activation;
 
   ASSERT_OK_AND_ASSIGN(CelValue result, cel_expr.Evaluate(activation, &arena));

eval/eval/evaluator_core.cc

@@ -112,14 +112,14 @@ absl::StatusOr<cel::Handle<cel::Value>> ExecutionFrame::Evaluate(
 
 FlatExpressionEvaluatorState FlatExpression::MakeEvaluatorState(
     cel::MemoryManagerRef manager) const {
-  return FlatExpressionEvaluatorState(path_.size(), comprehension_slots_size_,
-                                      type_provider_, manager);
+  return FlatExpressionEvaluatorState(
+      value_stack_size_, comprehension_slots_size_, type_provider_, manager);
 }
 
 FlatExpressionEvaluatorState FlatExpression::MakeEvaluatorState(
     cel::ValueFactory& value_factory) const {
-  return FlatExpressionEvaluatorState(path_.size(), comprehension_slots_size_,
-                                      value_factory);
+  return FlatExpressionEvaluatorState(value_stack_size_,
+                                      comprehension_slots_size_, value_factory);
 }
 
 absl::StatusOr<cel::Handle<cel::Value>> FlatExpression::EvaluateWithCallback(