Add ClearSlotsStep

PiperOrigin-RevId: 657215417
google · Jul 29, 2024 · 80b9594 · 80b9594
1 parent 7c81b49
commit 80b9594
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 0 deletions.
diff --git a/eval/eval/lazy_init_step.cc b/eval/eval/lazy_init_step.cc
@@ -141,6 +141,25 @@ class ClearSlotStep : public ExpressionStepBase {
   size_t slot_index_;
 };
 
+class ClearSlotsStep final : public ExpressionStepBase {
+ public:
+  explicit ClearSlotsStep(size_t slot_index, size_t slot_count, int64_t expr_id)
+      : ExpressionStepBase(expr_id),
+        slot_index_(slot_index),
+        slot_count_(slot_count) {}
+
+  absl::Status Evaluate(ExecutionFrame* frame) const override {
+    for (size_t i = 0; i < slot_count_; ++i) {
+      frame->comprehension_slots().ClearSlot(slot_index_ + i);
+    }
+    return absl::OkStatus();
+  }
+
+ private:
+  const size_t slot_index_;
+  const size_t slot_count_;
+};
+
 }  // namespace
 
 std::unique_ptr<DirectExpressionStep> CreateDirectBindStep(
@@ -172,4 +191,11 @@ std::unique_ptr<ExpressionStep> CreateClearSlotStep(size_t slot_index,
   return std::make_unique<ClearSlotStep>(slot_index, expr_id);
 }
 
+std::unique_ptr<ExpressionStep> CreateClearSlotsStep(size_t slot_index,
+                                                     size_t slot_count,
+                                                     int64_t expr_id) {
+  ABSL_DCHECK_GT(slot_count, 0);
+  return std::make_unique<ClearSlotsStep>(slot_index, slot_count, expr_id);
+}
+
 }  // namespace google::api::expr::runtime
diff --git a/eval/eval/lazy_init_step.h b/eval/eval/lazy_init_step.h
@@ -73,6 +73,10 @@ std::unique_ptr<ExpressionStep> CreateAssignSlotAndPopStep(size_t slot_index);
 std::unique_ptr<ExpressionStep> CreateClearSlotStep(size_t slot_index,
                                                     int64_t expr_id);
 
+std::unique_ptr<ExpressionStep> CreateClearSlotsStep(size_t slot_index,
+                                                     size_t slot_count,
+                                                     int64_t expr_id);
+
 }  // namespace google::api::expr::runtime
 
 #endif  // THIRD_PARTY_CEL_CPP_EVAL_EVAL_LAZY_INIT_STEP_H_
diff --git a/eval/eval/lazy_init_step_test.cc b/eval/eval/lazy_init_step_test.cc
@@ -39,6 +39,7 @@ using ::cel::RuntimeOptions;
 using ::cel::TypeProvider;
 using ::cel::ValueManager;
 using ::cel::extensions::ProtoMemoryManagerRef;
+using testing::IsNull;
 
 class LazyInitStepTest : public testing::Test {
  private:
@@ -142,5 +143,21 @@ TEST_F(LazyInitStepTest, CreateClearSlotStepBasic) {
   ASSERT_TRUE(slot == nullptr);
 }
 
+TEST_F(LazyInitStepTest, CreateClearSlotsStepBasic) {
+  ExecutionPath path;
+
+  path.push_back(CreateClearSlotsStep(0, 2, -1));
+
+  ExecutionFrame frame(path, activation_, runtime_options_, evaluator_state_);
+  frame.comprehension_slots().Set(0, value_factory().CreateIntValue(42));
+  frame.comprehension_slots().Set(1, value_factory().CreateIntValue(42));
+
+  // This will error because no return value, step will still evaluate.
+  frame.Evaluate().IgnoreError();
+
+  EXPECT_THAT(frame.comprehension_slots().Get(0), IsNull());
+  EXPECT_THAT(frame.comprehension_slots().Get(1), IsNull());
+}
+
 }  // namespace
 }  // namespace google::api::expr::runtime