Allow constant folding to recognize identifiers replaced with inlined constants.

eval/compiler/constant_folding.cc

@@ -190,6 +190,14 @@ absl::Status ConstantFoldingExtension::OnPostVisit(PlannerContext& context,
   is_const_.pop_back();
 
   if (is_const == IsConst::kNonConst) {
+    if (context.options().enable_subexpression_inlining &&
+        node.has_ident_expr()) {
+      ExecutionPathView subplan = context.GetSubplan(node);
+      if (subplan.size() == 1 &&
+          google::api::expr::runtime::IsConstant(subplan[0])) {
+        return absl::OkStatus();
+      }
+    }
     // update parent
     if (!is_const_.empty()) {
       is_const_.back() = IsConst::kNonConst;

eval/compiler/constant_folding_test.cc

@@ -564,6 +564,55 @@ TEST_F(UpdatedConstantFoldingTest, ErrorsOnUnexpectedOrder) {
               StatusIs(absl::StatusCode::kInternal));
 }
 
+TEST_F(UpdatedConstantFoldingTest, FoldsListWithInlinedConstantIdent) {
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<cel::Ast> list_ast,
+                       ParseFromCel("[a, 2]"));
+  const Expr& create_list = list_ast->root_expr();
+  const Expr& elem_a = create_list.list_expr().elements()[0].expr();
+  const Expr& elem_two = create_list.list_expr().elements()[1].expr();
+
+  ProgramBuilder program_builder;
+  program_builder.EnterSubexpression(&create_list);
+  // Simulate `a` having been inlined as a constant step by FlatExprVisitor.
+  program_builder.EnterSubexpression(&elem_a);
+  program_builder.AddStep(
+      ExpressionStep::MakeConstant(cel::IntValue(10), elem_a.id()));
+  program_builder.ExitSubexpression(&elem_a);
+
+  program_builder.EnterSubexpression(&elem_two);
+  program_builder.AddStep(
+      ExpressionStep::MakeConstant(cel::IntValue(2), elem_two.id()));
+  program_builder.ExitSubexpression(&elem_two);
+
+  program_builder.AddStep(CreateCreateListStep(
+      create_list.list_expr().elements().size(), {}, create_list.id()));
+  program_builder.ExitSubexpression(&create_list);
+
+  options_.enable_subexpression_inlining = true;
+  std::shared_ptr<google::protobuf::Arena> arena;
+  PlannerContext context(env_, resolver_, options_,
+                         type_registry_.GetComposedTypeProvider(),
+                         issue_collector_, program_builder, arena);
+
+  ProgramOptimizerFactory constant_folder_factory =
+      CreateConstantFoldingOptimizer();
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<ProgramOptimizer> constant_folder,
+                       constant_folder_factory(context, *list_ast));
+
+  ASSERT_THAT(constant_folder->OnPreVisit(context, create_list), IsOk());
+  ASSERT_THAT(constant_folder->OnPreVisit(context, elem_a), IsOk());
+  ASSERT_THAT(constant_folder->OnPostVisit(context, elem_a), IsOk());
+  ASSERT_THAT(constant_folder->OnPreVisit(context, elem_two), IsOk());
+  ASSERT_THAT(constant_folder->OnPostVisit(context, elem_two), IsOk());
+  ASSERT_THAT(constant_folder->OnPostVisit(context, create_list), IsOk());
+
+  ExecutionPath path = std::move(program_builder).FlattenMain();
+  EXPECT_THAT(path, SizeIs(1));
+  cel::Value folded_list;
+  ASSERT_TRUE(google::api::expr::runtime::GetIfConstant(path[0], folded_list));
+  EXPECT_TRUE(folded_list.IsList());
+}
+
 }  // namespace
 
 }  // namespace cel::runtime_internal

eval/compiler/flat_expr_builder_test.cc

@@ -3744,6 +3744,106 @@ TEST(FlatExprBuilderTest, BlockInlinesTrivialSubexpressionsDisabled) {
   // Here trace reports the init step on first access.
   EXPECT_THAT(traced_ids, testing::ElementsAre(3, 4, 6, 9, 5, 1));
 }
