Optimize bookkeeping in the program builder:

- don't apply incremental cleanup as program steps are rewritten, just delete everything with ProgramBuilder dtor.
- hint expected size for program nodes
- reserve size for the special function handlder map

PiperOrigin-RevId: 742338612

Author: jnthntatum

eval/compiler/constant_folding.cc

@@ -71,6 +71,11 @@ using ::google::api::expr::runtime::ProgramOptimizer;
 using ::google::api::expr::runtime::ProgramOptimizerFactory;
 using ::google::api::expr::runtime::Resolver;
 
+enum class IsConst {
+  kConditional,
+  kNonConst,
+};
+
 class ConstantFoldingExtension : public ProgramOptimizer {
  public:
   ConstantFoldingExtension(
@@ -92,10 +97,6 @@ class ConstantFoldingExtension : public ProgramOptimizer {
                            const Expr& node) override;
 
  private:
-  enum class IsConst {
-    kConditional,
-    kNonConst,
-  };
   // Most constant folding evaluations are simple
   // binary operators.
   static constexpr size_t kDefaultStackLimit = 4;
@@ -114,51 +115,42 @@ class ConstantFoldingExtension : public ProgramOptimizer {
   std::vector<IsConst> is_const_;
 };
 
-absl::Status ConstantFoldingExtension::OnPreVisit(PlannerContext& context,
-                                                  const Expr& node) {
-  struct IsConstVisitor {
-    IsConst operator()(const Constant&) { return IsConst::kConditional; }
-    IsConst operator()(const IdentExpr&) { return IsConst::kNonConst; }
-    IsConst operator()(const ComprehensionExpr&) {
+IsConst IsConstExpr(const Expr& expr, const Resolver& resolver) {
+  switch (expr.kind_case()) {
+    case ExprKindCase::kConstant:
+      return IsConst::kConditional;
+    case ExprKindCase::kIdentExpr:
+      return IsConst::kNonConst;
+    case ExprKindCase::kComprehensionExpr:
       // Not yet supported, need to identify whether range and
       // iter vars are compatible with const folding.
       return IsConst::kNonConst;
-    }
-    IsConst operator()(const StructExpr& create_struct) {
+    case ExprKindCase::kStructExpr:
       return IsConst::kNonConst;
-    }
-    IsConst operator()(const cel::MapExpr& map_expr) {
-      // Not yet supported but should be possible in the future.
+    case ExprKindCase::kMapExpr:
       // Empty maps are rare and not currently supported as they may eventually
       // have similar issues to empty list when used within comprehensions or
       // macros.
-      if (map_expr.entries().empty()) {
+      if (expr.map_expr().entries().empty()) {
         return IsConst::kNonConst;
       }
       return IsConst::kConditional;
-    }
-    IsConst operator()(const ListExpr& create_list) {
-      if (create_list.elements().empty()) {
-        // TODO: Don't fold for empty list to allow comprehension
+    case ExprKindCase::kListExpr:
+      if (expr.list_expr().elements().empty()) {
+        // Don't fold for empty list to allow comprehension
         // list append optimization.
         return IsConst::kNonConst;
       }
       return IsConst::kConditional;
-    }
-
-    IsConst operator()(const SelectExpr&) { return IsConst::kConditional; }
-
-    IsConst operator()(const cel::UnspecifiedExpr&) {
-      return IsConst::kNonConst;
-    }
-
-    IsConst operator()(const CallExpr& call) {
+    case ExprKindCase::kSelectExpr:
+      return IsConst::kConditional;
+    case ExprKindCase::kCallExpr: {
+      const auto& call = expr.call_expr();
       // Short Circuiting operators not yet supported.
       if (call.function() == kAnd || call.function() == kOr ||
           call.function() == kTernary) {
         return IsConst::kNonConst;
       }
-
       // For now we skip constant folding for cel.@block. We do not yet setup
       // slots. When we enable constant folding for comprehensions (like
       // cel.bind), we can address cel.@block.
@@ -167,23 +159,24 @@ absl::Status ConstantFoldingExtension::OnPreVisit(PlannerContext& context,
       }
 
       int arg_len = call.args().size() + (call.has_target() ? 1 : 0);
-      std::vector<cel::Kind> arg_matcher(arg_len, cel::Kind::kAny);
       // Check for any lazy overloads (activation dependant)
       if (!resolver
-               .FindLazyOverloads(call.function(), call.has_target(),
-                                  arg_matcher)
+               .FindLazyOverloads(call.function(), call.has_target(), arg_len)
                .empty()) {
         return IsConst::kNonConst;
       }
 
       return IsConst::kConditional;
     }
+    case ExprKindCase::kUnspecifiedExpr:
+    default:
+      return IsConst::kNonConst;
+  }
+}
 
-    const Resolver& resolver;
-  };
-
-  IsConst is_const =
-      absl::visit(IsConstVisitor{context.resolver()}, node.kind());
+absl::Status ConstantFoldingExtension::OnPreVisit(PlannerContext& context,
+                                                  const Expr& node) {
+  IsConst is_const = IsConstExpr(node, context.resolver());
   is_const_.push_back(is_const);
 
   return absl::OkStatus();

eval/compiler/constant_folding_test.cc

@@ -352,37 +352,37 @@ TEST_F(UpdatedConstantFoldingTest, CreatesLargeList) {
 
   ProgramBuilder program_builder;
   // Simulate the visitor order.
-  program_builder.EnterSubexpression(&create_list);
+  ASSERT_TRUE(program_builder.EnterSubexpression(&create_list) != nullptr);
 
   // 0
-  program_builder.EnterSubexpression(&elem0);
+  ASSERT_TRUE(program_builder.EnterSubexpression(&elem0) != nullptr);
   ASSERT_OK_AND_ASSIGN(auto step, CreateConstValueStep(cel::IntValue(1L), 1));
   program_builder.AddStep(std::move(step));
   program_builder.ExitSubexpression(&elem0);
 
   // 1
-  program_builder.EnterSubexpression(&elem1);
+  ASSERT_TRUE(program_builder.EnterSubexpression(&elem1));
   ASSERT_OK_AND_ASSIGN(step, CreateConstValueStep(cel::IntValue(2L), 2));
   program_builder.AddStep(std::move(step));
   program_builder.ExitSubexpression(&elem1);
 
   // 2
-  program_builder.EnterSubexpression(&elem2);
+  ASSERT_TRUE(program_builder.EnterSubexpression(&elem2) != nullptr);
   ASSERT_OK_AND_ASSIGN(step, CreateConstValueStep(cel::IntValue(3L), 3));
   program_builder.AddStep(std::move(step));
   program_builder.ExitSubexpression(&elem2);
 
   // 3
-  program_builder.EnterSubexpression(&elem2);
+  ASSERT_TRUE(program_builder.EnterSubexpression(&elem3) != nullptr);
   ASSERT_OK_AND_ASSIGN(step, CreateConstValueStep(cel::IntValue(4L), 4));
   program_builder.AddStep(std::move(step));
-  program_builder.ExitSubexpression(&elem2);
+  program_builder.ExitSubexpression(&elem3);
 
   // 4
-  program_builder.EnterSubexpression(&elem2);
+  ASSERT_TRUE(program_builder.EnterSubexpression(&elem4) != nullptr);
   ASSERT_OK_AND_ASSIGN(step, CreateConstValueStep(cel::IntValue(5L), 5));
   program_builder.AddStep(std::move(step));
-  program_builder.ExitSubexpression(&elem2);
+  program_builder.ExitSubexpression(&elem4);
 
   // createlist
   ASSERT_OK_AND_ASSIGN(step, CreateCreateListStep(create_list.list_expr(), 6));

eval/compiler/flat_expr_builder.cc

@@ -257,6 +257,31 @@ class ExhaustiveTernaryCondVisitor : public CondVisitor {
   FlatExprVisitor* visitor_;
 };
 
+// Returns a hint for the number of program nodes (steps or subexpressions) that
+// will be created for this expr.
+size_t SizeHint(const cel::Expr& expr) {
+  switch (expr.kind_case()) {
+    case cel::ExprKindCase::kConstant:
+      return 1;
+    case cel::ExprKindCase::kIdentExpr:
+      return 1;
+    case cel::ExprKindCase::kSelectExpr:
+      return 2;
+    case cel::ExprKindCase::kCallExpr:
+      return expr.call_expr().args().size() +
+             (expr.call_expr().has_target() ? 2 : 1);
+    case cel::ExprKindCase::kListExpr:
+      return expr.list_expr().elements().size() + 1;
+    case cel::ExprKindCase::kStructExpr:
+      return expr.struct_expr().fields().size() + 1;
+    case cel::ExprKindCase::kMapExpr:
+      return 2 * expr.struct_expr().fields().size() + 1;
+    default:
+      return 1;
+  }
+  return 0;
+}
+
 // Returns whether this comprehension appears to be a standard map/filter
 // macro implementation. It is not exhaustive, so it is unsafe to use with
 // custom comprehensions outside of the standard macros or hand crafted ASTs.
@@ -491,6 +516,8 @@ class FlatExprVisitor : public cel::AstVisitor {
         program_builder_(program_builder),
         extension_context_(extension_context),
         enable_optional_types_(enable_optional_types) {
+    constexpr size_t kCallHandlerSizeHint = 11;
+    call_handlers_.reserve(kCallHandlerSizeHint);
     call_handlers_[cel::builtin::kIndex] = [this](const cel::Expr& expr,
                                                   const cel::CallExpr& call) {
       return HandleIndex(expr, call);
@@ -565,7 +592,13 @@ class FlatExprVisitor : public cel::AstVisitor {
       }
     }
 
-    program_builder_.EnterSubexpression(&expr);
+    auto* subexpression =
+        program_builder_.EnterSubexpression(&expr, SizeHint(expr));
+    if (subexpression == nullptr) {
+      progress_status_.Update(
+          absl::InternalError("same CEL expr visited twice"));
+      return;
+    }
 
     for (const std::unique_ptr<ProgramOptimizer>& optimizer :
          program_optimizers_) {