From 07aef5f7a3538ebe88c955cac46564ed641f3e0c Mon Sep 17 00:00:00 2001 From: Jonathan Tatum Date: Tue, 5 Sep 2023 13:09:57 -0700 Subject: [PATCH] Simplify value stack operations for comprehensions. PiperOrigin-RevId: 562868576 --- eval/compiler/flat_expr_builder.cc | 26 +-- eval/eval/comprehension_step.cc | 269 +++++++++++++++------------ eval/eval/comprehension_step.h | 22 +-- eval/eval/comprehension_step_test.cc | 25 ++- 4 files changed, 182 insertions(+), 160 deletions(-) diff --git a/eval/compiler/flat_expr_builder.cc b/eval/compiler/flat_expr_builder.cc index 28e796208..5fb9592ca 100644 --- a/eval/compiler/flat_expr_builder.cc +++ b/eval/compiler/flat_expr_builder.cc @@ -1192,30 +1192,16 @@ int ComprehensionAccumulationReferences(const cel::ast_internal::Expr& expr, return absl::visit(handler, expr.expr_kind()); } -void ComprehensionVisitor::PreVisit(const cel::ast_internal::Expr*) { - constexpr int64_t kLoopStepPlaceholder = -10; - visitor_->AddStep(CreateConstValueStep( - visitor_->value_factory().CreateIntValue(kLoopStepPlaceholder), - kExprIdNotFromAst, false)); -} +void ComprehensionVisitor::PreVisit(const cel::ast_internal::Expr*) {} void ComprehensionVisitor::PostVisitArg( cel::ast_internal::ComprehensionArg arg_num, const cel::ast_internal::Expr* expr) { - // TODO(issues/20): Consider refactoring the comprehension prologue step. switch (arg_num) { case cel::ast_internal::ITER_RANGE: { - // Post-process iter_range to list its keys if it's a map. - visitor_->AddStep(CreateListKeysStep(expr->id())); - // Setup index stack position - visitor_->AddStep( - CreateConstValueStep(visitor_->value_factory().CreateIntValue(-1), - kExprIdNotFromAst, false)); - // Element at index. - constexpr int64_t kCurrentValuePlaceholder = -20; - visitor_->AddStep(CreateConstValueStep( - visitor_->value_factory().CreateIntValue(kCurrentValuePlaceholder), - kExprIdNotFromAst, false)); + // post process iter_range to list its keys if it's a map + // and initialize the loop index. + visitor_->AddStep(CreateComprehensionInitStep(expr->id())); break; } case cel::ast_internal::ACCU_INIT: { @@ -1245,8 +1231,8 @@ void ComprehensionVisitor::PostVisitArg( break; } case cel::ast_internal::RESULT: { - visitor_->AddStep(std::unique_ptr( - new ComprehensionFinish(slot_offset_, expr->id()))); + visitor_->AddStep( + CreateComprehensionFinishStep(slot_offset_, expr->id())); next_step_->set_error_jump_offset(visitor_->GetCurrentIndex() - next_step_pos_ - 1); cond_step_->set_error_jump_offset(visitor_->GetCurrentIndex() - diff --git a/eval/eval/comprehension_step.cc b/eval/eval/comprehension_step.cc index 0a060c8d5..916ce7a2a 100644 --- a/eval/eval/comprehension_step.cc +++ b/eval/eval/comprehension_step.cc @@ -9,6 +9,7 @@ #include "absl/status/status.h" #include "absl/strings/str_cat.h" #include "base/kind.h" +#include "base/values/error_value.h" #include "base/values/unknown_value.h" #include "eval/eval/attribute_trail.h" #include "eval/eval/comprehension_slots.h" @@ -18,6 +19,7 @@ #include "runtime/internal/mutable_list_impl.h" namespace google::api::expr::runtime { +namespace { using ::cel::Handle; using ::cel::UnknownValue; @@ -25,26 +27,119 @@ using ::cel::Value; using ::cel::runtime_internal::CreateNoMatchingOverloadError; using ::cel::runtime_internal::MutableListValue; +class ComprehensionFinish : public ExpressionStepBase { + public: + ComprehensionFinish(size_t slot_offset, int64_t expr_id); + + absl::Status Evaluate(ExecutionFrame* frame) const override; + + private: + size_t accu_slot_; +}; + +ComprehensionFinish::ComprehensionFinish(size_t slot_offset_, int64_t expr_id) + : ExpressionStepBase(expr_id), + accu_slot_(slot_offset_ + kComprehensionSlotsAccuOffset) {} + +// Stack changes of ComprehensionFinish. +// +// Stack size before: 3. +// Stack size after: 1. +absl::Status ComprehensionFinish::Evaluate(ExecutionFrame* frame) const { + if (!frame->value_stack().HasEnough(3)) { + return absl::Status(absl::StatusCode::kInternal, "Value stack underflow"); + } + Handle result = frame->value_stack().Peek(); + frame->value_stack().Pop(3); + if (frame->enable_comprehension_list_append() && + result->Is()) { + // We assume this is 'owned' by the evaluator stack so const cast is safe + // here. + // Convert the buildable list to an actual cel::ListValue. + MutableListValue& list_value = + const_cast(result->As()); + CEL_ASSIGN_OR_RETURN(result, std::move(list_value).Build()); + } + frame->value_stack().Push(std::move(result)); + frame->comprehension_slots().ClearSlot(accu_slot_); + return absl::OkStatus(); +} + +class ComprehensionInitStep : public ExpressionStepBase { + public: + explicit ComprehensionInitStep(int64_t expr_id) + : ExpressionStepBase(expr_id, false) {} + absl::Status Evaluate(ExecutionFrame* frame) const override; + + private: + absl::Status ProjectKeys(ExecutionFrame* frame) const; +}; + +absl::Status ComprehensionInitStep::ProjectKeys(ExecutionFrame* frame) const { + // Top of stack is map, but could be partially unknown. To tolerate cases when + // keys are not set for declared unknown values, convert to an unknown set. + if (frame->enable_unknowns()) { + absl::optional> unknown = + frame->attribute_utility().IdentifyAndMergeUnknowns( + frame->value_stack().GetSpan(1), + frame->value_stack().GetAttributeSpan(1), + /*use_partial=*/true); + if (unknown.has_value()) { + frame->value_stack().PopAndPush(*std::move(unknown)); + return absl::OkStatus(); + } + } + + CEL_ASSIGN_OR_RETURN( + auto list_keys, frame->value_stack().Peek().As()->ListKeys( + frame->value_factory())); + frame->value_stack().PopAndPush(std::move(list_keys)); + return absl::OkStatus(); +} + +// Setup the value stack for comprehension. +// Coerce the top of stack into a list and initilialize an index. +// This should happen after evaluating the iter_range part of the comprehension. +absl::Status ComprehensionInitStep::Evaluate(ExecutionFrame* frame) const { + if (!frame->value_stack().HasEnough(1)) { + return absl::Status(absl::StatusCode::kInternal, "Value stack underflow"); + } + if (frame->value_stack().Peek()->Is()) { + CEL_RETURN_IF_ERROR(ProjectKeys(frame)); + } + + const auto& range = frame->value_stack().Peek(); + if (!range->Is() && !range->Is() && + !range->Is()) { + frame->value_stack().PopAndPush(frame->value_factory().CreateErrorValue( + CreateNoMatchingOverloadError(""))); + } + + // Initialize current index. + // Error handling for wrong range type is deferred until the 'Next' step + // to simplify the number of jumps. + frame->value_stack().Push(frame->value_factory().CreateIntValue(-1)); + return absl::OkStatus(); +} + +} // namespace + // Stack variables during comprehension evaluation: -// 0. accu_init, then loop_step (any), available through accu_var -// 1. iter_range (list) -// 2. current index in iter_range (int64_t) -// 3. current_value from iter_range (any), available through iter_var -// 4. loop_condition (bool) OR loop_step (any) - -// What to put on ExecutionPath: stack size -// 0. (dummy) 1 -// 1. iter_range (dep) 2 -// 2. -1 3 -// 3. (dummy) 4 -// 4. accu_init (dep) 5 -// 5. ComprehensionNextStep 4 -// 6. loop_condition (dep) 5 -// 7. ComprehensionCondStep 4 -// 8. loop_step (dep) 5 -// 9. goto 5. 5 -// 10. result (dep) 2 -// 11. ComprehensionFinish 1 +// 0. iter_range (list) +// 1. current index in iter_range (int64_t) +// 2. current accumulator value or break condition + +// instruction stack size +// 0. iter_range (dep) 0 -> 1 +// 1. ComprehensionInit 1 -> 2 +// 2. accu_init (dep) 2 -> 3 +// 3. ComprehensionNextStep 3 -> 2 +// 4. loop_condition (dep) 2 -> 3 +// 5. ComprehensionCondStep 3 -> 2 +// 6. loop_step (dep) 2 -> 3 +// 7. goto 3. 3 -> 3 +// 8. result (dep) 2 -> 3 +// 9. ComprehensionFinish 3 -> 1 ComprehensionNextStep::ComprehensionNextStep(size_t slot_offset_, int64_t expr_id) @@ -63,20 +158,13 @@ void ComprehensionNextStep::set_error_jump_offset(int offset) { // Stack changes of ComprehensionNextStep. // // Stack before: -// 0. previous accu_init or "" on the first iteration -// 1. iter_range (list) -// 2. old current_index in iter_range (int64_t) -// 3. old current_value or "" on the first iteration -// 4. loop_step or accu_init (any) +// 0. iter_range (list) +// 1. old current_index in iter_range (int64_t) +// 2. loop_step or accu_init (any) // // Stack after: -// 0. loop_step or accu_init (any) -// 1. iter_range (list) -// 2. new current_index in iter_range (int64_t) -// 3. new current_value -// -// Stack on break: -// 0. loop_step or accu_init (any) +// 0. iter_range (list) +// 1. new current_index in iter_range (int64_t) // // When iter_range is not a list, this step jumps to error_jump_offset_ that is // controlled by set_error_jump_offset. In that case the stack is cleared @@ -86,28 +174,27 @@ void ComprehensionNextStep::set_error_jump_offset(int offset) { // 0. error absl::Status ComprehensionNextStep::Evaluate(ExecutionFrame* frame) const { enum { - POS_PREVIOUS_LOOP_STEP, POS_ITER_RANGE, POS_CURRENT_INDEX, - POS_CURRENT_VALUE, - POS_LOOP_STEP, + POS_LOOP_STEP_ACCU, }; - if (!frame->value_stack().HasEnough(5)) { + constexpr int kStackSize = 3; + if (!frame->value_stack().HasEnough(kStackSize)) { return absl::Status(absl::StatusCode::kInternal, "Value stack underflow"); } - auto state = frame->value_stack().GetSpan(5); + auto state = frame->value_stack().GetSpan(kStackSize); // Get range from the stack. auto iter_range = state[POS_ITER_RANGE]; if (!iter_range->Is()) { - frame->value_stack().Pop(5); + frame->value_stack().Pop(kStackSize); if (iter_range->Is() || iter_range->Is()) { frame->value_stack().Push(std::move(iter_range)); - return frame->JumpTo(error_jump_offset_); + } else { + frame->value_stack().Push(frame->value_factory().CreateErrorValue( + CreateNoMatchingOverloadError(""))); } - frame->value_stack().Push(frame->value_factory().CreateErrorValue( - CreateNoMatchingOverloadError(""))); return frame->JumpTo(error_jump_offset_); } @@ -125,16 +212,15 @@ absl::Status ComprehensionNextStep::Evaluate(ExecutionFrame* frame) const { AttributeTrail iter_range_attr; AttributeTrail iter_trail; if (frame->enable_unknowns()) { - auto attr = frame->value_stack().GetAttributeSpan(5); + auto attr = frame->value_stack().GetAttributeSpan(kStackSize); iter_range_attr = attr[POS_ITER_RANGE]; iter_trail = iter_range_attr.Step(cel::AttributeQualifier::OfInt(current_index + 1)); } // Pop invalidates references to the stack on the following line so copy. - Handle loop_step = state[POS_LOOP_STEP]; - frame->value_stack().Pop(5); - frame->value_stack().Push(loop_step); + Handle loop_step = state[POS_LOOP_STEP_ACCU]; + frame->value_stack().Pop(1); frame->comprehension_slots().Set(accu_slot_, std::move(loop_step)); // Make sure the iter var is out of scope. @@ -143,16 +229,15 @@ absl::Status ComprehensionNextStep::Evaluate(ExecutionFrame* frame) const { frame->comprehension_slots().ClearSlot(iter_slot_); return frame->JumpTo(jump_offset_); } - frame->value_stack().Push(iter_range, std::move(iter_range_attr)); + current_index += 1; CEL_ASSIGN_OR_RETURN( auto current_value, iter_range.As()->Get(frame->value_factory(), static_cast(current_index))); - frame->value_stack().Push( + frame->value_stack().PopAndPush( frame->value_factory().CreateIntValue(current_index)); - frame->value_stack().Push(current_value, iter_trail); frame->comprehension_slots().Set(iter_slot_, std::move(current_value), std::move(iter_trail)); return absl::OkStatus(); @@ -174,18 +259,23 @@ void ComprehensionCondStep::set_error_jump_offset(int offset) { error_jump_offset_ = offset; } +// Check the break condition for the comprehension. +// +// If the condition is false jump to the `result` subexpression. +// If not a bool, clear stack and jump past the result expression. +// Otherwise, continue to the accumulate step. // Stack changes by ComprehensionCondStep. // -// Stack size before: 5. -// Stack size after: 4. -// Stack size on break: 1. +// Stack size before: 3. +// Stack size after: 2. +// Stack size on error: 1. absl::Status ComprehensionCondStep::Evaluate(ExecutionFrame* frame) const { - if (!frame->value_stack().HasEnough(5)) { + if (!frame->value_stack().HasEnough(3)) { return absl::Status(absl::StatusCode::kInternal, "Value stack underflow"); } auto loop_condition_value = frame->value_stack().Peek(); if (!loop_condition_value->Is()) { - frame->value_stack().Pop(5); + frame->value_stack().Pop(3); if (loop_condition_value->Is() || loop_condition_value->Is()) { frame->value_stack().Push(std::move(loop_condition_value)); @@ -202,83 +292,18 @@ absl::Status ComprehensionCondStep::Evaluate(ExecutionFrame* frame) const { bool loop_condition = loop_condition_value.As()->value(); frame->value_stack().Pop(1); // loop_condition if (!loop_condition && shortcircuiting_) { - frame->value_stack().Pop(3); // current_value, current_index, iter_range return frame->JumpTo(jump_offset_); } return absl::OkStatus(); } -ComprehensionFinish::ComprehensionFinish(size_t slot_offset_, int64_t expr_id) - : ExpressionStepBase(expr_id), - accu_slot_(slot_offset_ + kComprehensionSlotsAccuOffset) {} - -// Stack changes of ComprehensionFinish. -// -// Stack size before: 2. -// Stack size after: 1. -absl::Status ComprehensionFinish::Evaluate(ExecutionFrame* frame) const { - if (!frame->value_stack().HasEnough(2)) { - return absl::Status(absl::StatusCode::kInternal, "Value stack underflow"); - } - Handle result = frame->value_stack().Peek(); - frame->value_stack().Pop(2); - if (frame->enable_comprehension_list_append() && - result->Is()) { - // We assume this is 'owned' by the evaluator stack so const cast is safe - // here. - // Convert the buildable list to an actual cel::ListValue. - MutableListValue& list_value = - const_cast(result->As()); - CEL_ASSIGN_OR_RETURN(result, std::move(list_value).Build()); - } - frame->value_stack().Push(std::move(result)); - frame->comprehension_slots().ClearSlot(accu_slot_); - return absl::OkStatus(); +std::unique_ptr CreateComprehensionFinishStep( + size_t slot_offset, int64_t expr_id) { + return std::make_unique(slot_offset, expr_id); } -class ListKeysStep : public ExpressionStepBase { - public: - explicit ListKeysStep(int64_t expr_id) : ExpressionStepBase(expr_id, false) {} - absl::Status Evaluate(ExecutionFrame* frame) const override; - - private: - absl::Status ProjectKeys(ExecutionFrame* frame) const; -}; - -std::unique_ptr CreateListKeysStep(int64_t expr_id) { - return std::make_unique(expr_id); -} - -absl::Status ListKeysStep::ProjectKeys(ExecutionFrame* frame) const { - // Top of stack is map, but could be partially unknown. To tolerate cases when - // keys are not set for declared unknown values, convert to an unknown set. - if (frame->enable_unknowns()) { - absl::optional> unknown = - frame->attribute_utility().IdentifyAndMergeUnknowns( - frame->value_stack().GetSpan(1), - frame->value_stack().GetAttributeSpan(1), - /*use_partial=*/true); - if (unknown.has_value()) { - frame->value_stack().PopAndPush(*std::move(unknown)); - return absl::OkStatus(); - } - } - - CEL_ASSIGN_OR_RETURN( - auto list_keys, frame->value_stack().Peek().As()->ListKeys( - frame->value_factory())); - frame->value_stack().PopAndPush(std::move(list_keys)); - return absl::OkStatus(); -} - -absl::Status ListKeysStep::Evaluate(ExecutionFrame* frame) const { - if (!frame->value_stack().HasEnough(1)) { - return absl::Status(absl::StatusCode::kInternal, "Value stack underflow"); - } - if (frame->value_stack().Peek()->Is()) { - return ProjectKeys(frame); - } - return absl::OkStatus(); +std::unique_ptr CreateComprehensionInitStep(int64_t expr_id) { + return std::make_unique(expr_id); } } // namespace google::api::expr::runtime diff --git a/eval/eval/comprehension_step.h b/eval/eval/comprehension_step.h index cb6ceafa0..e13765c25 100644 --- a/eval/eval/comprehension_step.h +++ b/eval/eval/comprehension_step.h @@ -45,19 +45,15 @@ class ComprehensionCondStep : public ExpressionStepBase { bool shortcircuiting_; }; -class ComprehensionFinish : public ExpressionStepBase { - public: - ComprehensionFinish(size_t slot_offset, int64_t expr_id); - - absl::Status Evaluate(ExecutionFrame* frame) const override; - - private: - size_t accu_slot_; -}; - -// Creates a step that lists the map keys if the top of the stack is a map, -// otherwise it's a no-op. -std::unique_ptr CreateListKeysStep(int64_t expr_id); +// Creates a cleanup step for the comprehension. +// Removes the comprehension context then pushes the 'result' sub expression to +// the top of the stack. +std::unique_ptr CreateComprehensionFinishStep( + size_t slot_offset, int64_t expr_id); + +// Creates a step that checks that the input is iterable and setups the loop +// context for the comprehension. +std::unique_ptr CreateComprehensionInitStep(int64_t expr_id); } // namespace google::api::expr::runtime diff --git a/eval/eval/comprehension_step_test.cc b/eval/eval/comprehension_step_test.cc index 9916b2de1..978450d9c 100644 --- a/eval/eval/comprehension_step_test.cc +++ b/eval/eval/comprehension_step_test.cc @@ -58,6 +58,16 @@ class ListKeysStepTest : public testing::Test { Expr dummy_expr_; }; +class GetListKeysResultStep : public ExpressionStepBase { + public: + GetListKeysResultStep() : ExpressionStepBase(-1, false) {} + + absl::Status Evaluate(ExecutionFrame* frame) const override { + frame->value_stack().Pop(1); + return absl::OkStatus(); + } +}; + MATCHER_P(CelStringValue, val, "") { const CelValue& to_match = arg; absl::string_view value = val; @@ -70,9 +80,10 @@ TEST_F(ListKeysStepTest, ListPassedThrough) { auto result = CreateIdentStep(ident, 0); ASSERT_OK(result); path.push_back(*std::move(result)); - result = CreateListKeysStep(1); + result = CreateComprehensionInitStep(1); ASSERT_OK(result); path.push_back(*std::move(result)); + path.push_back(std::make_unique()); auto expression = MakeExpression(std::move(path)); @@ -97,9 +108,10 @@ TEST_F(ListKeysStepTest, MapToKeyList) { auto result = CreateIdentStep(ident, 0); ASSERT_OK(result); path.push_back(*std::move(result)); - result = CreateListKeysStep(1); + result = CreateComprehensionInitStep(1); ASSERT_OK(result); path.push_back(*std::move(result)); + path.push_back(std::make_unique()); auto expression = MakeExpression(std::move(path)); @@ -133,9 +145,10 @@ TEST_F(ListKeysStepTest, MapPartiallyUnknown) { auto result = CreateIdentStep(ident, 0); ASSERT_OK(result); path.push_back(*std::move(result)); - result = CreateListKeysStep(1); + result = CreateComprehensionInitStep(1); ASSERT_OK(result); path.push_back(*std::move(result)); + path.push_back(std::make_unique()); auto expression = MakeExpression(std::move(path), /*unknown_attributes=*/true); @@ -171,9 +184,10 @@ TEST_F(ListKeysStepTest, ErrorPassedThrough) { auto result = CreateIdentStep(ident, 0); ASSERT_OK(result); path.push_back(*std::move(result)); - result = CreateListKeysStep(1); + result = CreateComprehensionInitStep(1); ASSERT_OK(result); path.push_back(*std::move(result)); + path.push_back(std::make_unique()); auto expression = MakeExpression(std::move(path)); @@ -196,9 +210,10 @@ TEST_F(ListKeysStepTest, UnknownSetPassedThrough) { auto result = CreateIdentStep(ident, 0); ASSERT_OK(result); path.push_back(*std::move(result)); - result = CreateListKeysStep(1); + result = CreateComprehensionInitStep(1); ASSERT_OK(result); path.push_back(*std::move(result)); + path.push_back(std::make_unique()); auto expression = MakeExpression(std::move(path), /*unknown_attributes=*/true);