Add tighter bound on value stack estimate.

PiperOrigin-RevId: 996788585

Author: jnthntatum

eval/compiler/BUILD

@@ -122,6 +122,7 @@ cc_library(
         "//runtime/internal:runtime_env",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base:nullability",
+        "@com_google_absl//absl/cleanup",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/container:node_hash_map",
@@ -146,6 +147,7 @@ cc_test(
         ":cel_expression_builder_flat_impl",
         ":constant_folding",
         ":flat_expr_builder",
+        ":flat_expr_builder_extensions",
         ":qualified_reference_resolver",
         "//base:builtins",
         "//checker:validation_result",
@@ -159,8 +161,8 @@ cc_test(
         "//compiler",
         "//compiler:compiler_factory",
         "//compiler:standard_library",
-        "//eval/compiler:flat_expr_builder_extensions",
         "//eval/eval:evaluator_core",
+        "//eval/eval:expression_step_logic",
         "//eval/public:activation",
         "//eval/public:ast_rewrite",
         "//eval/public:builtin_func_registrar",
@@ -181,6 +183,7 @@ cc_test(
         "//eval/public/structs:cel_proto_wrapper",
         "//eval/public/testing:matchers",
         "//eval/testutil:test_message_cc_proto",
+        "//extensions/protobuf:ast_converters",
         "//internal:proto_matchers",
         "//internal:status_macros",
         "//internal:testing",
@@ -202,7 +205,6 @@ cc_test(
         "@com_google_cel_spec//proto/cel/expr:checked_cc_proto",
         "@com_google_cel_spec//proto/cel/expr:syntax_cc_proto",
         "@com_google_cel_spec//proto/cel/expr/conformance/proto3:test_all_types_cc_proto",
-        "@com_google_protobuf//:field_mask_cc_proto",
         "@com_google_protobuf//:protobuf",
     ],
 )

eval/compiler/flat_expr_builder.cc

@@ -32,6 +32,7 @@
 
 #include "absl/algorithm/container.h"
 #include "absl/base/nullability.h"
+#include "absl/cleanup/cleanup.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
 #include "absl/container/node_hash_map.h"
@@ -971,13 +972,13 @@ class FlatExprVisitor : public cel::AstVisitor {
                                 *std::move(field_type), select_expr.test_only(),
                                 options_.enable_empty_wrapper_null_unboxing,
                                 enable_optional_types_),
-          expr.id());
+          expr.id(), /*stack_delta=*/0);
       return;
     }
     AddStep(CreateSelectStep(std::move(field), select_expr.test_only(),
                              options_.enable_empty_wrapper_null_unboxing,
                              enable_optional_types_),
-            expr.id());
+            expr.id(), /*stack_delta=*/0);
   }
 
   // Call node handler group.
@@ -1383,9 +1384,14 @@ class FlatExprVisitor : public cel::AstVisitor {
     std::vector<std::string> fields =
         std::move(status_or_resolved_fields.value().second);
 
+    size_t num_fields = fields.size();
+    int64_t stack_delta =
+        num_fields <= static_cast<size_t>(std::numeric_limits<int64_t>::max())
+            ? 1 - static_cast<int64_t>(num_fields)
+            : std::numeric_limits<int16_t>::max();
     AddStep(CreateCreateStructStep(std::move(resolved_name), std::move(fields),
                                    MakeOptionalIndicesSet(struct_expr)),
-            expr.id());
+            expr.id(), stack_delta);
   }
 
   void PostVisitMap(const cel::Expr& expr,
@@ -1406,9 +1412,15 @@ class FlatExprVisitor : public cel::AstVisitor {
       }
     }
 
