Skip to content

Commit a0ea380

Browse files
sbenzaquencopybara-github
authored andcommitted
Prepare CEL for releasing the public protobuf reflection API for maps.
PiperOrigin-RevId: 963475659
1 parent ba959b4 commit a0ea380

9 files changed

Lines changed: 250 additions & 33 deletions

File tree

common/values/parsed_json_map_value.cc

Lines changed: 57 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
#include "common/value.h"
3434
#include "common/values/parsed_json_value.h"
3535
#include "common/values/values.h"
36+
#include "extensions/protobuf/internal/map_reflection.h"
3637
#include "internal/json.h"
3738
#include "internal/message_equality.h"
3839
#include "internal/status_macros.h"
@@ -299,12 +300,14 @@ absl::Status ParsedJsonMapValue::ListKeys(
299300
well_known_types::GetStructReflectionOrDie(value_->GetDescriptor());
300301
auto builder = NewListValueBuilder(arena);
301302
builder->Reserve(static_cast<size_t>(reflection.FieldsSize(*value_)));
302-
auto keys_begin = reflection.BeginFields(*value_);
303-
const auto keys_end = reflection.EndFields(*value_);
304-
for (; keys_begin != keys_end; ++keys_begin) {
305-
CEL_RETURN_IF_ERROR(builder->Add(
306-
Value::WrapMapFieldKeyString(keys_begin.GetKey(), value_, arena)));
307-
}
303+
CEL_RETURN_IF_ERROR(cel::extensions::protobuf_internal::ForEachMapEntry(
304+
*value_->GetReflection(), *value_, *reflection.GetFieldsDescriptor(),
305+
[&](auto key_ref, auto value_ref) -> absl::Status {
306+
CEL_RETURN_IF_ERROR(
307+
builder->Add(Value::WrapMapFieldKeyString(key_ref, value_, arena)));
308+
return absl::OkStatus();
309+
}));
310+
308311
*result = std::move(*builder).Build();
309312
return absl::OkStatus();
310313
}
@@ -321,6 +324,19 @@ absl::Status ParsedJsonMapValue::ForEach(
321324
well_known_types::GetStructReflectionOrDie(value_->GetDescriptor());
322325
Value key_scratch;
323326
Value value_scratch;
327+
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
328+
for (auto entry : value_->GetReflection()->GetMap(
329+
*value_, reflection.GetFieldsDescriptor())) {
330+
// We have to copy until `google::protobuf::MapKey` is just a view.
331+
key_scratch = StringValue(arena, entry.key().GetStringValue());
332+
value_scratch = common_internal::ParsedJsonValue(
333+
&entry.value().GetMessageValue(), arena);
334+
CEL_ASSIGN_OR_RETURN(auto ok, callback(key_scratch, value_scratch));
335+
if (!ok) {
336+
break;
337+
}
338+
}
339+
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
324340
auto map_begin = reflection.BeginFields(*value_);
325341
const auto map_end = reflection.EndFields(*value_);
326342
for (; map_begin != map_end; ++map_begin) {
@@ -333,6 +349,7 @@ absl::Status ParsedJsonMapValue::ForEach(
333349
break;
334350
}
335351
}
352+
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
336353
return absl::OkStatus();
337354
}
338355

