Update `google::protobuf::MapKey` to be a view of string keys. After this change, `google::protobuf::MapKey` will no longer own string keys. It is simply a view over the keys. Callers of `google::protobuf::MapKey::SetStringValue` are required to ensure the supplied string outlives the instance of `google::protobuf::MapKey`. `google::protobuf::internal::DynamicMapField` now uses `google::protobuf::internal::DynamicMapKey` as its key value, instead of `google::protobuf::MapKey`.

PiperOrigin-RevId: 682051238

src/google/protobuf/BUILD.bazel

@@ -681,6 +681,7 @@ cc_library(
         "@com_google_absl//absl/types:optional",
         "@com_google_absl//absl/types:span",
         "@com_google_absl//absl/types:variant",
+        "@com_google_absl//absl/utility",
     ],
 )
 

src/google/protobuf/map.h

@@ -210,20 +210,22 @@ using KeyForBase = typename KeyForBaseImpl<T>::type;
 // only accept `key_type`.
 template <typename key_type>
 struct TransparentSupport {
+  static_assert(std::is_scalar<key_type>::value,
+                "Should only be used for ints.");
+
   // We hash all the scalars as uint64_t so that we can implement the same hash
   // function for VariantKey. This way we can have MapKey provide the same hash
   // as the underlying value would have.
-  using hash = absl::Hash<
-      std::conditional_t<std::is_scalar<key_type>::value, uint64_t, key_type>>;
+  using hash = absl::Hash<uint64_t>;
 
-  static bool Equals(const key_type& a, const key_type& b) { return a == b; }
+  static bool Equals(key_type a, key_type b) { return a == b; }
 
   template <typename K>
   using key_arg = key_type;
 
-  using ViewType = std::conditional_t<std::is_scalar<key_type>::value, key_type,
-                                      const key_type&>;
-  static ViewType ToView(const key_type& v) { return v; }
+  using ViewType = key_type;
+
+  static key_type ToView(key_type v) { return v; }
 };
 
 // We add transparent support for std::string keys. We use