-    AddStep(CreateCreateStructStepForMap(map_expr.entries().size(),
+    size_t num_entries = map_expr.entries().size();
+    int64_t stack_delta =
+        num_entries <=
+                static_cast<size_t>(std::numeric_limits<int64_t>::max() / 2)
+            ? 1 - 2 * static_cast<int64_t>(num_entries)
+            : std::numeric_limits<int16_t>::max();
+    AddStep(CreateCreateStructStepForMap(num_entries,
                                          MakeOptionalIndicesSet(map_expr)),
-            expr.id());
+            expr.id(), stack_delta);
   }
 
   absl::Status progress_status() const { return progress_status_; }
@@ -1471,21 +1483,22 @@ class FlatExprVisitor : public cel::AstVisitor {
   // may free the step at that point.
   template <typename T>
   std::enable_if_t<std::is_base_of_v<ExpressionStepLogic, T>, T*> AddStep(
-      std::unique_ptr<T> step, int64_t expr_id = -1) {
+      std::unique_ptr<T> step, int64_t expr_id = -1, int64_t stack_delta = 1) {
     if (progress_status_.ok() && !PlanningSuppressed()) {
       T* ptr = step.get();
-      program_builder_.AddStep(
-          ExpressionStep::MakeGenericStep(std::move(step), expr_id));
+      program_builder_.AddStep(ExpressionStep::MakeGenericStep(
+          std::move(step), expr_id, stack_delta));
       return ptr;
     }
     return nullptr;
   }
 
   template <typename T>
   std::enable_if_t<std::is_base_of_v<ExpressionStepLogic, T>, T*> AddStep(
-      absl::StatusOr<std::unique_ptr<T>> step, int64_t expr_id = -1) {
+      absl::StatusOr<std::unique_ptr<T>> step, int64_t expr_id = -1,
+      int64_t stack_delta = 1) {
     if (step.ok()) {
-      return AddStep(*std::move(step), expr_id);
+      return AddStep(*std::move(step), expr_id, stack_delta);
     } else {
       SetProgressStatusIfError(step.status());
     }
@@ -1728,7 +1741,7 @@ FlatExprVisitor::CallHandlerResult FlatExprVisitor::HandleIndex(
   }
 
   AddStep(CreateContainerAccessStep(call_expr, enable_optional_types_),
-          expr.id());
+          expr.id(), /*stack_delta=*/-1);
   return CallHandlerResult::kIntercepted;
 }
 