@@ -345,8 +362,19 @@ class ParsedJsonMapValueIterator final : public ValueIterator {
345362
: message_(message),
346363
reflection_(well_known_types::GetStructReflectionOrDie(
347364
message_->GetDescriptor())),
365+
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
366+
begin_(message->GetReflection()
367+
->GetMap(*message, reflection_.GetFieldsDescriptor())
368+
.begin()),
369+
end_(message->GetReflection()
370+
->GetMap(*message, reflection_.GetFieldsDescriptor())
371+
.end())
372+
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
348373
begin_(reflection_.BeginFields(*message_)),
349-
end_(reflection_.EndFields(*message_)) {}
374+
end_(reflection_.EndFields(*message_))
375+
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
376+
{
377+
}
350378

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

@@ -359,7 +387,11 @@ class ParsedJsonMapValueIterator final : public ValueIterator {
359387
"`ValueIterator::Next` called after `ValueIterator::HasNext` "
360388
"returned false");
361389
}
390+
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
391+
*result = Value::WrapMapFieldKeyString(begin_->key(), message_, arena);
392+
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
362393
*result = Value::WrapMapFieldKeyString(begin_.GetKey(), message_, arena);
394+
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
363395
++begin_;
364396
return absl::OkStatus();
365397
}
@@ -377,8 +409,13 @@ class ParsedJsonMapValueIterator final : public ValueIterator {
377409
if (begin_ == end_) {
378410
return false;
379411
}
412+
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
413+
*key_or_value =
414+
Value::WrapMapFieldKeyString(begin_->key(), message_, arena);
415+
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
380416
*key_or_value =
381417
Value::WrapMapFieldKeyString(begin_.GetKey(), message_, arena);
418+
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
382419
++begin_;
383420
return true;
384421
}
@@ -396,20 +433,33 @@ class ParsedJsonMapValueIterator final : public ValueIterator {
396433
if (begin_ == end_) {
397434
return false;
398435
}
436+
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
437+
*key = Value::WrapMapFieldKeyString(begin_->key(), message_, arena);
438+
if (value != nullptr) {
439+
*value = common_internal::ParsedJsonValue(
440+
&begin_->value().GetMessageValue(), arena);
441+
}
442+
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
399443
*key = Value::WrapMapFieldKeyString(begin_.GetKey(), message_, arena);
400444
if (value != nullptr) {
401445
*value = common_internal::ParsedJsonValue(
402446
&begin_.GetValueRef().GetMessageValue(), arena);
403447
}
448+
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
404449
++begin_;
405450
return true;
406451
}
407452

408453
private:
409454
const google::protobuf::Message* absl_nonnull const message_;
410455
const well_known_types::StructReflection reflection_;
456+
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
457+
proto2::GenericConstMapRef::iterator begin_;
458+
const proto2::GenericConstMapRef::iterator end_;
459+
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
411460
google::protobuf::ConstMapIterator begin_;
412461
const google::protobuf::ConstMapIterator end_;
462+
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
413463
std::string scratch_;
414464
};
415465

common/values/parsed_map_field_value.cc

Lines changed: 51 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -422,15 +422,14 @@ absl::Status ParsedMapFieldValue::ListKeys(
422422
field_->message_type()->map_key()));
423423
auto builder = NewListValueBuilder(arena);
424424
builder->Reserve(Size());
425-
auto begin = extensions::protobuf_internal::ConstMapBegin(*reflection,
426-
*message_, *field_);
427-
const auto end = extensions::protobuf_internal::ConstMapEnd(
428-
*reflection, *message_, *field_);
429-
for (; begin != end; ++begin) {
430-
Value scratch;
431-
(*key_accessor)(begin.GetKey(), message_, arena, &scratch);
432-
CEL_RETURN_IF_ERROR(builder->Add(std::move(scratch)));
433-
}
425+
CEL_RETURN_IF_ERROR(extensions::protobuf_internal::ForEachMapEntry(
426+
*reflection, *message_, *field_,
427+
[&](auto key_ref, auto value_ref) -> absl::Status {
428+
Value scratch;
429+
(*key_accessor)(key_ref, message_, arena, &scratch);
430+
CEL_RETURN_IF_ERROR(builder->Add(std::move(scratch)));
431+
return absl::OkStatus();
432+
}));
434433
*result = std::move(*builder).Build();
435434
return absl::OkStatus();
436435
}
@@ -453,12 +452,23 @@ absl::Status ParsedMapFieldValue::ForEach(
453452
CEL_ASSIGN_OR_RETURN(
454453
auto value_accessor,
455454
common_internal::MapFieldValueAccessorFor(value_field));
455+
Value key_scratch;
456+
Value value_scratch;
457+
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
458+
for (auto entry : reflection->GetMap(*message_, field_)) {
459+
(*key_accessor)(entry.key(), message_, arena, &key_scratch);
460+
(*value_accessor)(entry.value(), message_, value_field, descriptor_pool,
461+
message_factory, arena, &value_scratch);
462+
CEL_ASSIGN_OR_RETURN(auto ok, callback(key_scratch, value_scratch));
463+
if (!ok) {
464+
break;
465+
}
466+
}
467+
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
456468
auto begin = extensions::protobuf_internal::ConstMapBegin(
457469
*reflection, *message_, *field_);
458470
const auto end = extensions::protobuf_internal::ConstMapEnd(
459471
*reflection, *message_, *field_);
460-
Value key_scratch;
461-
Value value_scratch;
462472
for (; begin != end; ++begin) {
463473
(*key_accessor)(begin.GetKey(), message_, arena, &key_scratch);
464474
(*value_accessor)(begin.GetValueRef(), message_, value_field,
@@ -469,6 +479,7 @@ absl::Status ParsedMapFieldValue::ForEach(
469479
break;
470480
}
471481
}
482+
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
472483
}
473484
return absl::OkStatus();
474485
}
@@ -486,10 +497,17 @@ class ParsedMapFieldValueIterator final : public ValueIterator {
486497
value_field_(field->message_type()->map_value()),
487498
key_accessor_(key_accessor),
488499
value_accessor_(value_accessor),
500+
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
501+
begin_(message->GetReflection()->GetMap(*message, field).begin()),
502+
end_(message->GetReflection()->GetMap(*message, field).end())
503+
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
489504
begin_(extensions::protobuf_internal::ConstMapBegin(
490505
*message_->GetReflection(), *message_, *field)),
491506
end_(extensions::protobuf_internal::ConstMapEnd(
492-
*message_->GetReflection(), *message_, *field)) {}
507+
*message_->GetReflection(), *message_, *field))
508+
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
509+
{
510+
}
493511

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

