Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
64 changes: 57 additions & 7 deletions common/values/parsed_json_map_value.cc
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@
#include "common/value.h"
#include "common/values/parsed_json_value.h"
#include "common/values/values.h"
#include "extensions/protobuf/internal/map_reflection.h"
#include "internal/json.h"
#include "internal/message_equality.h"
#include "internal/status_macros.h"
Expand Down Expand Up @@ -299,12 +300,14 @@ absl::Status ParsedJsonMapValue::ListKeys(
well_known_types::GetStructReflectionOrDie(value_->GetDescriptor());
auto builder = NewListValueBuilder(arena);
builder->Reserve(static_cast<size_t>(reflection.FieldsSize(*value_)));
auto keys_begin = reflection.BeginFields(*value_);
const auto keys_end = reflection.EndFields(*value_);
for (; keys_begin != keys_end; ++keys_begin) {
CEL_RETURN_IF_ERROR(builder->Add(
Value::WrapMapFieldKeyString(keys_begin.GetKey(), value_, arena)));
}
CEL_RETURN_IF_ERROR(cel::extensions::protobuf_internal::ForEachMapEntry(
*value_->GetReflection(), *value_, *reflection.GetFieldsDescriptor(),
[&](auto key_ref, auto value_ref) -> absl::Status {
CEL_RETURN_IF_ERROR(
builder->Add(Value::WrapMapFieldKeyString(key_ref, value_, arena)));
return absl::OkStatus();
}));

*result = std::move(*builder).Build();
return absl::OkStatus();
}
Expand All @@ -321,6 +324,19 @@ absl::Status ParsedJsonMapValue::ForEach(
well_known_types::GetStructReflectionOrDie(value_->GetDescriptor());
Value key_scratch;
Value value_scratch;
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
for (auto entry : value_->GetReflection()->GetMap(
*value_, reflection.GetFieldsDescriptor())) {
// We have to copy until `google::protobuf::MapKey` is just a view.
key_scratch = StringValue(arena, entry.key().GetStringValue());
value_scratch = common_internal::ParsedJsonValue(
&entry.value().GetMessageValue(), arena);
CEL_ASSIGN_OR_RETURN(auto ok, callback(key_scratch, value_scratch));
if (!ok) {
break;
}
}
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
auto map_begin = reflection.BeginFields(*value_);
const auto map_end = reflection.EndFields(*value_);
for (; map_begin != map_end; ++map_begin) {
Expand All @@ -333,6 +349,7 @@ absl::Status ParsedJsonMapValue::ForEach(
break;
}
}
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
return absl::OkStatus();
}

Expand All @@ -345,8 +362,19 @@ class ParsedJsonMapValueIterator final : public ValueIterator {
: message_(message),
reflection_(well_known_types::GetStructReflectionOrDie(
message_->GetDescriptor())),
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
begin_(message->GetReflection()
->GetMap(*message, reflection_.GetFieldsDescriptor())
.begin()),
end_(message->GetReflection()
->GetMap(*message, reflection_.GetFieldsDescriptor())
.end())
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
begin_(reflection_.BeginFields(*message_)),
end_(reflection_.EndFields(*message_)) {}
end_(reflection_.EndFields(*message_))
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
{
}

bool HasNext() override { return begin_ != end_; }

Expand All @@ -359,7 +387,11 @@ class ParsedJsonMapValueIterator final : public ValueIterator {
"`ValueIterator::Next` called after `ValueIterator::HasNext` "
"returned false");
}
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
*result = Value::WrapMapFieldKeyString(begin_->key(), message_, arena);
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
*result = Value::WrapMapFieldKeyString(begin_.GetKey(), message_, arena);
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
++begin_;
return absl::OkStatus();
}
Expand All @@ -377,8 +409,13 @@ class ParsedJsonMapValueIterator final : public ValueIterator {
if (begin_ == end_) {
return false;
}
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
*key_or_value =
Value::WrapMapFieldKeyString(begin_->key(), message_, arena);
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
*key_or_value =
Value::WrapMapFieldKeyString(begin_.GetKey(), message_, arena);
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
++begin_;
return true;
}
Expand All @@ -396,20 +433,33 @@ class ParsedJsonMapValueIterator final : public ValueIterator {
if (begin_ == end_) {
return false;
}
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
*key = Value::WrapMapFieldKeyString(begin_->key(), message_, arena);
if (value != nullptr) {
*value = common_internal::ParsedJsonValue(
&begin_->value().GetMessageValue(), arena);
}
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
*key = Value::WrapMapFieldKeyString(begin_.GetKey(), message_, arena);
if (value != nullptr) {
*value = common_internal::ParsedJsonValue(
&begin_.GetValueRef().GetMessageValue(), arena);
}
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
++begin_;
return true;
}

private:
const google::protobuf::Message* absl_nonnull const message_;
const well_known_types::StructReflection reflection_;
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
proto2::GenericConstMapRef::iterator begin_;
const proto2::GenericConstMapRef::iterator end_;
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
google::protobuf::ConstMapIterator begin_;
const google::protobuf::ConstMapIterator end_;
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
std::string scratch_;
};

