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
5 months ago · 671e6c5ef6
parent 9f1722b117
commit 671e6c5ef6
6 changed files with 222 additions and 64 deletions
--- a/src/google/protobuf/BUILD.bazel
+++ b/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",
    ],
 )

--- a/src/google/protobuf/map.h
+++ b/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);
  }
--- a/src/google/protobuf/map_field.cc
+++ b/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.
--- a/src/google/protobuf/map_field.h
+++ b/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);
--- a/src/google/protobuf/map_field_inl.h
+++ b/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---------------
--- a/src/google/protobuf/map_test.inc
+++ b/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) {