From a4b5c7e9d14ea2dc5025f9fdcd5c143011edb26d Mon Sep 17 00:00:00 2001
From: Protobuf Team Bot <protobuf-github-bot@google.com>
Date: Tue, 12 Dec 2023 23:53:34 -0800
Subject: [PATCH] Avoid copies in proto map reflection.

PiperOrigin-RevId: 590472873
---
 src/google/protobuf/map_field_inl.h     | 34 ++++++++++++-------------
 src/google/protobuf/map_test.cc         |  1 +
 src/google/protobuf/reflection_tester.h | 10 ++++++++
 3 files changed, 28 insertions(+), 17 deletions(-)

diff --git a/src/google/protobuf/map_field_inl.h b/src/google/protobuf/map_field_inl.h
index 499d8ed2ec07..e6a7acfc9ba1 100644
--- 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
index 7e9a80df4271..59647affc72e 100644
--- 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
index b07e45ad9554..4d258485b109 100644
--- 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";
   }