Skip to content

Commit 672b10f

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

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
@@ -423,15 +423,14 @@ absl::Status ParsedMapFieldValue::ListKeys(
423423
field_->message_type()->map_key()));
424424
auto builder = NewListValueBuilder(arena);
425425
builder->Reserve(Size());
426-
auto begin = extensions::protobuf_internal::ConstMapBegin(*reflection,
427-
*message_, *field_);
428-
const auto end = extensions::protobuf_internal::ConstMapEnd(
429-
*reflection, *message_, *field_);
430-
for (; begin != end; ++begin) {
431-
Value scratch;
432-
(*key_accessor)(begin.GetKey(), message_, arena, &scratch);
433-
CEL_RETURN_IF_ERROR(builder->Add(std::move(scratch)));
434-
}
426+
CEL_RETURN_IF_ERROR(extensions::protobuf_internal::ForEachMapEntry(
427+
*reflection, *message_, *field_,
428+
[&](auto key_ref, auto value_ref) -> absl::Status {
429+
Value scratch;
430+
(*key_accessor)(key_ref, message_, arena, &scratch);
431+
CEL_RETURN_IF_ERROR(builder->Add(std::move(scratch)));
432+
return absl::OkStatus();
433+
}));
435434
*result = std::move(*builder).Build();
436435
return absl::OkStatus();
437436
}
@@ -454,12 +453,23 @@ absl::Status ParsedMapFieldValue::ForEach(
454453
CEL_ASSIGN_OR_RETURN(
455454
auto value_accessor,
456455
common_internal::MapFieldValueAccessorFor(value_field));
456+
Value key_scratch;
457+
Value value_scratch;
458+
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
459+
for (auto entry : reflection->GetMap(*message_, field_)) {
460+
(*key_accessor)(entry.key(), message_, arena, &key_scratch);
461+
(*value_accessor)(entry.value(), message_, value_field, descriptor_pool,
462+
message_factory, arena, &value_scratch);
463+
CEL_ASSIGN_OR_RETURN(auto ok, callback(key_scratch, value_scratch));
464+
if (!ok) {
465+
break;
466+
}
467+
}
468+
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
457469
auto begin = extensions::protobuf_internal::ConstMapBegin(
458470
*reflection, *message_, *field_);
459471
const auto end = extensions::protobuf_internal::ConstMapEnd(
460472
*reflection, *message_, *field_);
461-
Value key_scratch;
462-
Value value_scratch;
463473
for (; begin != end; ++begin) {
464474
(*key_accessor)(begin.GetKey(), message_, arena, &key_scratch);
465475
(*value_accessor)(begin.GetValueRef(), message_, value_field,
@@ -470,6 +480,7 @@ absl::Status ParsedMapFieldValue::ForEach(
470480
break;
471481
}
472482
}
483+
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
473484
}
474485
return absl::OkStatus();
475486
}
@@ -487,10 +498,17 @@ class ParsedMapFieldValueIterator final : public ValueIterator {
487498
value_field_(field->message_type()->map_value()),
488499
key_accessor_(key_accessor),
489500
value_accessor_(value_accessor),
501+
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
502+
begin_(message->GetReflection()->GetMap(*message, field).begin()),
503+
end_(message->GetReflection()->GetMap(*message, field).end())
504+
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
490505
begin_(extensions::protobuf_internal::ConstMapBegin(
491506
*message_->GetReflection(), *message_, *field)),
492507
end_(extensions::protobuf_internal::ConstMapEnd(
493-
*message_->GetReflection(), *message_, *field)) {}
508+
*message_->GetReflection(), *message_, *field))
509+
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
510+
{
511+
}
494512

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

@@ -503,7 +521,11 @@ class ParsedMapFieldValueIterator final : public ValueIterator {
503521
"ValueIterator::Next called after ValueIterator::HasNext returned "
504522
"false");
505523
}
524+
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
525+
(*key_accessor_)(begin_->key(), message_, arena, result);
526+
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
506527
(*key_accessor_)(begin_.GetKey(), message_, arena, result);
528+
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
507529
++begin_;
508530
return absl::OkStatus();
509531
}
@@ -521,7 +543,11 @@ class ParsedMapFieldValueIterator final : public ValueIterator {
521543
if (begin_ == end_) {
522544
return false;
523545
}
546+
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
547+
(*key_accessor_)(begin_->key(), message_, arena, key_or_value);
548+
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
524549
(*key_accessor_)(begin_.GetKey(), message_, arena, key_or_value);
550+
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
525551
++begin_;
526552
return true;
527553
}
@@ -539,11 +565,19 @@ class ParsedMapFieldValueIterator final : public ValueIterator {
539565
if (begin_ == end_) {
540566
return false;
541567
}
568+
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
569+
(*key_accessor_)(begin_->key(), message_, arena, key);
570+
if (value != nullptr) {
571+
(*value_accessor_)(begin_->value(), message_, value_field_,
572+
descriptor_pool, message_factory, arena, value);
573+
}
574+
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
542575
(*key_accessor_)(begin_.GetKey(), message_, arena, key);
543576
if (value != nullptr) {
544577
(*value_accessor_)(begin_.GetValueRef(), message_, value_field_,
545578
descriptor_pool, message_factory, arena, value);
546579
}
580+
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
547581
++begin_;
548582
return true;
549583
}
@@ -553,8 +587,13 @@ class ParsedMapFieldValueIterator final : public ValueIterator {
553587
const google::protobuf::FieldDescriptor* absl_nonnull const value_field_;
554588
const absl_nonnull common_internal::MapFieldKeyAccessor key_accessor_;
555589
const absl_nonnull common_internal::MapFieldValueAccessor value_accessor_;
590+
#if defined(PROTOBUF_HAS_MAP_REFLECTION_APIS)
591+
proto2::GenericConstMapRef::iterator begin_;
592+
const proto2::GenericConstMapRef::iterator end_;
593+
#else // PROTOBUF_HAS_MAP_REFLECTION_APIS
556594
google::protobuf::ConstMapIterator begin_;
557595
const google::protobuf::ConstMapIterator end_;
596+
#endif // PROTOBUF_HAS_MAP_REFLECTION_APIS
558597
};
559598

560599
} // 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)