Avoid copies in proto map reflection.

PiperOrigin-RevId: 590472873
protocolbuffers · Dec 13, 2023 · a4b5c7e · a4b5c7e
1 parent b429304
commit a4b5c7e
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 17 deletions.
diff --git a/src/google/protobuf/map_field_inl.h b/src/google/protobuf/map_field_inl.h
@@ -9,6 +9,7 @@
 #define GOOGLE_PROTOBUF_MAP_FIELD_INL_H__
 
 #include <cstddef>
+#include <cstdint>
 #include <memory>
 #include <string>
 #include <tuple>
@@ -31,38 +32,37 @@
 namespace google {
 namespace protobuf {
 namespace internal {
-// UnwrapMapKey template
-template <typename T>
-T UnwrapMapKey(const MapKey& map_key);
-template <>
-inline int32_t UnwrapMapKey(const MapKey& map_key) {
+// UnwrapMapKey template. We're using overloading rather than template
+// specialization so that we can return a value or reference type depending on
+// `T`.
+inline int32_t UnwrapMapKeyImpl(const MapKey& map_key, const int32_t*) {
   return map_key.GetInt32Value();
 }
-template <>
-inline uint32_t UnwrapMapKey(const MapKey& map_key) {
+inline uint32_t UnwrapMapKeyImpl(const MapKey& map_key, const uint32_t*) {
   return map_key.GetUInt32Value();
 }
-template <>
-inline int64_t UnwrapMapKey(const MapKey& map_key) {
+inline int64_t UnwrapMapKeyImpl(const MapKey& map_key, const int64_t*) {
   return map_key.GetInt64Value();
 }
-template <>
-inline uint64_t UnwrapMapKey(const MapKey& map_key) {
+inline uint64_t UnwrapMapKeyImpl(const MapKey& map_key, const uint64_t*) {
   return map_key.GetUInt64Value();
 }
-template <>
-inline bool UnwrapMapKey(const MapKey& map_key) {
+inline bool UnwrapMapKeyImpl(const MapKey& map_key, const bool*) {
   return map_key.GetBoolValue();
 }
-template <>
-inline std::string UnwrapMapKey(const MapKey& map_key) {
+inline const std::string& UnwrapMapKeyImpl(const MapKey& map_key,
+                                           const std::string*) {
   return map_key.GetStringValue();
 }
-template <>
-inline MapKey UnwrapMapKey(const MapKey& map_key) {
+inline const MapKey& UnwrapMapKeyImpl(const MapKey& map_key, const MapKey*) {
   return map_key;
 }
 
+template <typename T>
+decltype(auto) UnwrapMapKey(const MapKey& map_key) {
+  return UnwrapMapKeyImpl(map_key, static_cast<T*>(nullptr));
+}
+
 // SetMapKey
 inline void SetMapKey(MapKey* map_key, int32_t value) {
   map_key->SetInt32Value(value);

diff --git a/src/google/protobuf/map_test.cc b/src/google/protobuf/map_test.cc
@@ -22,6 +22,7 @@
 #include "absl/strings/str_cat.h"
 #include "google/protobuf/arena_test_util.h"
 #include "google/protobuf/internal_visibility_for_testing.h"
+#include "google/protobuf/map_field.h"
 #include "google/protobuf/map_proto2_unittest.pb.h"
 #include "google/protobuf/map_unittest.pb.h"
 #include "google/protobuf/reflection_tester.h"

diff --git a/src/google/protobuf/reflection_tester.h b/src/google/protobuf/reflection_tester.h
@@ -8,6 +8,7 @@
 #ifndef GOOGLE_PROTOBUF_REFLECTION_TESTER_H__
 #define GOOGLE_PROTOBUF_REFLECTION_TESTER_H__
 
+#include "google/protobuf/map_field.h"
 #include "google/protobuf/message.h"
 
 // Must be included last.
@@ -45,6 +46,15 @@ class MapReflectionTester {
   MapIterator MapEnd(Message* message, const std::string& field_name);
   int MapSize(const Message& message, const std::string& field_name);
 
+  static MapValueConstRef LookupMapValue(const Reflection& reflection,
+                                         const Message& message,
+                                         const FieldDescriptor& descriptor,
+                                         const MapKey& map_key) {
+    MapValueConstRef map_val_const;
+    reflection.LookupMapValue(message, &descriptor, map_key, &map_val_const);
+    return map_val_const;
+  }
+
   static std::string long_string() {
     return "This is a very long string that goes in the heap";
   }