@@ -502,7 +520,11 @@ class ParsedMapFieldValueIterator final : public ValueIterator {
502520
"ValueIterator::Next called after ValueIterator::HasNext returned "
503521
"false");
504522
}
523+
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
524+
(*key_accessor_)(begin_->key(), message_, arena, result);
525+
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
505526
(*key_accessor_)(begin_.GetKey(), message_, arena, result);
527+
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
506528
++begin_;
507529
return absl::OkStatus();
508530
}
@@ -520,7 +542,11 @@ class ParsedMapFieldValueIterator final : public ValueIterator {
520542
if (begin_ == end_) {
521543
return false;
522544
}
545+
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
546+
(*key_accessor_)(begin_->key(), message_, arena, key_or_value);
547+
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
523548
(*key_accessor_)(begin_.GetKey(), message_, arena, key_or_value);
549+
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
524550
++begin_;
525551
return true;
526552
}
@@ -538,11 +564,19 @@ class ParsedMapFieldValueIterator final : public ValueIterator {
538564
if (begin_ == end_) {
539565
return false;
540566
}
567+
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
568+
(*key_accessor_)(begin_->key(), message_, arena, key);
569+
if (value != nullptr) {
570+
(*value_accessor_)(begin_->value(), message_, value_field_,
571+
descriptor_pool, message_factory, arena, value);
572+
}
573+
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
541574
(*key_accessor_)(begin_.GetKey(), message_, arena, key);
542575
if (value != nullptr) {
543576
(*value_accessor_)(begin_.GetValueRef(), message_, value_field_,
544577
descriptor_pool, message_factory, arena, value);
545578
}
579+
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
546580
++begin_;
547581
return true;
548582
}
@@ -552,8 +586,13 @@ class ParsedMapFieldValueIterator final : public ValueIterator {
552586
const google::protobuf::FieldDescriptor* absl_nonnull const value_field_;
553587
const absl_nonnull common_internal::MapFieldKeyAccessor key_accessor_;
554588
const absl_nonnull common_internal::MapFieldValueAccessor value_accessor_;
589+
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
590+
proto2::GenericConstMapRef::iterator begin_;
591+
const proto2::GenericConstMapRef::iterator end_;
592+
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
555593
google::protobuf::ConstMapIterator begin_;
556594
const google::protobuf::ConstMapIterator end_;
595+
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
557596
};
558597

559598
} // namespace

extensions/protobuf/internal/BUILD

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ cc_library(
2828
deps = [
2929
"@com_google_absl//absl/base:core_headers",
3030
"@com_google_absl//absl/base:nullability",
31+
"@com_google_absl//absl/status",
3132
"@com_google_protobuf//:protobuf",
3233
],
3334
)

extensions/protobuf/internal/map_reflection.cc

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,54 @@
1919
#include "google/protobuf/map_field.h"
2020
#include "google/protobuf/message.h"
2121