+
+TEST(FlatExprBuilderTest, BindInlinesConstantAndFoldsMainExpression) {
+  ParsedExpr parsed_expr;
+  // Equivalent to: cel.bind(a, 1 + 2, (a + 3) + x)
+  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(
+      R"pb(
+        expr {
+          id: 1
+          comprehension_expr {
+            iter_var: "#unused"
+            iter_range {
+              id: 2
+              list_expr {}
+            }
+            accu_var: "a"
+            accu_init {
+              id: 3
+              call_expr {
+                function: "_+_"
+                args {
+                  id: 4
+                  const_expr { int64_value: 1 }
+                }
+                args {
+                  id: 5
+                  const_expr { int64_value: 2 }
+                }
+              }
+            }
+            loop_condition {
+              id: 6
+              const_expr { bool_value: false }
+            }
+            loop_step {
+              id: 7
+              ident_expr { name: "a" }
+            }
+            result {
+              id: 8
+              call_expr {
+                function: "_+_"
+                args {
+                  id: 9
+                  call_expr {
+                    function: "_+_"
+                    args {
+                      id: 10
+                      ident_expr { name: "a" }
+                    }
+                    args {
+                      id: 11
+                      const_expr { int64_value: 3 }
+                    }
+                  }
+                }
+                args {
+                  id: 12
+                  ident_expr { name: "x" }
+                }
+              }
+            }
+          }
+        }
+      )pb",
+      &parsed_expr));
+
+  InterpreterOptions options;
+  google::protobuf::Arena arena;
+  options.constant_folding = true;
+  options.constant_arena = &arena;
+  options.enable_subexpression_inlining = true;
+
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<CelExpressionBuilder> builder,
+      CreateConstantFoldingConformanceTestExprBuilder(options));
+  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_FALSE(flat_expr.subexpressions().empty());
+  ExecutionPathView main_path = flat_expr.subexpressions()[0];
+  // Main expression consists of:
+  //   0: Constant(6) (folded from `a + 3` after `a` was inlined as Constant(3))
+  //   1: Identifier(x)
+  //   2: _+_
+  //   3: ClearSlot
+  ASSERT_THAT(main_path, SizeIs(4));
+  cel::Value folded_val;
+  ASSERT_TRUE(GetIfConstant(main_path[0], folded_val));
+  EXPECT_EQ(folded_val.GetInt().NativeValue(), 6);
+  EXPECT_EQ(main_path[0].id(), 9);
+
+  Activation activation;
+  activation.InsertValue("x", CelValue::CreateInt64(10));
+  ASSERT_OK_AND_ASSIGN(CelValue result, cel_expr->Evaluate(activation, &arena));
+  EXPECT_THAT(result, test::IsCelInt64(16));
+}
 }  // namespace
 
 }  // namespace google::api::expr::runtime

extensions/BUILD

@@ -910,6 +910,7 @@ cc_test(
     name = "set_membership_optimization_test",
     srcs = ["set_membership_optimization_test.cc"],
     deps = [
+        ":bindings_ext",
         ":set_membership_optimization",
         ":sets_functions",
         "//base:attributes",

extensions/set_membership_optimization_test.cc

@@ -42,6 +42,7 @@
 #include "eval/compiler/flat_expr_builder.h"
 #include "eval/compiler/qualified_reference_resolver.h"
 #include "eval/eval/evaluator_core.h"
+#include "extensions/bindings_ext.h"
 #include "extensions/sets_functions.h"
 #include "internal/status_macros.h"
 #include "internal/testing.h"
@@ -89,6 +90,7 @@ absl::StatusOr<std::unique_ptr<Compiler>> CreateTestCompiler() {
       NewCompilerBuilder(cel::internal::GetTestingDescriptorPool()));
   CEL_RETURN_IF_ERROR(builder->AddLibrary(StandardCompilerLibrary()));
   CEL_RETURN_IF_ERROR(builder->AddLibrary(SetsCompilerLibrary()));
+  CEL_RETURN_IF_ERROR(builder->AddLibrary(BindingsCompilerLibrary()));
 
   TypeCheckerBuilder& checker = builder->GetCheckerBuilder();
   checker.set_container("cel.expr.conformance.proto2");
@@ -395,6 +397,40 @@ TEST(SetMembershipOptimizationPlanTest,
   }
 }
 
+TEST(SetMembershipOptimizationPlanTest,
+     EmitsCompactPlanWithInlinedBoundConstantSet) {
+  RuntimeOptions options;
+  options.container = "cel.expr.conformance.proto2";
+  options.enable_heterogeneous_equality = true;
+  options.enable_subexpression_inlining = true;
+  std::shared_ptr<cel::runtime_internal::RuntimeEnv> env =
+      NewTestingRuntimeEnv();
+  ASSERT_THAT(RegisterStandardFunctions(env->function_registry, options),
+              IsOk());
+  ASSERT_THAT(RegisterSetsFunctions(env->function_registry, options), IsOk());
+  FlatExprBuilder builder(std::move(env), options);
+  builder.set_container(options.container);
+  builder.AddAstTransform(
+      NewReferenceResolverExtension(ReferenceResolverOption::kAlways));
+  builder.AddProgramOptimizer(
+      cel::runtime_internal::CreateConstantFoldingOptimizer());
+  builder.AddProgramOptimizer(
+      CreateSetMembershipOptimizationProgramOptimizer());
+
+  for (absl::string_view expr_str : {
+           "cel.bind(set_var, [1, 2, 3, 4], x in set_var)",
+           "cel.bind(set_var, ['a', 'b', 'c', 'd'], s in set_var)",
+       }) {
+    ASSERT_OK_AND_ASSIGN(std::unique_ptr<Ast> ast, CompileTestExpr(expr_str));
+    ASSERT_OK_AND_ASSIGN(FlatExpression plan,
+                         builder.CreateExpressionImpl(std::move(ast), nullptr));
+    ASSERT_FALSE(plan.subexpressions().empty());
+    // Main expression has 3 steps: Identifier(x/s), Bitset/StringSet step, and
+    // ClearSlot.
+    EXPECT_THAT(plan.subexpressions()[0], SizeIs(3)) << expr_str;
+  }
+}
+
 TEST(SetMembershipOptimizationPlanTest, EmitsCompactBitsetPlanForUncheckedAst) {
   RuntimeOptions options;
   options.container = "cel.expr.conformance.proto2";