@@ -1017,7 +1019,7 @@ class KeyMapBase : public UntypedMapBase {
     // determined whether we are inserting into an empty list, a short list,
     // or whatever.  But it's probably cheap enough to recompute that here;
     // it's likely that we're inserting into an empty or short list.
-    ABSL_DCHECK(FindHelper(node->key()).node == nullptr);
+    ABSL_DCHECK(FindHelper(TS::ToView(node->key())).node == nullptr);
     if (TableEntryIsEmpty(b)) {
       InsertUniqueInList(b, node);
       index_of_first_non_null_ = (std::min)(index_of_first_non_null_, b);
@@ -1116,7 +1118,7 @@ class KeyMapBase : public UntypedMapBase {
   void TransferList(KeyNode* node) {
     do {
       auto* next = static_cast<KeyNode*>(node->next);
-      InsertUnique(BucketNumber(node->key()), node);
+      InsertUnique(BucketNumber(TS::ToView(node->key())), node);
       node = next;
     } while (node != nullptr);
   }
@@ -1152,7 +1154,7 @@ class KeyMapBase : public UntypedMapBase {
     // not.  Revalidate just to be sure.  This case is rare enough that we
     // don't worry about potential optimizations, such as having a custom
     // find-like method that compares Node* instead of the key.
-    auto res = FindHelper(node->key(), it);
+    auto res = FindHelper(TS::ToView(node->key()), it);
     bucket_index = res.bucket;
     return TableEntryIsList(bucket_index);
   }

src/google/protobuf/map_field.cc

@@ -8,11 +8,14 @@
 #include "google/protobuf/map_field.h"
 
 #include <atomic>
+#include <cstdint>
 #include <string>
 #include <utility>
 #include <vector>
 
 #include "absl/log/absl_check.h"
+#include "absl/types/variant.h"
+#include "absl/utility/utility.h"
 #include "google/protobuf/map.h"
 #include "google/protobuf/map_field_inl.h"
 #include "google/protobuf/port.h"
@@ -45,6 +48,58 @@ VariantKey RealKeyToVariantKey<MapKey>::operator()(const MapKey& value) const {
   }
 }
 
+DynamicMapKey::Variant DynamicMapKey::FromMapKey(google::protobuf::MapKey map_key) {
+  switch (map_key.type()) {
+    case FieldDescriptor::CPPTYPE_STRING:
+      return DynamicMapKey::Variant(absl::in_place_type<std::string>,
+                                    map_key.GetStringValue());
+    case FieldDescriptor::CPPTYPE_INT64:
+      return DynamicMapKey::Variant(map_key.GetInt64Value());
+    case FieldDescriptor::CPPTYPE_INT32:
+      return DynamicMapKey::Variant(map_key.GetInt32Value());
+    case FieldDescriptor::CPPTYPE_UINT64:
+      return DynamicMapKey::Variant(map_key.GetUInt64Value());
+    case FieldDescriptor::CPPTYPE_UINT32:
+      return DynamicMapKey::Variant(map_key.GetUInt32Value());
+    case FieldDescriptor::CPPTYPE_BOOL:
+      return DynamicMapKey::Variant(map_key.GetBoolValue());
+    default:
+      internal::Unreachable();
+  }
+}
+
+namespace {
+
+struct DynamicMapKeyToMapKey {
+  google::protobuf::MapKey* map_key;
+
+  void operator()(bool value) const { map_key->SetBoolValue(value); }
+
+  void operator()(int32_t value) const { map_key->SetInt32Value(value); }
+
+  void operator()(int64_t value) const { map_key->SetInt64Value(value); }
+
+  void operator()(uint32_t value) const { map_key->SetUInt32Value(value); }
+
+  void operator()(uint64_t value) const { map_key->SetUInt64Value(value); }
+
+  void operator()(const std::string& value) const {
+    map_key->SetStringValue(value);
+  }
+};
+
+}  // namespace
+
+google::protobuf::MapKey DynamicMapKey::ToMapKey() const {
+  google::protobuf::MapKey result;
+  absl::visit(DynamicMapKeyToMapKey{&result}, variant_);
+  return result;
+}
+
+VariantKey DynamicMapKey::ToVariantKey() const {
+  return absl::visit([](const auto& alt) { return VariantKey(alt); }, variant_);
+}
+
 MapFieldBase::~MapFieldBase() {
   ABSL_DCHECK_EQ(arena(), nullptr);
   delete maybe_payload();
@@ -422,7 +477,7 @@ DynamicMapField::DynamicMapField(const Message* default_entry)
       default_entry_(default_entry) {}
 
 DynamicMapField::DynamicMapField(const Message* default_entry, Arena* arena)
-    : TypeDefinedMapFieldBase<MapKey, MapValueRef>(&kVTable, arena),
+    : TypeDefinedMapFieldBase<DynamicMapKey, MapValueRef>(&kVTable, arena),
       default_entry_(default_entry) {}
 
 constexpr DynamicMapField::VTable DynamicMapField::kVTable =
@@ -485,7 +540,7 @@ bool DynamicMapField::InsertOrLookupMapValueNoSyncImpl(MapFieldBase& base,
                                                        const MapKey& map_key,
                                                        MapValueRef* val) {
   auto& self = static_cast<DynamicMapField&>(base);
-  Map<MapKey, MapValueRef>::iterator iter = self.map_.find(map_key);
+  auto iter = self.map_.find(map_key);
   if (iter == self.map_.end()) {
     MapValueRef& map_val = self.map_[map_key];
     self.AllocateMapValue(&map_val);
@@ -502,13 +557,12 @@ void DynamicMapField::MergeFromImpl(MapFieldBase& base,
                                     const MapFieldBase& other) {
   auto& self = static_cast<DynamicMapField&>(base);
   ABSL_DCHECK(self.IsMapValid() && other.IsMapValid());
-  Map<MapKey, MapValueRef>* map = self.MutableMap();
+  Map<DynamicMapKey, MapValueRef>* map = self.MutableMap();
   const DynamicMapField& other_field =
       reinterpret_cast<const DynamicMapField&>(other);
-  for (Map<MapKey, MapValueRef>::const_iterator other_it =
-           other_field.map_.begin();
+  for (auto other_it = other_field.map_.begin();
        other_it != other_field.map_.end(); ++other_it) {
-    Map<MapKey, MapValueRef>::iterator iter = map->find(other_it->first);
+    auto iter = map->find(other_it->first);
     MapValueRef* map_val;
     if (iter == map->end()) {
       map_val = &self.map_[other_it->first];
@@ -579,11 +633,11 @@ size_t DynamicMapField::SpaceUsedExcludingSelfNoLockImpl(
   }
   size_t map_size = self.map_.size();
   if (map_size) {
-    Map<MapKey, MapValueRef>::const_iterator it = self.map_.begin();
+    auto it = self.map_.begin();
     size += sizeof(it->first) * map_size;
     size += sizeof(it->second) * map_size;
     // If key is string, add the allocated space.
-    if (it->first.type() == FieldDescriptor::CPPTYPE_STRING) {
+    if (it->first.IsString()) {
       size += sizeof(std::string) * map_size;
     }
     // Add the allocated space in MapValueRef.

src/google/protobuf/map_field.h

@@ -16,13 +16,15 @@
 #include <type_traits>
 #include <utility>
 
+#include "absl/base/attributes.h"
+#include "absl/hash/hash.h"
 #include "absl/log/absl_check.h"
 #include "absl/log/absl_log.h"
 #include "absl/strings/string_view.h"
 #include "absl/synchronization/mutex.h"
+#include "absl/types/variant.h"
 #include "google/protobuf/arena.h"
 #include "google/protobuf/descriptor.h"
-#include "google/protobuf/explicitly_constructed.h"
 #include "google/protobuf/generated_message_reflection.h"
 #include "google/protobuf/generated_message_util.h"
 #include "google/protobuf/internal_visibility.h"
@@ -48,6 +50,10 @@ namespace protobuf {
 class DynamicMessage;
 class MapIterator;
 
+namespace internal {
+class DynamicMapKey;
+}  // namespace internal
+
 // Microsoft compiler complains about non-virtual destructor,
 // even when the destructor is private.
 #ifdef _MSC_VER
@@ -65,23 +71,14 @@ class MapIterator;
                     << FieldDescriptor::CppTypeName(type());              \
   }
 
-// MapKey is an union type for representing any possible
-// map key.
+// MapKey is an union type for representing any possible map key. For strings,
+// map key does not own the underlying data. It is up to the caller to ensure
+// any supplied strings outlive any instance of this class.
 class PROTOBUF_EXPORT MapKey {
  public:
-  MapKey() : type_() {}
-  MapKey(const MapKey& other) : type_() { CopyFrom(other); }
-
-  MapKey& operator=(const MapKey& other) {
-    CopyFrom(other);
-    return *this;
-  }
-
-  ~MapKey() {
-    if (type_ == FieldDescriptor::CPPTYPE_STRING) {
-      val_.string_value.Destruct();
-    }
-  }
+  MapKey() = default;
+  MapKey(const MapKey&) = default;
+  MapKey& operator=(const MapKey&) = default;
 
   FieldDescriptor::CppType type() const {
     if (type_ == FieldDescriptor::CppType()) {
@@ -114,7 +111,7 @@ class PROTOBUF_EXPORT MapKey {
   }
   void SetStringValue(absl::string_view val) {
     SetType(FieldDescriptor::CPPTYPE_STRING);
-    val_.string_value.get_mutable()->assign(val.data(), val.size());
+    val_.string_value = val;
   }
 
   int64_t GetInt64Value() const {
@@ -139,7 +136,7 @@ class PROTOBUF_EXPORT MapKey {
   }
   absl::string_view GetStringValue() const {
     TYPE_CHECK(FieldDescriptor::CPPTYPE_STRING, "MapKey::GetStringValue");
-    return val_.string_value.get();
+    return val_.string_value;
   }
 
   bool operator<(const MapKey& other) const {
@@ -156,7 +153,7 @@ class PROTOBUF_EXPORT MapKey {
         ABSL_LOG(FATAL) << "Unsupported";
         return false;
       case FieldDescriptor::CPPTYPE_STRING:
-        return val_.string_value.get() < other.val_.string_value.get();
+        return val_.string_value < other.val_.string_value;
       case FieldDescriptor::CPPTYPE_INT64:
         return val_.int64_value < other.val_.int64_value;
       case FieldDescriptor::CPPTYPE_INT32:
@@ -184,7 +181,7 @@ class PROTOBUF_EXPORT MapKey {
         ABSL_LOG(FATAL) << "Unsupported";
         break;
       case FieldDescriptor::CPPTYPE_STRING:
-        return val_.string_value.get() == other.val_.string_value.get();
+        return val_.string_value == other.val_.string_value;
       case FieldDescriptor::CPPTYPE_INT64:
         return val_.int64_value == other.val_.int64_value;
       case FieldDescriptor::CPPTYPE_INT32:
@@ -210,7 +207,7 @@ class PROTOBUF_EXPORT MapKey {
         ABSL_LOG(FATAL) << "Unsupported";
         break;
       case FieldDescriptor::CPPTYPE_STRING:
-        *val_.string_value.get_mutable() = other.val_.string_value.get();
+        val_.string_value = other.val_.string_value;
         break;
       case FieldDescriptor::CPPTYPE_INT64:
         val_.int64_value = other.val_.int64_value;
@@ -239,7 +236,7 @@ class PROTOBUF_EXPORT MapKey {
 
   union KeyValue {
     KeyValue() {}
-    internal::ExplicitlyConstructed<std::string> string_value;
+    absl::string_view string_value;
     int64_t int64_value;
     int32_t int32_value;
     uint64_t uint64_value;
@@ -247,20 +244,11 @@ class PROTOBUF_EXPORT MapKey {
     bool bool_value;
   } val_;
 
-  void SetType(FieldDescriptor::CppType type) {
-    if (type_ == type) return;
-    if (type_ == FieldDescriptor::CPPTYPE_STRING) {
-      val_.string_value.Destruct();
-    }
-    type_ = type;
-    if (type_ == FieldDescriptor::CPPTYPE_STRING) {
-      val_.string_value.DefaultConstruct();
-    }
-  }
+  void SetType(FieldDescriptor::CppType type) { type_ = type; }
 
   // type_ is 0 or a valid FieldDescriptor::CppType.
   // Use "CppType()" to indicate zero.
-  FieldDescriptor::CppType type_;
+  FieldDescriptor::CppType type_ = FieldDescriptor::CppType();
 };
 
 namespace internal {
@@ -692,8 +680,90 @@ bool AllAreInitialized(const TypeDefinedMapFieldBase<Key, T>& field) {
   return true;
 }
 
-class PROTOBUF_EXPORT DynamicMapField final
-    : public TypeDefinedMapFieldBase<MapKey, MapValueRef> {
+// Used by DynamicMapField for it's key type.
+//
+// This is a lite wrapper around `absl::variant`. We do not use the variant
+// directly to prevent accidental hashing or equality of the implicitly provided
+// operators.
+class PROTOBUF_EXPORT DynamicMapKey {
+ public:
+  DynamicMapKey() = default;
+  DynamicMapKey(const DynamicMapKey&) = default;
+  DynamicMapKey(DynamicMapKey&&) = default;
+  DynamicMapKey& operator=(const DynamicMapKey&) = default;
+  DynamicMapKey& operator=(DynamicMapKey&&) = default;
+
+  explicit DynamicMapKey(google::protobuf::MapKey map_key)
+      : variant_(FromMapKey(map_key)) {}
+
+  google::protobuf::MapKey ToMapKey() const ABSL_ATTRIBUTE_LIFETIME_BOUND;
+
+  VariantKey ToVariantKey() const ABSL_ATTRIBUTE_LIFETIME_BOUND;
+
+  bool IsString() const {
+    return absl::holds_alternative<std::string>(variant_);
+  }
+
+  friend void swap(DynamicMapKey& lhs, DynamicMapKey& rhs) noexcept {
+    using std::swap;
+    swap(lhs.variant_, rhs.variant_);
+  }
+
+ private:
+  using Variant =
+      absl::variant<bool, int32_t, int64_t, uint32_t, uint64_t, std::string>;
+
+  static Variant FromMapKey(google::protobuf::MapKey map_key);
+
+  Variant variant_;
+};
+
+template <>
+struct is_internal_map_key_type<DynamicMapKey> : std::true_type {};
+
+template <>
+struct RealKeyToVariantKey<DynamicMapKey> : public RealKeyToVariantKey<MapKey> {
+  // Bring in for heterogeneous lookups.
+  using RealKeyToVariantKey<MapKey>::operator();
+
+  VariantKey operator()(const DynamicMapKey& value) const {
+    return value.ToVariantKey();
+  }
+};
+
+template <>
+struct RealKeyToVariantKeyAlternative<DynamicMapKey>
+    : public RealKeyToVariantKeyAlternative<MapKey> {
+  using RealKeyToVariantKeyAlternative<MapKey>::operator();
+
+  VariantKey operator()(const DynamicMapKey& value) const {
+    return RealKeyToVariantKey<DynamicMapKey>{}(value);
+  }
+};
+
+template <>
+struct TransparentSupport<DynamicMapKey> {
+  using hash = absl::Hash<DynamicMapKey>;
+
+  template <typename T, typename U>
+  static bool Equals(T&& t, U&& u) {
+    return ToView(t) == ToView(u);
+  }
+
+  template <typename K>
+  using key_arg = K;
+
+  using ViewType = google::protobuf::MapKey;
+
+  static ViewType ToView(ViewType v) { return v; }
+
+  static ViewType ToView(const DynamicMapKey& v ABSL_ATTRIBUTE_LIFETIME_BOUND) {
+    return v.ToMapKey();
+  }
+};
+
+class DynamicMapField final
+    : public TypeDefinedMapFieldBase<DynamicMapKey, MapValueRef> {
  public:
   explicit DynamicMapField(const Message* default_entry);
   DynamicMapField(const Message* default_entry, Arena* arena);

src/google/protobuf/map_field_inl.h

@@ -58,6 +58,10 @@ inline absl::string_view UnwrapMapKeyImpl(const MapKey& map_key,
 inline const MapKey& UnwrapMapKeyImpl(const MapKey& map_key, const MapKey*) {
   return map_key;
 }
+inline const MapKey& UnwrapMapKeyImpl(const MapKey& map_key,
+                                      const DynamicMapKey*) {
+  return map_key;
+}
 
 template <typename T>
 decltype(auto) UnwrapMapKey(const MapKey& map_key) {
@@ -80,11 +84,14 @@ inline void SetMapKey(MapKey* map_key, uint64_t value) {
 inline void SetMapKey(MapKey* map_key, bool value) {
   map_key->SetBoolValue(value);
 }
-inline void SetMapKey(MapKey* map_key, const std::string& value) {
+inline void SetMapKey(MapKey* map_key, absl::string_view value) {
   map_key->SetStringValue(value);
 }
 inline void SetMapKey(MapKey* map_key, const MapKey& value) {
-  map_key->CopyFrom(value);
+  *map_key = value;
+}
+inline void SetMapKey(MapKey* map_key, const DynamicMapKey& value) {
+  *map_key = value.ToMapKey();
 }
 
 // ------------------------TypeDefinedMapFieldBase---------------

src/google/protobuf/map_test.inc

@@ -113,6 +113,8 @@ struct ConstructorTag {
 
   bool operator<(const ConstructorTag&) const { return false; }
 
+  bool operator==(const ConstructorTag&) const { return true; }
+
   ConstructorType invoked_constructor;
 };
 
@@ -315,16 +317,6 @@ TEST_F(MapImplTest, OperatorBracket) {
 }  // namespace protobuf
 }  // namespace google
 
-namespace std {
-
-template <>  // NOLINT
-struct hash<google::protobuf::internal::MoveTestKey> {
-  size_t operator()(const google::protobuf::internal::MoveTestKey& key) const {
-    return hash<uint64_t>{}(key.data);
-  }
-};
-}  // namespace std
-
 namespace google {
 namespace protobuf {
 namespace internal {
@@ -343,6 +335,38 @@ struct RealKeyToVariantKeyAlternative<MoveTestKey> {
   }
 };
 
+template <>
+struct TransparentSupport<MoveTestKey> {
+  using hash = absl::Hash<MoveTestKey>;
+
+  static bool Equals(const MoveTestKey& a, const MoveTestKey& b) {
+    return a == b;
+  }
+
+  template <typename K>
+  using key_arg = MoveTestKey;
+
+  using ViewType = const MoveTestKey&;
+
+  static ViewType ToView(ViewType v) { return v; }
+};
+
+template <>
+struct TransparentSupport<ConstructorTag> {
+  using hash = absl::Hash<ConstructorTag>;
+
+  static bool Equals(const ConstructorTag& a, const ConstructorTag& b) {
+    return a == b;
+  }
+
+  template <typename K>
+  using key_arg = ConstructorTag;
+
+  using ViewType = const ConstructorTag&;
+
+  static ViewType ToView(ViewType v) { return v; }
+};
+
 namespace {
 
 TEST_F(MapImplTest, OperatorBracketRValue) {