22+
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
23+
namespace cel::extensions::protobuf_internal {
24+
25+
bool LookupMapValue(const google::protobuf::Reflection& reflection,
26+
const google::protobuf::Message& message,
27+
const google::protobuf::FieldDescriptor& field,
28+
const google::protobuf::MapKey& key,
29+
google::protobuf::MapValueConstRef* value) {
30+
auto map = reflection.GetMap(message, &field);
31+
auto it = map.find(key);
32+
if (it == map.end()) return false;
33+
*value = it->value();
34+
return true;
35+
}
36+
37+
bool ContainsMapKey(const google::protobuf::Reflection& reflection,
38+
const google::protobuf::Message& message,
39+
const google::protobuf::FieldDescriptor& field,
40+
const google::protobuf::MapKey& key) {
41+
return reflection.GetMap(message, &field).contains(key);
42+
}
43+
44+
int MapSize(const google::protobuf::Reflection& reflection,
45+
const google::protobuf::Message& message,
46+
const google::protobuf::FieldDescriptor& field) {
47+
return reflection.GetMap(message, &field).size();
48+
}
49+
50+
bool InsertOrLookupMapValue(const google::protobuf::Reflection& reflection,
51+
google::protobuf::Message* message,
52+
const google::protobuf::FieldDescriptor& field,
53+
const google::protobuf::MapKey& key,
54+
google::protobuf::MapValueRef* value) {
55+
auto map = reflection.MutableMap(message, &field);
56+
auto res = map.try_emplace(key);
57+
*value = res.first->value();
58+
return res.second;
59+
}
60+
61+
bool DeleteMapValue(const google::protobuf::Reflection* absl_nonnull reflection,
62+
google::protobuf::Message* absl_nonnull message,
63+
const google::protobuf::FieldDescriptor* absl_nonnull field,
64+
const google::protobuf::MapKey& key) {
65+
return reflection->MutableMap(message, field).erase(key);
66+
}
67+
68+
} // namespace cel::extensions::protobuf_internal
69+
#else
2270
namespace google::protobuf::expr {
2371

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

132180
} // namespace cel::extensions::protobuf_internal
181+
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS

extensions/protobuf/internal/map_reflection.h

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,13 @@
1717

1818
#include "absl/base/attributes.h"
1919
#include "absl/base/nullability.h"
20+
#include "absl/status/status.h"
2021
#include "google/protobuf/descriptor.h"
2122
#include "google/protobuf/map_field.h"
2223
#include "google/protobuf/message.h"
2324

24-
#ifndef GOOGLE_PROTOBUF_HAS_CEL_MAP_REFLECTION_FRIEND
25+
#if !defined(GOOGLE_PROTOBUF_HAS_CEL_MAP_REFLECTION_FRIEND) && \
26+
!defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
2527
#error "protobuf library is too old, please update to version 3.15.0 or newer"
2628
#endif
2729

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

47+
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
48+
template <typename F>
49+
absl::Status ForEachMapEntry(const google::protobuf::Reflection& reflection,
50+
const google::protobuf::Message& message,
51+
const google::protobuf::FieldDescriptor& field, F f) {
52+
for (auto entry : reflection.GetMap(message, &field)) {
53+
if (auto status = f(entry.key(), entry.value()); !status.ok()) {
54+
return status;
55+
}
56+
}
57+
return absl::OkStatus();
58+
}
59+
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
4560
google::protobuf::ConstMapIterator ConstMapBegin(const google::protobuf::Reflection& reflection,
4661
const google::protobuf::Message& message,
4762
const google::protobuf::FieldDescriptor& field);
4863

4964
google::protobuf::ConstMapIterator ConstMapEnd(const google::protobuf::Reflection& reflection,
5065
const google::protobuf::Message& message,
5166
const google::protobuf::FieldDescriptor& field);
67+
template <typename F>
68+
absl::Status ForEachMapEntry(const google::protobuf::Reflection& reflection,
69+
const google::protobuf::Message& message,
70+
const google::protobuf::FieldDescriptor& field, F f) {
71+
auto it = ConstMapBegin(reflection, message, field);
72+
auto end = ConstMapEnd(reflection, message, field);
73+
for (; it != end; ++it) {
74+
if (auto status = f(it.GetKey(), it.GetValueRef()); !status.ok()) {
75+
return status;
76+
}
77+
}
78+
return absl::OkStatus();
79+
}
80+
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
5281

5382
bool InsertOrLookupMapValue(const google::protobuf::Reflection& reflection,
5483
google::protobuf::Message* message,

0 commit comments

Comments
 (0)