@@ -2041,7 +2054,7 @@ void ExhaustiveTernaryCondVisitor::PreVisit(const cel::Expr* expr) {
 }
 
 void ExhaustiveTernaryCondVisitor::PostVisit(const cel::Expr* expr) {
-  visitor_->AddStep(CreateTernaryStep(), expr->id());
+  visitor_->AddStep(CreateTernaryStep(), expr->id(), /*stack_delta=*/-2);
 }
 
 void ComprehensionVisitor::PreVisit(const cel::Expr* expr) {
@@ -2232,6 +2245,52 @@ std::vector<ExecutionPathView> FlattenExpressionTable(
   return subexpression_indexes;
 }
 
+std::optional<int64_t> CheckedDeltaAdd(int64_t current,
+                                       std::optional<int64_t> delta) {
+  if (!delta.has_value()) {
+    return std::nullopt;
+  }
+  if (*delta > 0 && current > std::numeric_limits<int64_t>::max() - *delta) {
+    return std::nullopt;
+  }
+  if (*delta < 0 && current < std::numeric_limits<int64_t>::min() - *delta) {
+    return std::nullopt;
+  }
+  current += *delta;
+  if (current < 0) {
+    return std::nullopt;
+  }
+  return current;
+}
+
+// Conservative estimate of the maximum value stack size needed for the given
+// subexpressions.
+//
+// If overflow occurs, returns fallback_size, which is the total number of
+// steps in the program.
+size_t EstimateMaxStackSize(absl::Span<const ExecutionPathView> subexpressions,
+                            size_t fallback_size) {
+  size_t total_max_stack = 0;
+  for (ExecutionPathView path : subexpressions) {
+    int64_t current = 0;
+    int64_t max_depth = 0;
+    for (const ExpressionStep& step : path) {
+      std::optional<int64_t> next = CheckedDeltaAdd(current, step.StackDelta());
+      if (!next.has_value()) {
+        return fallback_size;
+      }
+      current = *next;
+      max_depth = std::max(max_depth, current);
+    }
+    if (static_cast<uint64_t>(max_depth) >
+        std::numeric_limits<size_t>::max() - total_max_stack) {
+      return fallback_size;
+    }
+    total_max_stack += static_cast<size_t>(max_depth);
+  }
+  return total_max_stack;
+}
+
 absl::Status CheckAstExtensions(
     const std::vector<cel::ExtensionSpec>& extensions,
     bool enable_select_optimization) {
@@ -2266,11 +2325,27 @@ absl::StatusOr<FlatExpression> FlatExprBuilder::CreateExpressionImpl(
                                             ? RuntimeIssue::Severity::kWarning
                                             : RuntimeIssue::Severity::kError;
   IssueCollector issue_collector(max_severity);
+  auto record_issues = absl::MakeCleanup([issues, &issue_collector] {
+    if (issues == nullptr) {
+      return;
+    }
+    *issues = issue_collector.ExtractIssues();
+  });
 
-  CEL_ASSIGN_OR_RETURN(std::vector<cel::ExtensionSpec> runtime_extensions,
-                       ExtractAndValidateRuntimeExtensions(*ast));
-  CEL_RETURN_IF_ERROR(
-      CheckAstExtensions(runtime_extensions, enable_select_optimization_));
+  absl::StatusOr<std::vector<cel::ExtensionSpec>> runtime_extensions =
+      ExtractAndValidateRuntimeExtensions(*ast);
+
+  if (!runtime_extensions.ok()) {
+    CEL_RETURN_IF_ERROR(issue_collector.AddIssue(
+        RuntimeIssue::CreateError(runtime_extensions.status())));
+  } else {
+    absl::Status status =
+        CheckAstExtensions(*runtime_extensions, enable_select_optimization_);
+    if (!status.ok()) {
+      CEL_RETURN_IF_ERROR(
+          issue_collector.AddIssue(RuntimeIssue::CreateError(status)));
+    }
+  }
 
   Resolver resolver(container_, function_registry_, type_registry_,
                     GetTypeProvider(),
@@ -2318,17 +2393,16 @@ absl::StatusOr<FlatExpression> FlatExprBuilder::CreateExpressionImpl(
         "Expression too large to be executed, exceeds uint32_t limits.");
   }
 
-  if (issues != nullptr) {
-    (*issues) = issue_collector.ExtractIssues();
-  }
-
   ExecutionPath execution_path;
   std::vector<ExecutionPathView> subexpressions =
       FlattenExpressionTable(program_builder, execution_path);
+  size_t value_stack_size =
+      EstimateMaxStackSize(subexpressions, execution_path.size());
 
   return FlatExpression(std::move(execution_path), std::move(subexpressions),
                         visitor.slot_count(), visitor.iter_stack_count(),
-                        GetTypeProvider(), options_, std::move(arena));
+                        GetTypeProvider(), options_, std::move(arena),
+                        value_stack_size);
 }
 const cel::TypeProvider& FlatExprBuilder::GetTypeProvider() const {
   return use_legacy_type_provider_

eval/compiler/flat_expr_builder_extensions.h

@@ -334,9 +334,9 @@ class PlannerContext {
   absl::Status AddSubplanStep(const cel::Expr& node, ExpressionStep step);
   absl::Status AddSubplanStep(const cel::Expr& node,
                               std::unique_ptr<ExpressionStepLogic> step,
-                              int64_t expr_id = -1) {
-    return AddSubplanStep(
-        node, ExpressionStep::MakeGenericStep(std::move(step), expr_id));
+                              int64_t expr_id = -1, int64_t stack_delta = 1) {
+    return AddSubplanStep(node, ExpressionStep::MakeGenericStep(
+                                    std::move(step), expr_id, stack_delta));
   }
 
   const Resolver& resolver() const { return resolver_; }