Ensure the lifecycle of rewritten expressions is preserved in the output CelExpression

PiperOrigin-RevId: 347652258

Author: TristonianJones

eval/compiler/flat_expr_builder.cc

@@ -1,5 +1,7 @@
 #include "eval/compiler/flat_expr_builder.h"
 
+#include <memory>
+
 #include "google/api/expr/v1alpha1/checked.pb.h"
 #include "stack"
 #include "absl/container/node_hash_map.h"
@@ -658,8 +660,8 @@ void ComprehensionVisitor::PreVisit(const Expr*) {
 
 void ComprehensionVisitor::PostVisitArg(int arg_num, const Expr* expr) {
   const Comprehension* comprehension = &expr->comprehension_expr();
-  const auto accu_var = comprehension->accu_var();
-  const auto iter_var = comprehension->iter_var();
+  const auto& accu_var = comprehension->accu_var();
+  const auto& iter_var = comprehension->iter_var();
   // TODO(issues/20): Consider refactoring the comprehension prologue step.
   switch (arg_num) {
     case ITER_RANGE: {
@@ -732,7 +734,7 @@ FlatExprBuilder::CreateExpressionImpl(
 
   const Expr* effective_expr = expr;
   // transformed expression preserving expression IDs
-  Expr rewrite_buffer;
+  std::unique_ptr<Expr> rewrite_buffer = nullptr;
   // TODO(issues/98): A type checker may perform these rewrites, but there
   // currently  isn't a signal to expose that in an expression. If that becomes
   // available, we can skip the reference resolve step here if it's already
@@ -745,19 +747,19 @@ FlatExprBuilder::CreateExpressionImpl(
       return rewritten.status();
     }
     if (rewritten.value().has_value()) {
-      rewrite_buffer = std::move(rewritten)->value();
-      effective_expr = &rewrite_buffer;
+      rewrite_buffer =
+          std::make_unique<Expr>(std::move(rewritten).value().value());
+      effective_expr = rewrite_buffer.get();
     }
     // TODO(issues/99): we could setup a check step here that confirms all of
     // references are defined before actually evaluating.
   }
 
+  Expr const_fold_buffer;
   if (constant_folding_) {
-    Expr buffer;
     FoldConstants(*effective_expr, *this->GetRegistry(), constant_arena_,
-                  idents, &buffer);
-    rewrite_buffer = std::move(buffer);
-    effective_expr = &rewrite_buffer;
+                  idents, &const_fold_buffer);
+    effective_expr = &const_fold_buffer;
   }
 
   std::set<std::string> iter_variable_names;
@@ -776,7 +778,8 @@ FlatExprBuilder::CreateExpressionImpl(
       absl::make_unique<CelExpressionFlatImpl>(
           expr, std::move(execution_path), comprehension_max_iterations_,
           std::move(iter_variable_names), enable_unknowns_,
-          enable_unknown_function_results_, enable_missing_attribute_errors_);
+          enable_unknown_function_results_, enable_missing_attribute_errors_,
+          std::move(rewrite_buffer));
 
   if (warnings != nullptr) {
     *warnings = std::move(warnings_builder).warnings();

eval/compiler/flat_expr_builder_test.cc

@@ -614,6 +614,60 @@ TEST(FlatExprBuilderTest, CheckedExprActivationMissesReferences) {
   EXPECT_FALSE(result.BoolOrDie());
 }
 
+TEST(FlatExprBuilderTest, CheckedExprWithReferenceMapAndConstantFolding) {
+  CheckedExpr expr;
+  // {`var1`: 'hello'}
+  google::protobuf::TextFormat::ParseFromString(R"(
+    reference_map {
+      key: 3
+      value {
+        name: "var1"
+        value {
+          int64_value: 1
+        }
+      }
+    }
+    expr {
+      id: 1
+      struct_expr {
+        entries {
+          id: 2
+          map_key {
+            id: 3
+            ident_expr {
+              name: "var1"
+            }
+          }
+          value {
+            id: 4
+            const_expr {
+              string_value: "hello"
+            }
+          }
+        }
+      }
+    })",
+                                      &expr);
+
+  FlatExprBuilder builder;
+  google::protobuf::Arena arena;
+  builder.set_constant_folding(true, &arena);
+  ASSERT_OK(RegisterBuiltinFunctions(builder.GetRegistry()));
+  auto build_status = builder.CreateExpression(&expr);
+  ASSERT_OK(build_status);
+
+  auto cel_expr = std::move(build_status.value());
+
+  Activation activation;
+  auto result_or = cel_expr->Evaluate(activation, &arena);
+  ASSERT_OK(result_or);
+  CelValue result = result_or.value();
+  ASSERT_TRUE(result.IsMap());
+  auto m = result.MapOrDie();
+  auto v = (*m)[CelValue::CreateInt64(1L)];
+  EXPECT_THAT(v.value().StringOrDie().value(), Eq("hello"));
+}
+
 TEST(FlatExprBuilderTest, ComprehensionWorksForError) {
   Expr expr;
   // {}[0].all(x, x) should evaluate OK but return an error value

eval/eval/evaluator_core.h

@@ -67,7 +67,7 @@ using ExecutionPath = std::vector<std::unique_ptr<const ExpressionStep>>;
 // stack as Span<>.
 class ValueStack {
  public:
-  ValueStack(size_t max_size) : current_size_(0) {
+  explicit ValueStack(size_t max_size) : current_size_(0) {
     stack_.resize(max_size);
     attribute_stack_.resize(max_size);
   }
@@ -336,8 +336,10 @@ class CelExpressionFlatImpl : public CelExpression {
                         std::set<std::string> iter_variable_names,
                         bool enable_unknowns = false,
                         bool enable_unknown_function_results = false,
-                        bool enable_missing_attribute_errors = false)
-      : path_(std::move(path)),
+                        bool enable_missing_attribute_errors = false,
+                        std::unique_ptr<Expr> rewritten_expr = nullptr)
+      : rewritten_expr_(std::move(rewritten_expr)),
+        path_(std::move(path)),
         max_iterations_(max_iterations),
         iter_variable_names_(std::move(iter_variable_names)),
         enable_unknowns_(enable_unknowns),
@@ -372,6 +374,8 @@ class CelExpressionFlatImpl : public CelExpression {
                                  CelEvaluationListener callback) const override;
 
  private:
+  // Maintain lifecycle of a modified expression.
+  std::unique_ptr<Expr> rewritten_expr_;
   const ExecutionPath path_;
   const int max_iterations_;
   const std::set<std::string> iter_variable_names_;