Expand Down
63 changes: 51 additions & 12 deletions common/values/parsed_map_field_value.cc
Original file line number Diff line number Diff line change
Expand Up @@ -423,15 +423,14 @@ absl::Status ParsedMapFieldValue::ListKeys(
field_->message_type()->map_key()));
auto builder = NewListValueBuilder(arena);
builder->Reserve(Size());
auto begin = extensions::protobuf_internal::ConstMapBegin(*reflection,
*message_, *field_);
const auto end = extensions::protobuf_internal::ConstMapEnd(
*reflection, *message_, *field_);
for (; begin != end; ++begin) {
Value scratch;
(*key_accessor)(begin.GetKey(), message_, arena, &scratch);
CEL_RETURN_IF_ERROR(builder->Add(std::move(scratch)));
}
CEL_RETURN_IF_ERROR(extensions::protobuf_internal::ForEachMapEntry(
*reflection, *message_, *field_,
[&](auto key_ref, auto value_ref) -> absl::Status {
Value scratch;
(*key_accessor)(key_ref, message_, arena, &scratch);
CEL_RETURN_IF_ERROR(builder->Add(std::move(scratch)));
return absl::OkStatus();
}));
*result = std::move(*builder).Build();
return absl::OkStatus();
}
Expand All @@ -454,12 +453,23 @@ absl::Status ParsedMapFieldValue::ForEach(
CEL_ASSIGN_OR_RETURN(
auto value_accessor,
common_internal::MapFieldValueAccessorFor(value_field));
Value key_scratch;
Value value_scratch;
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
for (auto entry : reflection->GetMap(*message_, field_)) {
(*key_accessor)(entry.key(), message_, arena, &key_scratch);
(*value_accessor)(entry.value(), message_, value_field, descriptor_pool,
message_factory, arena, &value_scratch);
CEL_ASSIGN_OR_RETURN(auto ok, callback(key_scratch, value_scratch));
if (!ok) {
break;
}
}
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
auto begin = extensions::protobuf_internal::ConstMapBegin(
*reflection, *message_, *field_);
const auto end = extensions::protobuf_internal::ConstMapEnd(
*reflection, *message_, *field_);
Value key_scratch;
Value value_scratch;
for (; begin != end; ++begin) {
(*key_accessor)(begin.GetKey(), message_, arena, &key_scratch);
(*value_accessor)(begin.GetValueRef(), message_, value_field,
Expand All @@ -470,6 +480,7 @@ absl::Status ParsedMapFieldValue::ForEach(
break;
}
}
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
}
return absl::OkStatus();
}
Expand All @@ -487,10 +498,17 @@ class ParsedMapFieldValueIterator final : public ValueIterator {
value_field_(field->message_type()->map_value()),
key_accessor_(key_accessor),
value_accessor_(value_accessor),
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
begin_(message->GetReflection()->GetMap(*message, field).begin()),
end_(message->GetReflection()->GetMap(*message, field).end())
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
begin_(extensions::protobuf_internal::ConstMapBegin(
*message_->GetReflection(), *message_, *field)),
end_(extensions::protobuf_internal::ConstMapEnd(
*message_->GetReflection(), *message_, *field)) {}
*message_->GetReflection(), *message_, *field))
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
{
}

bool HasNext() override { return begin_ != end_; }

Expand All @@ -503,7 +521,11 @@ class ParsedMapFieldValueIterator final : public ValueIterator {
"ValueIterator::Next called after ValueIterator::HasNext returned "
"false");
}
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
(*key_accessor_)(begin_->key(), message_, arena, result);
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
(*key_accessor_)(begin_.GetKey(), message_, arena, result);
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
++begin_;
return absl::OkStatus();
}
Expand All @@ -521,7 +543,11 @@ class ParsedMapFieldValueIterator final : public ValueIterator {
if (begin_ == end_) {
return false;
}
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
(*key_accessor_)(begin_->key(), message_, arena, key_or_value);
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
(*key_accessor_)(begin_.GetKey(), message_, arena, key_or_value);
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
++begin_;
return true;
}
Expand All @@ -539,11 +565,19 @@ class ParsedMapFieldValueIterator final : public ValueIterator {
if (begin_ == end_) {
return false;
}
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
(*key_accessor_)(begin_->key(), message_, arena, key);
if (value != nullptr) {
(*value_accessor_)(begin_->value(), message_, value_field_,
descriptor_pool, message_factory, arena, value);
}
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
(*key_accessor_)(begin_.GetKey(), message_, arena, key);
if (value != nullptr) {
(*value_accessor_)(begin_.GetValueRef(), message_, value_field_,
descriptor_pool, message_factory, arena, value);
}
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
++begin_;
return true;
}
Expand All @@ -553,8 +587,13 @@ class ParsedMapFieldValueIterator final : public ValueIterator {
const google::protobuf::FieldDescriptor* absl_nonnull const value_field_;
const absl_nonnull common_internal::MapFieldKeyAccessor key_accessor_;
const absl_nonnull common_internal::MapFieldValueAccessor value_accessor_;
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
proto2::GenericConstMapRef::iterator begin_;
const proto2::GenericConstMapRef::iterator end_;
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
google::protobuf::ConstMapIterator begin_;
const google::protobuf::ConstMapIterator end_;
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
};

} // namespace
Expand Down
1 change: 1 addition & 0 deletions extensions/protobuf/internal/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ cc_library(
deps = [
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/base:nullability",
"@com_google_absl//absl/status",
"@com_google_protobuf//:protobuf",
],
)
Expand Down
49 changes: 49 additions & 0 deletions extensions/protobuf/internal/map_reflection.cc
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,54 @@
#include "google/protobuf/map_field.h"
#include "google/protobuf/message.h"

