Move part of the `TypeDefinedMapFieldBase` implementation into `MapFieldBase` implemented on top of `UntypedMapBase` visitation.

Reduces code duplication in large binaries.

More to come in future changes.

PiperOrigin-RevId: 707269328

src/google/protobuf/BUILD.bazel

@@ -672,6 +672,7 @@ cc_library(
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/functional:any_invocable",
         "@com_google_absl//absl/functional:function_ref",
+        "@com_google_absl//absl/functional:overload",
         "@com_google_absl//absl/hash",
         "@com_google_absl//absl/log:absl_check",
         "@com_google_absl//absl/log:absl_log",

src/google/protobuf/dynamic_message.cc

@@ -123,11 +123,6 @@ class DynamicMapKey {
   Variant variant_;
 };
 
-// The other overloads for SetMapKey are located in map_field_inl.h
-inline void SetMapKey(MapKey* map_key, const DynamicMapKey& value) {
-  *map_key = value.ToMapKey();
-}
-
 template <>
 struct is_internal_map_key_type<DynamicMapKey> : std::true_type {};
 
@@ -218,6 +213,9 @@ class DynamicMapField final
                                                MapValueRef* val);
   static void ClearMapNoSyncImpl(MapFieldBase& base);
   static bool DeleteMapValueImpl(MapFieldBase& map, const MapKey& map_key);
