Simplify program layout for variadic and/or.

eval/compiler/flat_expr_builder.cc

@@ -2172,21 +2172,27 @@ void LogicalCondVisitor::PostVisitArg(int arg_num, const cel::Expr* expr) {
     return;
   }
   const int last_arg_index = expr->call_expr().args().size() - 1;
-  if (arg_num > 0) {
+  const size_t num_args = expr->call_expr().args().size();
+  if (arg_num == last_arg_index) {
     if (is_or_) {
-      visitor_->AddStep(CreateOrStep(expr->id()));
+      visitor_->AddStep(CreateOrStep(num_args, expr->id()));
     } else {
-      visitor_->AddStep(CreateAndStep(expr->id()));
+      visitor_->AddStep(CreateAndStep(num_args, expr->id()));
     }
     if (short_circuiting_ && !jump_steps_.empty()) {
-      visitor_->SetProgressStatusIfError(
-          jump_steps_.back().set_target(visitor_->GetCurrentIndex()));
+      for (auto& jump : jump_steps_) {
+        visitor_->SetProgressStatusIfError(
+            jump.set_target(visitor_->GetCurrentIndex()));
+      }
     }
   }
   if (short_circuiting_ && arg_num < last_arg_index) {
     std::unique_ptr<JumpStepBase> jump_step =
-        is_or_ ? CreateCondJumpStep(true, {}, expr->id())
-               : CreateCondJumpStep(false, {}, expr->id());
+        is_or_
+            ? CreateCondJumpStep(true, {}, /*expected_stack_size=*/arg_num + 1,
+                                 expr->id())
+            : CreateCondJumpStep(false, {}, /*expected_stack_size=*/arg_num + 1,
+                                 expr->id());
     ProgramStepIndex index = visitor_->GetCurrentIndex();
     if (JumpStepBase* jump_step_ptr = visitor_->AddStep(std::move(jump_step));
         jump_step_ptr) {

eval/eval/BUILD

@@ -418,7 +418,6 @@ cc_library(
         "//common:value",
         "//eval/internal:errors",
         "@com_google_absl//absl/status",
-        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/types:optional",
         "@com_google_cel_spec//proto/cel/expr:syntax_cc_proto",
     ],

eval/eval/jump_step.cc

@@ -143,9 +143,10 @@ class BoolCheckJumpStep : public JumpStepBase {
 
 // Factory method for Conditional Jump step.
 std::unique_ptr<JumpStepBase> CreateCondJumpStep(
-    bool jump_condition, absl::optional<int> jump_offset, int64_t expr_id) {
-  return std::make_unique<CondJumpStep>(jump_condition, jump_offset, 1,
-                                        expr_id);
+    bool jump_condition, absl::optional<int> jump_offset,
+    size_t expected_stack_size, int64_t expr_id) {
+  return std::make_unique<CondJumpStep>(jump_condition, jump_offset,
+                                        expected_stack_size, expr_id);
 }
 
 // Factory method for Ternary Conditional Jump step.

eval/eval/jump_step.h

@@ -53,7 +53,8 @@ std::unique_ptr<JumpStepBase> CreateJumpStep(absl::optional<int> jump_offset,
 // It is compared to jump_condition, and if matched, jump is performed.
 // The boolean value is left on top of the stack.
 std::unique_ptr<JumpStepBase> CreateCondJumpStep(
-    bool jump_condition, absl::optional<int> jump_offset, int64_t expr_id);
+    bool jump_condition, absl::optional<int> jump_offset,
+    size_t expected_stack_size, int64_t expr_id);
 
 // Factory method for Ternary Conditional Jump step.
 // Requires a boolean condition value on top of the stack.

eval/eval/logic_step.cc

@@ -450,13 +450,15 @@ std::unique_ptr<DirectExpressionStep> CreateDirectOrStep(
 }
 
 // Factory method for "And" Execution step
-absl::StatusOr<std::unique_ptr<ExpressionStep>> CreateAndStep(int64_t expr_id) {
-  return std::make_unique<LogicalOpStep>(OpType::kAnd, 2, expr_id);
+absl::StatusOr<std::unique_ptr<ExpressionStep>> CreateAndStep(size_t num_args,
+                                                              int64_t expr_id) {
+  return std::make_unique<LogicalOpStep>(OpType::kAnd, num_args, expr_id);
 }
 
 // Factory method for "Or" Execution step
-absl::StatusOr<std::unique_ptr<ExpressionStep>> CreateOrStep(int64_t expr_id) {
-  return std::make_unique<LogicalOpStep>(OpType::kOr, 2, expr_id);
+absl::StatusOr<std::unique_ptr<ExpressionStep>> CreateOrStep(size_t num_args,
+                                                             int64_t expr_id) {
+  return std::make_unique<LogicalOpStep>(OpType::kOr, num_args, expr_id);
 }
 
 // Factory method for recursive logical not "!" Execution step

eval/eval/logic_step.h

@@ -23,10 +23,12 @@ std::unique_ptr<DirectExpressionStep> CreateDirectOrStep(
     bool shortcircuiting);
 
 // Factory method for "And" Execution step
-absl::StatusOr<std::unique_ptr<ExpressionStep>> CreateAndStep(int64_t expr_id);
+absl::StatusOr<std::unique_ptr<ExpressionStep>> CreateAndStep(size_t num_args,
+                                                              int64_t expr_id);
 
 // Factory method for "Or" Execution step
-absl::StatusOr<std::unique_ptr<ExpressionStep>> CreateOrStep(int64_t expr_id);
+absl::StatusOr<std::unique_ptr<ExpressionStep>> CreateOrStep(size_t num_args,
+                                                             int64_t expr_id);
 
 // Factory method for recursive logical not "!" Execution step
 std::unique_ptr<DirectExpressionStep> CreateDirectNotStep(

eval/eval/logic_step_test.cc

@@ -73,7 +73,9 @@ class LogicStepTest : public testing::TestWithParam<bool> {
     CEL_ASSIGN_OR_RETURN(step, CreateIdentStep("name1", /*expr_id=*/-1));
     path.push_back(std::move(step));
 
-    CEL_ASSIGN_OR_RETURN(step, (is_or) ? CreateOrStep(2) : CreateAndStep(2));
+    CEL_ASSIGN_OR_RETURN(
+        step, (is_or) ? CreateOrStep(/*num_args=*/2, /*expr_id=*/2)
+                      : CreateAndStep(/*num_args=*/2, /*expr_id=*/2));
     path.push_back(std::move(step));
 
     auto dummy_expr = std::make_unique<Expr>();