Map root identifiers to slots and memoize.

eval/compiler/BUILD

@@ -149,8 +149,11 @@ cc_test(
         "//common:function_descriptor",
         "//common:kind",
         "//common:value",
+        "//eval/eval:cel_expression_flat_impl",
+        "//eval/eval:evaluator_core",
         "//eval/public:activation",
         "//eval/public:ast_rewrite",
+        "//eval/public:base_activation",
         "//eval/public:builtin_func_registrar",
         "//eval/public:cel_attribute",
         "//eval/public:cel_expr_builder_factory",

eval/compiler/constant_folding.cc

@@ -201,18 +201,18 @@ absl::Status ConstantFoldingExtension::OnPostVisit(PlannerContext& context,
     // This subexpression is already optimized out or suppressed.
     return absl::OkStatus();
   }
-  // copy string to managed handle if backed by the original program.
+  // Ensure that string constants are copied to the constant folding arena
   Value value;
   if (node.has_const_expr()) {
     CEL_ASSIGN_OR_RETURN(value,
                          ConvertConstant(node.const_expr(), state_.arena()));
   } else {
-    ExecutionFrame frame(subplan, empty_, context.options(), state_);
     state_.Reset();
     // Update stack size to accommodate sub expression.
     // This only results in a vector resize if the new maxsize is greater than
     // the current capacity.
     state_.value_stack().SetMaxSize(subplan.size());
+    ExecutionFrame frame(subplan, empty_, context.options(), state_);
 
     auto result = frame.Evaluate();
     // If this would be a runtime error, then don't adjust the program plan, but

eval/compiler/flat_expr_builder.cc

@@ -70,6 +70,7 @@
 #include "eval/eval/evaluator_core.h"
 #include "eval/eval/expression_step_logic.h"
 #include "eval/eval/function_step.h"
+#include "eval/eval/ident_step.h"
 #include "eval/eval/optional_or_step.h"
 #include "eval/eval/select_step.h"
 #include "eval/eval/shadowable_value_step.h"
@@ -892,8 +893,12 @@ class FlatExprVisitor : public cel::AstVisitor {
       return;
     }
 
+    if (PlanningSuppressed()) {
+      return;
+    }
     absl::string_view ident_name = absl::StripPrefix(ident_expr.name(), ".");
-      AddStep(ExpressionStep::MakeIdentifierStep(ident_name, expr.id()));
+    size_t ident_index = AddOrUpdateIdentifier(ident_name);
+    AddStep(ExpressionStep::MakeIdentifierStep(ident_index, expr.id()));
   }
 
   void PreVisitSelect(const cel::Expr& expr,
@@ -1547,6 +1552,15 @@ class FlatExprVisitor : public cel::AstVisitor {
     return index_manager_.IterStackRequirement();
   }
 
+  IdentifierTable BuildIdentifierTable(size_t& total_slot_count) && {
+    IdentifierTable table;
+    table.reserve(identifiers_.size());
+    for (std::string& name : identifiers_) {
+      table.push_back(Identifier{std::move(name), total_slot_count++});
+    }
+    return table;
+  }
+
   void AddOptimizer(std::unique_ptr<ProgramOptimizer> optimizer) {
     program_optimizers_.push_back(std::move(optimizer));
   }
@@ -1698,6 +1712,8 @@ class FlatExprVisitor : public cel::AstVisitor {
 
   void MaybeResolveType(const cel::Expr& expr);
 
+  size_t AddOrUpdateIdentifier(absl::string_view name);
+
   const Resolver& resolver_;
   const cel::TypeProvider& type_provider_;
   absl::Status progress_status_;
@@ -1726,6 +1742,8 @@ class FlatExprVisitor : public cel::AstVisitor {
   ProgramBuilder& program_builder_;
   PlannerContext& extension_context_;
   IndexManager index_manager_;
+  std::vector<std::string> identifiers_;
+  absl::flat_hash_map<std::string, size_t> identifier_map_;
 
   bool enable_optional_types_;
   std::optional<FlatExprVisitor::BlockInfo> block_;
@@ -1872,6 +1890,16 @@ void FlatExprVisitor::MaybeResolveType(const cel::Expr& expr) {
   SetResolvedType(expr, *type);
 }
 
+size_t FlatExprVisitor::AddOrUpdateIdentifier(absl::string_view name) {
+  auto it = identifier_map_.find(name);
+  if (it == identifier_map_.end()) {
+    size_t index = identifiers_.size();
+    it = identifier_map_.emplace(name, index).first;
+    identifiers_.push_back(std::string(name));
+  }
+  return it->second;
+}
+
 void LogicalCondVisitor::PreVisit(const cel::Expr* expr) {
   visitor_->ValidateOrError(
       !expr->call_expr().has_target() && expr->call_expr().args().size() >= 2,
@@ -2329,7 +2357,12 @@ absl::StatusOr<FlatExpression> FlatExprBuilder::CreateExpressionImpl(
     return visitor.progress_status();
   }
 
-  if (visitor.slot_count() > std::numeric_limits<uint32_t>::max() ||
+  size_t total_slot_count = visitor.slot_count();
+  size_t iterator_stack_size = visitor.iter_stack_count();
+  IdentifierTable identifiers =
+      std::move(visitor).BuildIdentifierTable(total_slot_count);
+
+  if (total_slot_count > std::numeric_limits<uint32_t>::max() ||
       program_builder.ExtractedSubexpressionCount() >
           std::numeric_limits<uint32_t>::max()) {
     // Impractical to trigger (we'd run out of memory first), but assuming this
@@ -2345,10 +2378,21 @@ absl::StatusOr<FlatExpression> FlatExprBuilder::CreateExpressionImpl(
   ExecutionPath execution_path;
   std::vector<ExecutionPathView> subexpressions =
       FlattenExpressionTable(program_builder, execution_path);
+  // Iterate over execution path and update identifier steps with resolved slot.
+  // TODO(b/572083203): Should invert this so the local slots are stored
+  // after the global ones.
+  for (ExpressionStep& step : execution_path) {
+    if (IdentStepInfo* ident_step = GetIfIdentStep(step);
+        ident_step != nullptr) {
+      ident_step->slot_index = static_cast<uint32_t>(
+          identifiers[ident_step->ident_index].slot_index);
+    }
+  }
 
   return FlatExpression(std::move(execution_path), std::move(subexpressions),
-                        visitor.slot_count(), visitor.iter_stack_count(),
-                        GetTypeProvider(), options_, std::move(arena));
+                        total_slot_count, iterator_stack_size,
+                        GetTypeProvider(), options_, std::move(arena),
+                        std::move(identifiers));
 }
 const cel::TypeProvider& FlatExprBuilder::GetTypeProvider() const {
   return use_legacy_type_provider_

eval/compiler/flat_expr_builder_test.cc

@@ -17,6 +17,7 @@
 #include <cstdint>
 #include <functional>
 #include <memory>
+#include <optional>
 #include <string>
 #include <utility>
 #include <vector>
@@ -41,8 +42,12 @@
 #include "eval/compiler/cel_expression_builder_flat_impl.h"
 #include "eval/compiler/constant_folding.h"
 #include "eval/compiler/qualified_reference_resolver.h"
+#include "eval/eval/cel_expression_flat_impl.h"
+#include "eval/eval/evaluator_core.h"
+#include "eval/eval/ident_step.h"
 #include "eval/public/activation.h"
 #include "eval/public/ast_rewrite.h"
+#include "eval/public/base_activation.h"
 #include "eval/public/builtin_func_registrar.h"
 #include "eval/public/cel_attribute.h"
 #include "eval/public/cel_expr_builder_factory.h"
@@ -3148,6 +3153,107 @@ TEST(FlatExprBuilderTest, ListAppendComprehensionWithTargetFailsPlanning) {
       StatusIs(absl::StatusCode::kInvalidArgument));
 }
 
+class CountingActivation : public BaseActivation {
+ public:
+  void InsertInt(absl::string_view name, int64_t val) {
+    values_[name] = CelValue::CreateInt64(val);
+  }
+
+  void InsertBool(absl::string_view name, bool val) {
+    values_[name] = CelValue::CreateBool(val);
+  }
+
+  std::vector<const CelFunction*> FindFunctionOverloads(
+      absl::string_view) const override {
+    return {};
+  }
+
+  std::optional<CelValue> FindValue(absl::string_view name,
+                                    google::protobuf::Arena*) const override {
+    lookup_counts_[name]++;
+    auto it = values_.find(name);
+    if (it != values_.end()) {
+      return it->second;
+    }
+    return std::nullopt;
+  }
+
+  int LookupCount(absl::string_view name) const {
+    auto it = lookup_counts_.find(name);
+    return it != lookup_counts_.end() ? it->second : 0;
+  }
+
+ private:
+  absl::flat_hash_map<std::string, CelValue> values_;
+  mutable absl::flat_hash_map<std::string, int> lookup_counts_;
+};
+
+TEST(FlatExprBuilderTest, IdentifierTableRepeatedRootIdentifierReceivesSlot) {
+  ASSERT_OK_AND_ASSIGN(ParsedExpr parsed_expr, parser::Parse("x + y + x"));
+  CelExpressionBuilderFlatImpl builder(NewTestingRuntimeEnv());
+  ASSERT_THAT(RegisterBuiltinFunctions(builder.GetRegistry()), IsOk());
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<CelExpression> cel_expr,
+                       builder.CreateExpression(&parsed_expr.expr(),
+                                                &parsed_expr.source_info()));
+
+  const FlatExpression& flat_expr =
+      static_cast<const CelExpressionFlatImpl*>(cel_expr.get())
+          ->flat_expression();
+  ASSERT_THAT(flat_expr.identifiers(), SizeIs(2));
+  EXPECT_EQ(flat_expr.identifiers()[0].name, "x");
+  EXPECT_EQ(flat_expr.identifiers()[0].slot_index, 0);
+  EXPECT_EQ(flat_expr.identifiers()[1].name, "y");
+  EXPECT_EQ(flat_expr.identifiers()[1].slot_index, 1);
+  EXPECT_EQ(flat_expr.comprehension_slots_size(), 2);
+  EXPECT_EQ(flat_expr.iterator_stack_size(), 0);
+
+  CountingActivation activation;
+  activation.InsertInt("x", 10);
+  activation.InsertInt("y", 20);
+  google::protobuf::Arena arena;
+  ASSERT_OK_AND_ASSIGN(CelValue result, cel_expr->Evaluate(activation, &arena));
+  EXPECT_THAT(result, test::IsCelInt64(40));
+  EXPECT_EQ(activation.LookupCount("x"), 1);
+  EXPECT_EQ(activation.LookupCount("y"), 1);
+}
+
+TEST(FlatExprBuilderTest,
+     IdentifierTableInComprehensionReceivesSlotAboveComprehensionSlots) {
+  ASSERT_OK_AND_ASSIGN(
+      ParsedExpr parsed_expr,
+      parser::Parse("[1, 2, 3].all(i, i < limit) && single_var"));
+  CelExpressionBuilderFlatImpl builder(NewTestingRuntimeEnv());
+  ASSERT_THAT(RegisterBuiltinFunctions(builder.GetRegistry()), IsOk());
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<CelExpression> cel_expr,
+                       builder.CreateExpression(&parsed_expr.expr(),
+                                                &parsed_expr.source_info()));
+
+  const FlatExpression& flat_expr =
+      static_cast<const CelExpressionFlatImpl*>(cel_expr.get())
+          ->flat_expression();
+  // Local comprehension variables (i, @result) are not in the root identifier
+  // table; only "limit" and "single_var" are present.
+  ASSERT_THAT(flat_expr.identifiers(), SizeIs(2));
+  EXPECT_EQ(flat_expr.identifiers()[0].name, "limit");
+  // Comprehension uses 2 slots (accu_var at 0, iter_var at 1), so "limit" gets
+  // slot 2 and "single_var" gets slot 3.
+  EXPECT_EQ(flat_expr.identifiers()[0].slot_index, 2);
+  EXPECT_EQ(flat_expr.identifiers()[1].name, "single_var");
+  EXPECT_EQ(flat_expr.identifiers()[1].slot_index, 3);
+  EXPECT_EQ(flat_expr.comprehension_slots_size(), 4);
+  EXPECT_EQ(flat_expr.iterator_stack_size(), 1);
+
+  CountingActivation activation;
+  activation.InsertInt("limit", 5);
+  activation.InsertBool("single_var", true);
+  google::protobuf::Arena arena;
+  ASSERT_OK_AND_ASSIGN(CelValue result, cel_expr->Evaluate(activation, &arena));
+  EXPECT_THAT(result, test::IsCelBool(true));
+  EXPECT_EQ(activation.LookupCount("limit"), 1);
+  EXPECT_EQ(activation.LookupCount("single_var"), 1);
+  EXPECT_EQ(activation.LookupCount("i"), 0);
+}
+
 }  // namespace
 
 }  // namespace google::api::expr::runtime

eval/eval/BUILD

@@ -85,6 +85,7 @@ cc_library(
         "//runtime/standard:equality_functions",
         "@com_google_absl//absl/base",
         "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/base:no_destructor",
         "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/log:absl_check",
@@ -448,6 +449,7 @@ cc_test(
         ":cel_expression_flat_impl",
         ":container_access_step",
         ":evaluator_core",
+        ":expression_step_logic",
         "//base:builtins",
         "//base:data",
         "//common:ast",
@@ -471,6 +473,7 @@ cc_test(
         "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/log:absl_check",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_cel_spec//proto/cel/expr:syntax_cc_proto",
         "@com_google_protobuf//:protobuf",
         "@com_google_protobuf//:struct_cc_proto",
@@ -510,7 +513,9 @@ cc_test(
         "//base:data",
         "//common:memory",
         "//eval/public:activation",
+        "//eval/public:base_activation",
         "//eval/public:cel_attribute",
+        "//eval/public:cel_function",
         "//eval/public:cel_value",
         "//internal:testing",
         "//runtime:runtime_options",

eval/eval/comprehension_slots.h

@@ -150,7 +150,7 @@ class ComprehensionSlots final {
   size_t size() const { return slots_.size(); }
 
  private:
-  absl::FixedArray<ComprehensionSlot, 0> slots_;
+  absl::FixedArray<ComprehensionSlot, 4> slots_;
 };
 
 }  // namespace google::api::expr::runtime

eval/eval/comprehension_step_test.cc

@@ -50,8 +50,9 @@ class ListKeysStepTest : public testing::Test {
     auto env = NewTestingRuntimeEnv();
     return std::make_unique<CelExpressionFlatImpl>(
         env,
-        FlatExpression(std::move(path), /*comprehension_slot_count=*/0,
-                       env->type_registry.GetComposedTypeProvider(), options));
+        FlatExpression(std::move(path), /*comprehension_slot_count=*/1,
+                       env->type_registry.GetComposedTypeProvider(), options,
+                       /*arena=*/nullptr, IdentifierTable{{"var", 0}}));
   }
 
  private:
@@ -75,7 +76,7 @@ MATCHER_P(CelStringValue, val, "") {
 
 TEST_F(ListKeysStepTest, MapPartiallyUnknown) {
   ExecutionPath path;
-  path.push_back(ExpressionStep::MakeIdentifierStep("var"));
+  path.push_back(ExpressionStep::MakeIdentifierStep(0));
   auto init_step =
       std::make_unique<ComprehensionInitStep>(/*iter_slot=*/0, /*accu_slot=*/0);
   init_step->set_error_jump_offset(1);
@@ -115,7 +116,7 @@ TEST_F(ListKeysStepTest, MapPartiallyUnknown) {
 
 TEST_F(ListKeysStepTest, ErrorPassedThrough) {
   ExecutionPath path;
-  path.push_back(ExpressionStep::MakeIdentifierStep("var"));
+  path.push_back(ExpressionStep::MakeIdentifierStep(0));
   auto init_step =
       std::make_unique<ComprehensionInitStep>(/*iter_slot=*/0, /*accu_slot=*/0);
   init_step->set_error_jump_offset(1);
@@ -140,7 +141,7 @@ TEST_F(ListKeysStepTest, ErrorPassedThrough) {
 
 TEST_F(ListKeysStepTest, UnknownSetPassedThrough) {
   ExecutionPath path;
-  path.push_back(ExpressionStep::MakeIdentifierStep("var"));
+  path.push_back(ExpressionStep::MakeIdentifierStep(0));
   auto init_step =
       std::make_unique<ComprehensionInitStep>(/*iter_slot=*/0, /*accu_slot=*/0);
   init_step->set_error_jump_offset(1);

eval/eval/container_access_step_test.cc

@@ -11,12 +11,15 @@
 #include "absl/base/nullability.h"
 #include "absl/log/absl_check.h"
 #include "absl/status/status.h"
+#include "absl/status/statusor.h"
 #include "base/builtins.h"
 #include "base/type_provider.h"
 #include "common/ast.h"
 #include "common/expr.h"
 #include "eval/eval/cel_expression_flat_impl.h"
 #include "eval/eval/evaluator_core.h"
+#include "eval/eval/expression_step_logic.h"
+#include "eval/eval/ident_step.h"
 #include "eval/public/activation.h"
 #include "eval/public/cel_attribute.h"
 #include "eval/public/cel_expr_builder_factory.h"
@@ -75,18 +78,21 @@ CelValue EvaluateAttributeHelper(
   container_expr.mutable_ident_expr().set_name("container");
   key_expr.mutable_ident_expr().set_name("key");
 
-    path.push_back(ExpressionStep::MakeIdentifierStep("container", 1));
-    path.push_back(ExpressionStep::MakeIdentifierStep("key", 2));
-    path.push_back(ExpressionStep::MakeGenericStep(
-        std::move(CreateContainerAccessStep(call).value()), 3));
+  path.push_back(ExpressionStep::MakeIdentifierStep(0, 1));
+  path.push_back(ExpressionStep::MakeIdentifierStep(1, 2));
+  absl::StatusOr<std::unique_ptr<ExpressionStepLogic>> access_step =
+      CreateContainerAccessStep(call);
+  ABSL_CHECK_OK(access_step);  // Crash OK
+  path.push_back(ExpressionStep::MakeGenericStep(*std::move(access_step), 3));
 
   cel::RuntimeOptions options;
   options.unknown_processing = cel::UnknownProcessingOptions::kAttributeOnly;
   options.enable_heterogeneous_equality = false;
   CelExpressionFlatImpl cel_expr(
-      env,
-      FlatExpression(std::move(path), /*comprehension_slot_count=*/0,
-                     env->type_registry.GetComposedTypeProvider(), options));
+      env, FlatExpression(std::move(path), /*comprehension_slot_count=*/2,
+                          env->type_registry.GetComposedTypeProvider(), options,
+                          /*arena=*/nullptr,
+                          IdentifierTable{{"container", 0}, {"key", 1}}));
   Activation activation;
 
   activation.InsertValue("container", container);