#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
namespace cel::extensions::protobuf_internal {

bool LookupMapValue(const google::protobuf::Reflection& reflection,
const google::protobuf::Message& message,
const google::protobuf::FieldDescriptor& field,
const google::protobuf::MapKey& key,
google::protobuf::MapValueConstRef* value) {
auto map = reflection.GetMap(message, &field);
auto it = map.find(key);
if (it == map.end()) return false;
*value = it->value();
return true;
}

bool ContainsMapKey(const google::protobuf::Reflection& reflection,
const google::protobuf::Message& message,
const google::protobuf::FieldDescriptor& field,
const google::protobuf::MapKey& key) {
return reflection.GetMap(message, &field).contains(key);
}

int MapSize(const google::protobuf::Reflection& reflection,
const google::protobuf::Message& message,
const google::protobuf::FieldDescriptor& field) {
return reflection.GetMap(message, &field).size();
}

bool InsertOrLookupMapValue(const google::protobuf::Reflection& reflection,
google::protobuf::Message* message,
const google::protobuf::FieldDescriptor& field,
const google::protobuf::MapKey& key,
google::protobuf::MapValueRef* value) {
auto map = reflection.MutableMap(message, &field);
auto res = map.try_emplace(key);
*value = res.first->value();
return res.second;
}

bool DeleteMapValue(const google::protobuf::Reflection* absl_nonnull reflection,
google::protobuf::Message* absl_nonnull message,
const google::protobuf::FieldDescriptor* absl_nonnull field,
const google::protobuf::MapKey& key) {
return reflection->MutableMap(message, field).erase(key);
}

} // namespace cel::extensions::protobuf_internal
#else
namespace google::protobuf::expr {

class CelMapReflectionFriend final {
Expand Down Expand Up @@ -130,3 +178,4 @@ bool DeleteMapValue(const google::protobuf::Reflection* absl_nonnull reflection,
}

} // namespace cel::extensions::protobuf_internal
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
31 changes: 30 additions & 1 deletion extensions/protobuf/internal/map_reflection.h
Original file line number Diff line number Diff line change
Expand Up @@ -17,11 +17,13 @@

#include "absl/base/attributes.h"
#include "absl/base/nullability.h"
#include "absl/status/status.h"
#include "google/protobuf/descriptor.h"
#include "google/protobuf/map_field.h"
#include "google/protobuf/message.h"

#ifndef GOOGLE_PROTOBUF_HAS_CEL_MAP_REFLECTION_FRIEND
#if !defined(GOOGLE_PROTOBUF_HAS_CEL_MAP_REFLECTION_FRIEND) && \
!defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
#error "protobuf library is too old, please update to version 3.15.0 or newer"
#endif

Expand All @@ -42,13 +44,40 @@ int MapSize(const google::protobuf::Reflection& reflection,
const google::protobuf::Message& message,
const google::protobuf::FieldDescriptor& field);

#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
template <typename F>
absl::Status ForEachMapEntry(const google::protobuf::Reflection& reflection,
const google::protobuf::Message& message,
const google::protobuf::FieldDescriptor& field, F f) {
for (auto entry : reflection.GetMap(message, &field)) {
if (auto status = f(entry.key(), entry.value()); !status.ok()) {
return status;
}
}
return absl::OkStatus();
}
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
google::protobuf::ConstMapIterator ConstMapBegin(const google::protobuf::Reflection& reflection,
const google::protobuf::Message& message,
const google::protobuf::FieldDescriptor& field);

google::protobuf::ConstMapIterator ConstMapEnd(const google::protobuf::Reflection& reflection,
const google::protobuf::Message& message,
const google::protobuf::FieldDescriptor& field);
template <typename F>
absl::Status ForEachMapEntry(const google::protobuf::Reflection& reflection,
const google::protobuf::Message& message,
const google::protobuf::FieldDescriptor& field, F f) {
auto it = ConstMapBegin(reflection, message, field);
auto end = ConstMapEnd(reflection, message, field);
for (; it != end; ++it) {
if (auto status = f(it.GetKey(), it.GetValueRef()); !status.ok()) {
return status;
}
}
return absl::OkStatus();
}
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS

bool InsertOrLookupMapValue(const google::protobuf::Reflection& reflection,
google::protobuf::Message* message,
Expand Down
Loading
Loading