+  static void SetMapIteratorValueImpl(MapIterator* map_iter);
+  static bool LookupMapValueImpl(const MapFieldBase& self,
+                                 const MapKey& map_key, MapValueConstRef* val);
 
   static void UnsafeShallowSwapImpl(MapFieldBase& lhs, MapFieldBase& rhs) {
     static_cast<DynamicMapField&>(lhs).Swap(
@@ -273,6 +271,27 @@ bool DynamicMapField::DeleteMapValueImpl(MapFieldBase& base,
   return true;
 }
 
+void DynamicMapField::SetMapIteratorValueImpl(MapIterator* map_iter) {
+  if (map_iter->iter_.Equals(UntypedMapBase::EndIterator())) return;
+  auto iter = typename decltype(map_)::const_iterator(map_iter->iter_);
+  map_iter->key_ = iter->first.ToMapKey();
+  map_iter->value_.CopyFrom(iter->second);
+}
+
+bool DynamicMapField::LookupMapValueImpl(const MapFieldBase& self,
+                                         const MapKey& map_key,
+                                         MapValueConstRef* val) {
+  const auto& map = static_cast<const DynamicMapField&>(self).GetMap();
+  auto iter = map.find(map_key);
+  if (map.end() == iter) {
+    return false;
+  }
+  if (val != nullptr) {
+    val->CopyFrom(iter->second);
+  }
+  return true;
+}
+
 void DynamicMapField::AllocateMapValue(MapValueRef* map_val) {
   const FieldDescriptor* val_des = default_entry_->GetDescriptor()->map_value();
   map_val->SetType(val_des->cpp_type());

src/google/protobuf/map.h

@@ -73,6 +73,7 @@ class TableDriven;
 
 template <typename Key, typename T>
 class MapFieldLite;
+class MapFieldBase;
 
 template <typename Derived, typename Key, typename T,
           WireFormatLite::FieldType key_wire_type,
@@ -339,13 +340,16 @@ class PROTOBUF_EXPORT UntypedMapBase {
     return reinterpret_cast<T*>(node->GetVoidKey());
   }
 
+  void* GetVoidValue(NodeBase* node) const {
+    return reinterpret_cast<char*>(node) + type_info_.value_offset;
+  }
+
   template <typename T>
   T* GetValue(NodeBase* node) const {
     // Debug check that `T` matches what we expect from the type info.
     ABSL_DCHECK_EQ(static_cast<int>(StaticTypeKind<T>()),
                    static_cast<int>(type_info_.value_type));
-    return reinterpret_cast<T*>(reinterpret_cast<char*>(node) +
-                                type_info_.value_offset);
+    return reinterpret_cast<T*>(GetVoidValue(node));
   }
 
   void ClearTable(bool reset, void (*destroy)(NodeBase*)) {
@@ -401,6 +405,7 @@ class PROTOBUF_EXPORT UntypedMapBase {
   void VisitAllNodes(F f) const;
 
  protected:
+  friend class MapFieldBase;
   friend class TcParser;
   friend struct MapTestPeer;
   friend struct MapBenchmarkPeer;
@@ -668,6 +673,7 @@ class KeyMapBase : public UntypedMapBase {
   using KeyNode = internal::KeyNode<Key>;
 
  protected:
+  friend class MapFieldBase;
   friend class TcParser;
   friend struct MapTestPeer;
   friend struct MapBenchmarkPeer;

src/google/protobuf/map_field.cc

@@ -12,6 +12,7 @@
 #include <cstdint>
 #include <string>
 
+#include "absl/functional/overload.h"
 #include "absl/log/absl_check.h"
 #include "absl/synchronization/mutex.h"
 #include "google/protobuf/arena.h"
@@ -34,6 +35,58 @@ MapFieldBase::~MapFieldBase() {
   delete maybe_payload();
 }
 
+void MapFieldBase::ClearMapNoSyncImpl(MapFieldBase& self) {
+  self.GetMapRaw().ClearTable(true, nullptr);
+}
+
+void MapFieldBase::SetMapIteratorValueImpl(MapIterator* map_iter) {
+  if (map_iter->iter_.Equals(UntypedMapBase::EndIterator())) return;
+
+  const UntypedMapBase& map = *map_iter->iter_.m_;
+  NodeBase* node = map_iter->iter_.node_;
+  auto& key = map_iter->key_;
+  map.VisitKey(node,
+               absl::Overload{
+                   [&](const std::string* v) { key.val_.string_value = *v; },
+                   [&](const auto* v) {
+                     // Memcpy the scalar into the union.
+                     memcpy(static_cast<void*>(&key.val_), v, sizeof(*v));
+                   },
+               });
+  map_iter->value_.SetValue(map.GetVoidValue(node));
+}
+
+bool MapFieldBase::LookupMapValueImpl(const MapFieldBase& self,
+                                      const MapKey& map_key,
+                                      MapValueConstRef* val) {
+  auto& map = self.GetMap();
+  if (map.empty()) return false;
+
+  switch (map_key.type()) {
+#define HANDLE_TYPE(CPPTYPE, Type, KeyBaseType)                       \
+  case FieldDescriptor::CPPTYPE_##CPPTYPE: {                          \
+    auto& key_map = static_cast<const KeyMapBase<KeyBaseType>&>(map); \
+    auto res = key_map.FindHelper(map_key.Get##Type##Value());        \
+    if (res.node == nullptr) {                                        \
+      return false;                                                   \
+    }                                                                 \
+    if (val != nullptr) {                                             \
+      val->SetValue(map.GetVoidValue(res.node));                      \
+    }                                                                 \
+    return true;                                                      \
+  }
+    HANDLE_TYPE(INT32, Int32, uint32_t);
+    HANDLE_TYPE(UINT32, UInt32, uint32_t);
+    HANDLE_TYPE(INT64, Int64, uint64_t);
+    HANDLE_TYPE(UINT64, UInt64, uint64_t);
+    HANDLE_TYPE(BOOL, Bool, bool);
+    HANDLE_TYPE(STRING, String, std::string);
+#undef HANDLE_TYPE
+    default:
+      Unreachable();
+  }
+}
+
 size_t MapFieldBase::SpaceUsedExcludingSelfNoLockImpl(const MapFieldBase& map) {
   return map.GetMapRaw().SpaceUsedExcludingSelfLong();
 }

src/google/protobuf/map_field.h

@@ -51,10 +51,6 @@ 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
@@ -505,6 +501,10 @@ class PROTOBUF_EXPORT MapFieldBase : public MapFieldBaseForParse {
 
   static const UntypedMapBase& GetMapImpl(const MapFieldBaseForParse& map,
                                           bool is_mutable);
+  static void ClearMapNoSyncImpl(MapFieldBase& self);
+  static void SetMapIteratorValueImpl(MapIterator* map_iter);
+  static bool LookupMapValueImpl(const MapFieldBase& self,
+                                 const MapKey& map_key, MapValueConstRef* val);
 
  private:
   friend class ContendedMapCleanTest;
@@ -601,10 +601,6 @@ class TypeDefinedMapFieldBase : public MapFieldBase {
     return &map_;
   }
 
-  static void ClearMapNoSyncImpl(MapFieldBase& map) {
-    static_cast<TypeDefinedMapFieldBase&>(map).map_.clear();
-  }
-
   void InternalSwap(TypeDefinedMapFieldBase* other);
 
   static constexpr size_t InternalGetArenaOffsetAlt(
@@ -619,9 +615,6 @@ class TypeDefinedMapFieldBase : public MapFieldBase {
   using Iter = typename Map<Key, T>::const_iterator;
 
   static bool DeleteMapValueImpl(MapFieldBase& map, const MapKey& map_key);
-  static bool LookupMapValueImpl(const MapFieldBase& self,
-                                 const MapKey& map_key, MapValueConstRef* val);
-  static void SetMapIteratorValueImpl(MapIterator* map_iter);
   static bool InsertOrLookupMapValueNoSyncImpl(MapFieldBase& map,
                                                const MapKey& map_key,
                                                MapValueRef* val);

src/google/protobuf/map_field_inl.h

@@ -55,54 +55,13 @@ inline absl::string_view UnwrapMapKeyImpl(const MapKey& map_key,
                                           const std::string*) {
   return map_key.GetStringValue();
 }
-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) {
   return UnwrapMapKeyImpl(map_key, static_cast<T*>(nullptr));
 }
 
-// SetMapKey
-inline void SetMapKey(MapKey* map_key, int32_t value) {
-  map_key->SetInt32Value(value);
-}
-inline void SetMapKey(MapKey* map_key, uint32_t value) {
-  map_key->SetUInt32Value(value);
-}
-inline void SetMapKey(MapKey* map_key, int64_t value) {
-  map_key->SetInt64Value(value);
-}
-inline void SetMapKey(MapKey* map_key, uint64_t value) {
-  map_key->SetUInt64Value(value);
-}
-inline void SetMapKey(MapKey* map_key, bool value) {
-  map_key->SetBoolValue(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 = value;
-}
-// The overload of SetMapKey for DynamicMapKey is located in dynamic_message.cc
-// and is discovered via ADL.
-
 // ------------------------TypeDefinedMapFieldBase---------------
-template <typename Key, typename T>
-void TypeDefinedMapFieldBase<Key, T>::SetMapIteratorValueImpl(
-    MapIterator* map_iter) {
-  if (map_iter->iter_.Equals(UntypedMapBase::EndIterator())) return;
-  auto iter = typename Map<Key, T>::const_iterator(map_iter->iter_);
-  SetMapKey(&map_iter->key_, iter->first);
-  map_iter->value_.SetValueOrCopy(&iter->second);
-}
-
 template <typename Key, typename T>
 bool TypeDefinedMapFieldBase<Key, T>::InsertOrLookupMapValueNoSyncImpl(
     MapFieldBase& map, const MapKey& map_key, MapValueRef* val) {
@@ -112,20 +71,6 @@ bool TypeDefinedMapFieldBase<Key, T>::InsertOrLookupMapValueNoSyncImpl(
   return res.second;
 }
 
-template <typename Key, typename T>
-bool TypeDefinedMapFieldBase<Key, T>::LookupMapValueImpl(
-    const MapFieldBase& self, const MapKey& map_key, MapValueConstRef* val) {
-  const auto& map = static_cast<const TypeDefinedMapFieldBase&>(self).GetMap();
-  auto iter = map.find(UnwrapMapKey<Key>(map_key));
-  if (map.end() == iter) {
-    return false;
-  }
-  if (val != nullptr) {
-    val->SetValueOrCopy(&iter->second);
-  }
-  return true;
-}
-
 template <typename Key, typename T>
 bool TypeDefinedMapFieldBase<Key, T>::DeleteMapValueImpl(
     MapFieldBase& map, const MapKey& map_key) {