diff --git a/common/values/parsed_json_list_value.cc b/common/values/parsed_json_list_value.cc index 9418aebad..902deab33 100644 --- a/common/values/parsed_json_list_value.cc +++ b/common/values/parsed_json_list_value.cc @@ -41,6 +41,7 @@ #include "google/protobuf/descriptor.h" #include "google/protobuf/io/zero_copy_stream.h" #include "google/protobuf/message.h" +#include "google/protobuf/message_lite.h" namespace cel { @@ -204,6 +205,12 @@ size_t ParsedJsonListValue::Size() const { if (value_ == nullptr) { return 0; } + if (const auto* generated_message = + google::protobuf::DynamicCastMessage(value_); + ABSL_PREDICT_TRUE(generated_message != nullptr)) { + return static_cast( + well_known_types::ListValueReflection::ValuesSize(*generated_message)); + } return static_cast( well_known_types::GetListValueReflectionOrDie(value_->GetDescriptor()) .ValuesSize(*value_)); @@ -223,6 +230,21 @@ absl::Status ParsedJsonListValue::Get( *result = IndexOutOfBoundsError(index, arena); return absl::OkStatus(); } + if (const auto* generated_message = + google::protobuf::DynamicCastMessage(value_); + ABSL_PREDICT_TRUE(generated_message != nullptr)) { + const int size = + well_known_types::ListValueReflection::ValuesSize(*generated_message); + if (ABSL_PREDICT_FALSE(index >= static_cast(size))) { + *result = IndexOutOfBoundsError(index, arena); + return absl::OkStatus(); + } + *result = common_internal::ParsedJsonValue( + &well_known_types::ListValueReflection::Values(*generated_message, + static_cast(index)), + arena); + return absl::OkStatus(); + } const auto reflection = well_known_types::GetListValueReflectionOrDie(value_->GetDescriptor()); if (ABSL_PREDICT_FALSE(index >= @@ -248,6 +270,22 @@ absl::Status ParsedJsonListValue::ForEach( return absl::OkStatus(); } Value scratch; + if (const auto* generated_message = + google::protobuf::DynamicCastMessage(value_); + ABSL_PREDICT_TRUE(generated_message != nullptr)) { + const int size = + well_known_types::ListValueReflection::ValuesSize(*generated_message); + for (int i = 0; i < size; ++i) { + scratch = common_internal::ParsedJsonValue( + &well_known_types::ListValueReflection::Values(*generated_message, i), + arena); + CEL_ASSIGN_OR_RETURN(auto ok, callback(static_cast(i), scratch)); + if (!ok) { + break; + } + } + return absl::OkStatus(); + } const auto reflection = well_known_types::GetListValueReflectionOrDie(value_->GetDescriptor()); const int size = reflection.ValuesSize(*value_); @@ -269,9 +307,16 @@ class ParsedJsonListValueIterator final : public ValueIterator { explicit ParsedJsonListValueIterator( const google::protobuf::Message* absl_nonnull message) : message_(message), - reflection_(well_known_types::GetListValueReflectionOrDie( - message_->GetDescriptor())), - size_(reflection_.ValuesSize(*message_)) {} + generated_( + google::protobuf::DynamicCastMessage(message_)) { + if (ABSL_PREDICT_TRUE(generated_ != nullptr)) { + size_ = well_known_types::ListValueReflection::ValuesSize(*generated_); + } else { + reflection_ = well_known_types::GetListValueReflectionOrDie( + message_->GetDescriptor()); + size_ = reflection_.ValuesSize(*message_); + } + } bool HasNext() override { return index_ < size_; } @@ -289,8 +334,14 @@ class ParsedJsonListValueIterator final : public ValueIterator { "`ValueIterator::Next` called after `ValueIterator::HasNext` " "returned false"); } - *result = common_internal::ParsedJsonValue( - &reflection_.Values(*message_, index_), arena); + if (ABSL_PREDICT_TRUE(generated_ != nullptr)) { + *result = common_internal::ParsedJsonValue( + &well_known_types::ListValueReflection::Values(*generated_, index_), + arena); + } else { + *result = common_internal::ParsedJsonValue( + &reflection_.Values(*message_, index_), arena); + } ++index_; return absl::OkStatus(); } @@ -308,8 +359,14 @@ class ParsedJsonListValueIterator final : public ValueIterator { if (index_ >= size_) { return false; } - *key_or_value = common_internal::ParsedJsonValue( - &reflection_.Values(*message_, index_), arena); + if (ABSL_PREDICT_TRUE(generated_ != nullptr)) { + *key_or_value = common_internal::ParsedJsonValue( + &well_known_types::ListValueReflection::Values(*generated_, index_), + arena); + } else { + *key_or_value = common_internal::ParsedJsonValue( + &reflection_.Values(*message_, index_), arena); + } ++index_; return true; } @@ -328,8 +385,14 @@ class ParsedJsonListValueIterator final : public ValueIterator { return false; } if (value != nullptr) { - *value = common_internal::ParsedJsonValue( - &reflection_.Values(*message_, index_), arena); + if (ABSL_PREDICT_TRUE(generated_ != nullptr)) { + *value = common_internal::ParsedJsonValue( + &well_known_types::ListValueReflection::Values(*generated_, index_), + arena); + } else { + *value = common_internal::ParsedJsonValue( + &reflection_.Values(*message_, index_), arena); + } } *key = IntValue(index_); ++index_; @@ -338,8 +401,9 @@ class ParsedJsonListValueIterator final : public ValueIterator { private: const google::protobuf::Message* absl_nonnull const message_; - const well_known_types::ListValueReflection reflection_; - const int size_; + const google::protobuf::ListValue* absl_nullable const generated_; + well_known_types::ListValueReflection reflection_; + int size_; int index_ = 0; }; diff --git a/common/values/parsed_json_map_value.cc b/common/values/parsed_json_map_value.cc index 1b568e046..c1d9ed6cb 100644 --- a/common/values/parsed_json_map_value.cc +++ b/common/values/parsed_json_map_value.cc @@ -203,6 +203,12 @@ size_t ParsedJsonMapValue::Size() const { if (value_ == nullptr) { return 0; } + if (const auto* generated_message = + google::protobuf::DynamicCastMessage(value_); + ABSL_PREDICT_TRUE(generated_message != nullptr)) { + return static_cast( + well_known_types::StructReflection::FieldsSize(*generated_message)); + } return static_cast( well_known_types::GetStructReflectionOrDie(value_->GetDescriptor()) .FieldsSize(*value_)); @@ -237,6 +243,18 @@ absl::StatusOr ParsedJsonMapValue::Find( return false; } std::string key_scratch; + if (const auto* generated_message = + google::protobuf::DynamicCastMessage(value_); + ABSL_PREDICT_TRUE(generated_message != nullptr)) { + if (const auto* value = well_known_types::StructReflection::FindField( + *generated_message, string_key->ToStringView(&key_scratch)); + value != nullptr) { + *result = common_internal::ParsedJsonValue(value, arena); + return true; + } + *result = NullValue(); + return false; + } if (const auto* value = well_known_types::GetStructReflectionOrDie( value_->GetDescriptor()) @@ -269,6 +287,13 @@ absl::Status ParsedJsonMapValue::Has( return absl::OkStatus(); } std::string key_scratch; + if (const auto* generated_message = + google::protobuf::DynamicCastMessage(value_); + ABSL_PREDICT_TRUE(generated_message != nullptr)) { + *result = BoolValue(well_known_types::StructReflection::ContainsField( + *generated_message, string_key->ToStringView(&key_scratch))); + return absl::OkStatus(); + } if (const auto* value = well_known_types::GetStructReflectionOrDie( value_->GetDescriptor()) diff --git a/common/values/parsed_json_value.cc b/common/values/parsed_json_value.cc index fdcd5a1d0..44bf872d1 100644 --- a/common/values/parsed_json_value.cc +++ b/common/values/parsed_json_value.cc @@ -17,8 +17,10 @@ #include #include +#include "google/protobuf/struct.pb.h" #include "absl/base/attributes.h" #include "absl/base/nullability.h" +#include "absl/base/optimization.h" #include "absl/functional/overload.h" #include "absl/status/status.h" #include "absl/strings/cord.h" @@ -29,6 +31,7 @@ #include "internal/well_known_types.h" #include "google/protobuf/arena.h" #include "google/protobuf/message.h" +#include "google/protobuf/message_lite.h" namespace cel::common_internal { @@ -36,6 +39,7 @@ namespace { using ::cel::well_known_types::AsVariant; using ::cel::well_known_types::GetValueReflectionOrDie; +using ::cel::well_known_types::ValueReflection; google::protobuf::Arena* absl_nonnull MessageArenaOr( const google::protobuf::Message* absl_nonnull message, @@ -51,6 +55,45 @@ google::protobuf::Arena* absl_nonnull MessageArenaOr( Value ParsedJsonValue(const google::protobuf::Message* absl_nonnull message, google::protobuf::Arena* absl_nonnull arena) { + if (const auto* generated_message = + google::protobuf::DynamicCastMessage(message); + ABSL_PREDICT_TRUE(generated_message != nullptr)) { + const auto kind_case = ValueReflection::GetKindCase(*generated_message); + switch (kind_case) { + case google::protobuf::Value::KIND_NOT_SET: + ABSL_FALLTHROUGH_INTENDED; + case google::protobuf::Value::kNullValue: + return NullValue(); + case google::protobuf::Value::kBoolValue: + return BoolValue(ValueReflection::GetBoolValue(*generated_message)); + case google::protobuf::Value::kNumberValue: + return DoubleValue(ValueReflection::GetNumberValue(*generated_message)); + case google::protobuf::Value::kStringValue: { + absl::string_view string = + ValueReflection::GetStringValue(*generated_message); + if (string.empty()) { + return StringValue(); + } + if (google::protobuf::Arena* message_arena = generated_message->GetArena(); + message_arena != nullptr) { + return StringValue::Wrap(string, message_arena); + } + return StringValue::From(string, arena); + } + case google::protobuf::Value::kListValue: + return ParsedJsonListValue( + &ValueReflection::GetListValue(*generated_message), + MessageArenaOr(generated_message, arena)); + case google::protobuf::Value::kStructValue: + return ParsedJsonMapValue( + &ValueReflection::GetStructValue(*generated_message), + MessageArenaOr(generated_message, arena)); + default: + return ErrorValue::From(absl::InvalidArgumentError(absl::StrCat( + "unexpected value kind case: ", kind_case)), + arena); + } + } const auto reflection = GetValueReflectionOrDie(message->GetDescriptor()); const auto kind_case = reflection.GetKindCase(*message); switch (kind_case) {