From 938396120fab24f10a345d6c08b0dc976b44ef19 Mon Sep 17 00:00:00 2001 From: Jonathan Tatum Date: Wed, 15 Apr 2026 11:41:20 -0700 Subject: [PATCH] Refactor select step impl to behave consistently between iterative and recursive programs. PiperOrigin-RevId: 900271414 --- eval/compiler/flat_expr_builder.cc | 8 +- eval/eval/select_step.cc | 364 ++++++++++------------------- eval/eval/select_step.h | 3 +- eval/eval/select_step_test.cc | 25 +- 4 files changed, 149 insertions(+), 251 deletions(-) diff --git a/eval/compiler/flat_expr_builder.cc b/eval/compiler/flat_expr_builder.cc index fc7289bd2..d1435b4a3 100644 --- a/eval/compiler/flat_expr_builder.cc +++ b/eval/compiler/flat_expr_builder.cc @@ -976,6 +976,7 @@ class FlatExprVisitor : public cel::AstVisitor { return; } + StringValue field = cel::StringValue(select_expr.field()); if (auto depth = RecursionEligible(); depth.has_value()) { auto deps = ExtractRecursiveDependencies(); if (deps.size() != 1) { @@ -983,7 +984,6 @@ class FlatExprVisitor : public cel::AstVisitor { "unexpected number of dependencies for select operation.")); return; } - StringValue field = cel::StringValue(select_expr.field()); SetRecursiveStep( CreateDirectSelectStep(std::move(deps[0]), std::move(field), @@ -994,9 +994,9 @@ class FlatExprVisitor : public cel::AstVisitor { return; } - AddStep(CreateSelectStep(select_expr, expr.id(), - options_.enable_empty_wrapper_null_unboxing, - enable_optional_types_)); + AddStep(CreateSelectStep( + std::move(field), select_expr.test_only(), expr.id(), + options_.enable_empty_wrapper_null_unboxing, enable_optional_types_)); } // Call node handler group. diff --git a/eval/eval/select_step.cc b/eval/eval/select_step.cc index b815f5d87..806e20b31 100644 --- a/eval/eval/select_step.cc +++ b/eval/eval/select_step.cc @@ -5,14 +5,12 @@ #include #include -#include "absl/base/nullability.h" #include "absl/log/absl_check.h" #include "absl/log/absl_log.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/string_view.h" #include "absl/types/optional.h" -#include "common/expr.h" #include "common/value.h" #include "common/value_kind.h" #include "eval/eval/attribute_trail.h" @@ -32,7 +30,6 @@ namespace { using ::cel::BoolValue; using ::cel::ErrorValue; using ::cel::MapValue; -using ::cel::NullValue; using ::cel::OptionalValue; using ::cel::ProtoWrapperTypeOptions; using ::cel::StringValue; @@ -75,35 +72,97 @@ absl::optional CheckForMarkedAttributes(const AttributeTrail& trail, return std::nullopt; } -void TestOnlySelect(const StructValue& msg, const std::string& field, - const google::protobuf::DescriptorPool* absl_nonnull descriptor_pool, - google::protobuf::MessageFactory* absl_nonnull message_factory, - google::protobuf::Arena* absl_nonnull arena, - Value* absl_nonnull result) { - absl::StatusOr has_field = msg.HasFieldByName(field); +absl::Status PerformHas(const Value& target, absl::string_view field, + const StringValue& field_value, + const google::protobuf::DescriptorPool* descriptor_pool, + google::protobuf::MessageFactory* message_factory, + google::protobuf::Arena* arena, Value& result) { + switch (target.kind()) { + case ValueKind::kMap: { + CEL_RETURN_IF_ERROR(target.GetMap().Has(field_value, descriptor_pool, + message_factory, arena, &result)); + return absl::OkStatus(); + } + case ValueKind::kStruct: { + auto has_field = target.GetStruct().HasFieldByName(field); + if (!has_field.ok()) { + result = ErrorValue(std::move(has_field).status()); + } else { + result = BoolValue{*has_field}; + } + return absl::OkStatus(); + } + default: + return InvalidSelectTargetError(); + } +} - if (!has_field.ok()) { - *result = ErrorValue(std::move(has_field).status()); - return; +absl::Status PerformGet(const Value& target, absl::string_view field, + const StringValue& field_value, + ProtoWrapperTypeOptions unboxing_option, + const google::protobuf::DescriptorPool* descriptor_pool, + google::protobuf::MessageFactory* message_factory, + google::protobuf::Arena* arena, Value& result) { + switch (target.kind()) { + case ValueKind::kMap: { + auto status = target.GetMap().Get(field_value, descriptor_pool, + message_factory, arena, &result); + if (!status.ok()) { + result = ErrorValue(std::move(status)); + } + return absl::OkStatus(); + } + case ValueKind::kStruct: { + auto status = target.GetStruct().GetFieldByName( + field, unboxing_option, descriptor_pool, message_factory, arena, + &result); + if (!status.ok()) { + result = ErrorValue(std::move(status)); + } + return absl::OkStatus(); + } + default: + return InvalidSelectTargetError(); } - *result = BoolValue{*has_field}; } -void TestOnlySelect(const MapValue& map, const StringValue& field_name, - const google::protobuf::DescriptorPool* absl_nonnull descriptor_pool, - google::protobuf::MessageFactory* absl_nonnull message_factory, - google::protobuf::Arena* absl_nonnull arena, - Value* absl_nonnull result) { - // Field presence only supports string keys containing valid identifier - // characters. - absl::Status presence = - map.Has(field_name, descriptor_pool, message_factory, arena, result); - - if (!presence.ok()) { - *result = ErrorValue(std::move(presence)); - return; +absl::Status PerformOptionalGet(const Value& target, absl::string_view field, + const StringValue& field_value, + ProtoWrapperTypeOptions unboxing_option, + const google::protobuf::DescriptorPool* descriptor_pool, + google::protobuf::MessageFactory* message_factory, + google::protobuf::Arena* arena, Value& result) { + switch (target.kind()) { + case ValueKind::kMap: { + CEL_ASSIGN_OR_RETURN( + bool found, target.GetMap().Find(field_value, descriptor_pool, + message_factory, arena, &result)); + if (!found) { + result = OptionalValue::None(); + return absl::OkStatus(); + } + ABSL_DCHECK(!result.IsUnknown()); + result = OptionalValue::Of(std::move(result), arena); + return absl::OkStatus(); + } + case ValueKind::kStruct: { + CEL_ASSIGN_OR_RETURN(bool found, + target.GetStruct().HasFieldByName(field)); + if (!found) { + result = OptionalValue::None(); + return absl::OkStatus(); + } + CEL_RETURN_IF_ERROR(target.GetStruct().GetFieldByName( + field, unboxing_option, descriptor_pool, message_factory, arena, + &result)); + + ABSL_DCHECK(!result.IsUnknown()); + result = OptionalValue::Of(std::move(result), arena); + return absl::OkStatus(); + } + default: + return InvalidSelectTargetError(); } - ABSL_DCHECK(!result->IsUnknown()); } // SelectStep performs message field access specified by Expr::Select @@ -124,11 +183,6 @@ class SelectStep : public ExpressionStepBase { absl::Status Evaluate(ExecutionFrame* frame) const override; private: - absl::Status PerformTestOnlySelect(ExecutionFrame* frame, - const Value& arg) const; - absl::StatusOr PerformSelect(ExecutionFrame* frame, const Value& arg, - Value& result) const; - cel::StringValue field_value_; std::string field_; bool test_field_presence_; @@ -177,128 +231,48 @@ absl::Status SelectStep::Evaluate(ExecutionFrame* frame) const { return absl::OkStatus(); } - // Handle test only Select. + Value result; if (test_field_presence_) { + const Value* target = &arg; if (optional_arg) { if (!optional_arg->HasValue()) { - frame->value_stack().PopAndPush(cel::BoolValue{false}); + frame->value_stack().PopAndPush(cel::BoolValue{false}, + std::move(result_trail)); return absl::OkStatus(); } - Value value; - optional_arg->Value(&value); - return PerformTestOnlySelect(frame, value); + optional_arg->Value(&result); + target = &result; } - return PerformTestOnlySelect(frame, arg); + CEL_RETURN_IF_ERROR( + PerformHas(*target, field_, field_value_, frame->descriptor_pool(), + frame->message_factory(), frame->arena(), result)); + frame->value_stack().PopAndPush(std::move(result), std::move(result_trail)); + return absl::OkStatus(); } - // Normal select path. - // Select steps can be applied to either maps or messages if (optional_arg) { if (!optional_arg->HasValue()) { - // Leave optional_arg at the top of the stack. Its empty. + frame->value_stack().PopAndPush(OptionalValue::None(), + std::move(result_trail)); return absl::OkStatus(); } Value value; - Value result; - bool ok; optional_arg->Value(&value); - CEL_ASSIGN_OR_RETURN(ok, PerformSelect(frame, value, result)); - if (!ok) { - frame->value_stack().PopAndPush(cel::OptionalValue::None(), - std::move(result_trail)); - return absl::OkStatus(); + auto status = PerformOptionalGet( + value, field_, field_value_, unboxing_option_, frame->descriptor_pool(), + frame->message_factory(), frame->arena(), result); + if (!status.ok()) { + result = ErrorValue(std::move(status)); } - frame->value_stack().PopAndPush( - cel::OptionalValue::Of(std::move(result), frame->arena()), - std::move(result_trail)); + frame->value_stack().PopAndPush(std::move(result), std::move(result_trail)); return absl::OkStatus(); } - // Normal select path. - // Select steps can be applied to either maps or messages - switch (arg.kind()) { - case ValueKind::kStruct: { - Value result; - auto status = arg.GetStruct().GetFieldByName( - field_, unboxing_option_, frame->descriptor_pool(), - frame->message_factory(), frame->arena(), &result); - if (!status.ok()) { - result = ErrorValue(std::move(status)); - } - frame->value_stack().PopAndPush(std::move(result), - std::move(result_trail)); - return absl::OkStatus(); - } - case ValueKind::kMap: { - Value result; - auto status = - arg.GetMap().Get(field_value_, frame->descriptor_pool(), - frame->message_factory(), frame->arena(), &result); - if (!status.ok()) { - result = ErrorValue(std::move(status)); - } - frame->value_stack().PopAndPush(std::move(result), - std::move(result_trail)); - return absl::OkStatus(); - } - default: - // Control flow should have returned earlier. - return InvalidSelectTargetError(); - } -} - -absl::Status SelectStep::PerformTestOnlySelect(ExecutionFrame* frame, - const Value& arg) const { - switch (arg.kind()) { - case ValueKind::kMap: { - Value result; - TestOnlySelect(arg.GetMap(), field_value_, frame->descriptor_pool(), - frame->message_factory(), frame->arena(), &result); - frame->value_stack().PopAndPush(std::move(result)); - return absl::OkStatus(); - } - case ValueKind::kMessage: { - Value result; - TestOnlySelect(arg.GetStruct(), field_, frame->descriptor_pool(), - frame->message_factory(), frame->arena(), &result); - frame->value_stack().PopAndPush(std::move(result)); - return absl::OkStatus(); - } - default: - // Control flow should have returned earlier. - return InvalidSelectTargetError(); - } -} - -absl::StatusOr SelectStep::PerformSelect(ExecutionFrame* frame, - const Value& arg, - Value& result) const { - switch (arg->kind()) { - case ValueKind::kStruct: { - const auto& struct_value = arg.GetStruct(); - CEL_ASSIGN_OR_RETURN(auto ok, struct_value.HasFieldByName(field_)); - if (!ok) { - result = NullValue{}; - return false; - } - CEL_RETURN_IF_ERROR(struct_value.GetFieldByName( - field_, unboxing_option_, frame->descriptor_pool(), - frame->message_factory(), frame->arena(), &result)); - ABSL_DCHECK(!result.IsUnknown()); - return true; - } - case ValueKind::kMap: { - CEL_ASSIGN_OR_RETURN( - auto found, - arg.GetMap().Find(field_value_, frame->descriptor_pool(), - frame->message_factory(), frame->arena(), &result)); - ABSL_DCHECK(!found || !result.IsUnknown()); - return found; - } - default: - // Control flow should have returned earlier. - return InvalidSelectTargetError(); - } + CEL_RETURN_IF_ERROR(PerformGet( + arg, field_, field_value_, unboxing_option_, frame->descriptor_pool(), + frame->message_factory(), frame->arena(), result)); + frame->value_stack().PopAndPush(std::move(result), std::move(result_trail)); + return absl::OkStatus(); } class DirectSelectStep : public DirectExpressionStep { @@ -362,11 +336,11 @@ class DirectSelectStep : public DirectExpressionStep { } Value value; optional_arg->Value(&value); - PerformTestOnlySelect(frame, value, result); - return absl::OkStatus(); + return PerformHas(value, field_, field_value_, frame.descriptor_pool(), + frame.message_factory(), frame.arena(), result); } - PerformTestOnlySelect(frame, result, result); - return absl::OkStatus(); + return PerformHas(result, field_, field_value_, frame.descriptor_pool(), + frame.message_factory(), frame.arena(), result); } if (optional_arg) { @@ -376,26 +350,24 @@ class DirectSelectStep : public DirectExpressionStep { } Value value; optional_arg->Value(&value); - return PerformOptionalSelect(frame, value, result); + auto status = + PerformOptionalGet(value, field_, field_value_, unboxing_option_, + frame.descriptor_pool(), frame.message_factory(), + frame.arena(), result); + if (!status.ok()) { + result = ErrorValue(std::move(status)); + } + return absl::OkStatus(); } - auto status = PerformSelect(frame, result, result); - if (!status.ok()) { - result = ErrorValue(std::move(status)); - } - return absl::OkStatus(); + return PerformGet(result, field_, field_value_, unboxing_option_, + frame.descriptor_pool(), frame.message_factory(), + frame.arena(), result); } private: std::unique_ptr operand_; - void PerformTestOnlySelect(ExecutionFrameBase& frame, const Value& value, - Value& result) const; - absl::Status PerformOptionalSelect(ExecutionFrameBase& frame, - const Value& value, Value& result) const; - absl::Status PerformSelect(ExecutionFrameBase& frame, const Value& value, - Value& result) const; - // Field name in formats supported by each of the map and struct field access // APIs. // @@ -410,84 +382,6 @@ class DirectSelectStep : public DirectExpressionStep { bool enable_optional_types_; }; -void DirectSelectStep::PerformTestOnlySelect(ExecutionFrameBase& frame, - const cel::Value& value, - Value& result) const { - switch (value.kind()) { - case ValueKind::kMap: - TestOnlySelect(value.GetMap(), field_value_, frame.descriptor_pool(), - frame.message_factory(), frame.arena(), &result); - return; - case ValueKind::kMessage: - TestOnlySelect(value.GetStruct(), field_, frame.descriptor_pool(), - frame.message_factory(), frame.arena(), &result); - return; - default: - // Control flow should have returned earlier. - result = cel::ErrorValue(InvalidSelectTargetError()); - return; - } -} - -absl::Status DirectSelectStep::PerformOptionalSelect(ExecutionFrameBase& frame, - const Value& value, - Value& result) const { - switch (value.kind()) { - case ValueKind::kStruct: { - auto struct_value = value.GetStruct(); - CEL_ASSIGN_OR_RETURN(auto ok, struct_value.HasFieldByName(field_)); - if (!ok) { - result = OptionalValue::None(); - return absl::OkStatus(); - } - CEL_RETURN_IF_ERROR(struct_value.GetFieldByName( - field_, unboxing_option_, frame.descriptor_pool(), - frame.message_factory(), frame.arena(), &result)); - ABSL_DCHECK(!result.IsUnknown()); - result = OptionalValue::Of(std::move(result), frame.arena()); - return absl::OkStatus(); - } - case ValueKind::kMap: { - CEL_ASSIGN_OR_RETURN( - auto found, - value.GetMap().Find(field_value_, frame.descriptor_pool(), - frame.message_factory(), frame.arena(), &result)); - if (!found) { - result = OptionalValue::None(); - return absl::OkStatus(); - } - ABSL_DCHECK(!result.IsUnknown()); - result = OptionalValue::Of(std::move(result), frame.arena()); - return absl::OkStatus(); - } - default: - // Control flow should have returned earlier. - return InvalidSelectTargetError(); - } -} - -absl::Status DirectSelectStep::PerformSelect(ExecutionFrameBase& frame, - const cel::Value& value, - Value& result) const { - switch (value.kind()) { - case ValueKind::kStruct: - CEL_RETURN_IF_ERROR(value.GetStruct().GetFieldByName( - field_, unboxing_option_, frame.descriptor_pool(), - frame.message_factory(), frame.arena(), &result)); - ABSL_DCHECK(!result.IsUnknown()); - return absl::OkStatus(); - case ValueKind::kMap: - CEL_RETURN_IF_ERROR( - value.GetMap().Get(field_value_, frame.descriptor_pool(), - frame.message_factory(), frame.arena(), &result)); - ABSL_DCHECK(!result.IsUnknown()); - return absl::OkStatus(); - default: - // Control flow should have returned earlier. - return InvalidSelectTargetError(); - } -} - } // namespace std::unique_ptr CreateDirectSelectStep( @@ -501,11 +395,11 @@ std::unique_ptr CreateDirectSelectStep( // Factory method for Select - based Execution step absl::StatusOr> CreateSelectStep( - const cel::SelectExpr& select_expr, int64_t expr_id, + cel::StringValue field, bool test_only, int64_t expr_id, bool enable_wrapper_type_null_unboxing, bool enable_optional_types) { - return std::make_unique( - cel::StringValue(select_expr.field()), select_expr.test_only(), expr_id, - enable_wrapper_type_null_unboxing, enable_optional_types); + return std::make_unique(std::move(field), test_only, expr_id, + enable_wrapper_type_null_unboxing, + enable_optional_types); } } // namespace google::api::expr::runtime diff --git a/eval/eval/select_step.h b/eval/eval/select_step.h index 6eaaf9487..2fe35ecb0 100644 --- a/eval/eval/select_step.h +++ b/eval/eval/select_step.h @@ -5,7 +5,6 @@ #include #include "absl/status/statusor.h" -#include "common/expr.h" #include "common/value.h" #include "eval/eval/direct_expression_step.h" #include "eval/eval/evaluator_core.h" @@ -20,7 +19,7 @@ std::unique_ptr CreateDirectSelectStep( // Factory method for Select - based Execution step absl::StatusOr> CreateSelectStep( - const cel::SelectExpr& select_expr, int64_t expr_id, + cel::StringValue field, bool test_only, int64_t expr_id, bool enable_wrapper_type_null_unboxing, bool enable_optional_types = false); } // namespace google::api::expr::runtime diff --git a/eval/eval/select_step_test.cc b/eval/eval/select_step_test.cc index ce532eabd..d5ca2aecd 100644 --- a/eval/eval/select_step_test.cc +++ b/eval/eval/select_step_test.cc @@ -133,8 +133,8 @@ class SelectStepTest : public testing::Test { CEL_ASSIGN_OR_RETURN(auto step0, CreateIdentStep(ident.name(), expr0.id())); CEL_ASSIGN_OR_RETURN( auto step1, - CreateSelectStep(select, expr.id(), - options.enable_wrapper_type_null_unboxing)); + CreateSelectStep(cel::StringValue(select.field()), select.test_only(), + expr.id(), options.enable_wrapper_type_null_unboxing)); path.push_back(std::move(step0)); path.push_back(std::move(step1)); @@ -330,11 +330,13 @@ TEST_F(SelectStepTest, MapPresenseIsErrorTest) { ASSERT_OK_AND_ASSIGN(auto step0, CreateIdentStep(ident.name(), expr0.id())); ASSERT_OK_AND_ASSIGN( auto step1, - CreateSelectStep(select_map, expr1.id(), + CreateSelectStep(cel::StringValue(select_map.field()), + select_map.test_only(), expr1.id(), /*enable_wrapper_type_null_unboxing=*/false)); ASSERT_OK_AND_ASSIGN( auto step2, - CreateSelectStep(select, select_expr.id(), + CreateSelectStep(cel::StringValue(select.field()), select.test_only(), + select_expr.id(), /*enable_wrapper_type_null_unboxing=*/false)); ExecutionPath path; @@ -836,7 +838,8 @@ TEST_P(SelectStepConformanceTest, CelErrorAsArgument) { ASSERT_OK_AND_ASSIGN(auto step0, CreateIdentStep(ident.name(), expr0.id())); ASSERT_OK_AND_ASSIGN( auto step1, - CreateSelectStep(select, dummy_expr.id(), + CreateSelectStep(cel::StringValue(select.field()), select.test_only(), + dummy_expr.id(), /*enable_wrapper_type_null_unboxing=*/false)); path.push_back(std::move(step0)); @@ -877,7 +880,8 @@ TEST_F(SelectStepTest, DisableMissingAttributeOK) { ASSERT_OK_AND_ASSIGN(auto step0, CreateIdentStep(ident.name(), expr0.id())); ASSERT_OK_AND_ASSIGN( auto step1, - CreateSelectStep(select, dummy_expr.id(), + CreateSelectStep(cel::StringValue(select.field()), select.test_only(), + dummy_expr.id(), /*enable_wrapper_type_null_unboxing=*/false)); path.push_back(std::move(step0)); @@ -919,7 +923,8 @@ TEST_F(SelectStepTest, UnrecoverableUnknownValueProducesError) { ASSERT_OK_AND_ASSIGN(auto step0, CreateIdentStep(ident.name(), expr0.id())); ASSERT_OK_AND_ASSIGN( auto step1, - CreateSelectStep(select, dummy_expr.id(), + CreateSelectStep(cel::StringValue(select.field()), select.test_only(), + dummy_expr.id(), /*enable_wrapper_type_null_unboxing=*/false)); path.push_back(std::move(step0)); @@ -965,9 +970,9 @@ TEST_F(SelectStepTest, UnknownPatternResolvesToUnknown) { auto& ident = expr0.mutable_ident_expr(); ident.set_name("message"); auto step0_status = CreateIdentStep(ident.name(), expr0.id()); - auto step1_status = - CreateSelectStep(select, dummy_expr.id(), - /*enable_wrapper_type_null_unboxing=*/false); + auto step1_status = CreateSelectStep( + cel::StringValue(select.field()), select.test_only(), dummy_expr.id(), + /*enable_wrapper_type_null_unboxing=*/false); ASSERT_THAT(step0_status, IsOk()); ASSERT_THAT(step1_status, IsOk());