Add support for nested lazy init bind definitions.

PiperOrigin-RevId: 594103352
google · Dec 27, 2023 · 750c731 · 750c731
1 parent 5fb141d
commit 750c731
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 21 deletions.
diff --git a/eval/compiler/flat_expr_builder.cc b/eval/compiler/flat_expr_builder.cc
@@ -430,14 +430,6 @@ class FlatExprVisitor : public cel::ast_internal::AstVisitor {
           int subexpression = -1;
           if (record.should_lazy_eval) {
             subexpression = record.subexpression;
-            // Invalidate any other binds that depend on this var since we
-            // don't support nested lazy init yet.
-            for (int j = comprehension_stack_.size() - 1; j > i; j--) {
-              ComprehensionStackRecord& other_record = comprehension_stack_[j];
-              if (other_record.in_accu_init && other_record.should_lazy_eval) {
-                other_record.should_lazy_eval = false;
-              }
-            }
           }
           return {slot, subexpression};
         }

diff --git a/eval/eval/evaluator_core.cc b/eval/eval/evaluator_core.cc
@@ -62,11 +62,12 @@ const ExpressionStep* ExecutionFrame::Next() {
   size_t end_pos = execution_path_.size();
 
   if (pc_ < end_pos) return execution_path_[pc_++].get();
-  if (pc_ == end_pos && sub_frame_.has_value()) {
-    pc_ = sub_frame_->return_pc;
-    execution_path_ = sub_frame_->return_expression;
-    ABSL_DCHECK_EQ(value_stack().size(), sub_frame_->expected_stack_size);
-    sub_frame_.reset();
+  if (pc_ == end_pos && !call_stack_.empty()) {
+    pc_ = call_stack_.back().return_pc;
+    execution_path_ = call_stack_.back().return_expression;
+    ABSL_DCHECK_EQ(value_stack().size(),
+                   call_stack_.back().expected_stack_size);
+    call_stack_.pop_back();
     return Next();
   }
   if (pc_ > end_pos) {

diff --git a/eval/eval/evaluator_core.h b/eval/eval/evaluator_core.h
@@ -145,8 +145,7 @@ class ExecutionFrame {
                            state_.value_factory()),
         max_iterations_(options_.comprehension_max_iterations),
         iterations_(0),
-        subexpressions_(),
-        sub_frame_(absl::nullopt) {}
+        subexpressions_() {}
 
   ExecutionFrame(absl::Span<const ExecutionPathView> subexpressions,
                  const cel::ActivationInterface& activation,
@@ -162,8 +161,7 @@ class ExecutionFrame {
                            state_.value_factory()),
         max_iterations_(options_.comprehension_max_iterations),
         iterations_(0),
-        subexpressions_(subexpressions),
-        sub_frame_(absl::nullopt) {
+        subexpressions_(subexpressions) {
     ABSL_DCHECK(!subexpressions.empty());
   }
 
@@ -197,16 +195,15 @@ class ExecutionFrame {
   // Only intended for use in built-in notion of lazily evaluated
   // subexpressions.
   void Call(int return_pc_offset, size_t subexpression_index) {
-    ABSL_DCHECK(!sub_frame_.has_value());
     ABSL_DCHECK_LT(subexpression_index, subexpressions_.size());
     ExecutionPathView subexpression = subexpressions_[subexpression_index];
     ABSL_DCHECK(subexpression != execution_path_);
     int return_pc = static_cast<int>(pc_) + return_pc_offset;
     // return pc == size() is supported (a tail call).
     ABSL_DCHECK_GE(return_pc, 0);
     ABSL_DCHECK_LE(return_pc, static_cast<int>(execution_path_.size()));
-    sub_frame_ = SubFrame{static_cast<size_t>(return_pc),
-                          value_stack().size() + 1, execution_path_};
+    call_stack_.push_back(SubFrame{static_cast<size_t>(return_pc),
+                                   value_stack().size() + 1, execution_path_});
     pc_ = 0UL;
     execution_path_ = subexpression;
   }
@@ -291,7 +288,7 @@ class ExecutionFrame {
   const int max_iterations_;
   int iterations_;
   absl::Span<const ExecutionPathView> subexpressions_;
-  absl::optional<SubFrame> sub_frame_;
+  std::vector<SubFrame> call_stack_;
 };
 
 // A flattened representation of the input CEL AST.