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: 712658107

src/google/protobuf/map.cc

@@ -10,7 +10,6 @@
 #include <algorithm>
 #include <cstddef>
 #include <cstdint>
-#include <iterator>
 #include <string>
 
 #include "absl/base/optimization.h"
@@ -30,6 +29,78 @@ namespace internal {
 
 NodeBase* const kGlobalEmptyTable[kGlobalEmptyTableSize] = {};
 
+void UntypedMapBase::UntypedMergeFrom(const UntypedMapBase& other) {
+  if (other.empty()) return;
+
+  // Do the merging in steps to avoid Key*Value number of instantiations and
+  // reduce code duplication per instantation.
+  NodeBase* nodes = nullptr;
+
+  // First, allocate all the nodes without types.
+  for (size_t i = 0; i < other.num_elements_; ++i) {
+    NodeBase* new_node = AllocNode();
+    new_node->next = nodes;
+    nodes = new_node;
+  }
+
+  // Then, copy the values.
+  VisitValueType([&](auto value_type) {
+    using Value = typename decltype(value_type)::type;
+    NodeBase* out_node = nodes;
+
+    // Get the ClassData once to avoid redundant virtual function calls.
+    const internal::ClassData* class_data =
+        std::is_same_v<MessageLite, Value>
+            ? GetClassData(*other.GetValue<MessageLite>(other.begin().node_))
+            : nullptr;
+
+    for (auto it = other.begin(); !it.Equals(EndIterator()); it.PlusPlus()) {
+      Value* out = GetValue<Value>(out_node);
+      out_node = out_node->next;
+      auto& in = *other.GetValue<Value>(it.node_);
+      if constexpr (std::is_same_v<MessageLite, Value>) {
+        class_data->PlacementNew(out, arena())->CheckTypeAndMergeFrom(in);
+      } else {
+        Arena::CreateInArenaStorage(out, this->arena_, in);
+      }
+    }
+  });
+
+  // Finally, copy the keys and insert the nodes.
+  VisitKeyType([&](auto key_type) {
+    using Key = typename decltype(key_type)::type;
+    for (auto it = other.begin(); !it.Equals(EndIterator()); it.PlusPlus()) {
+      NodeBase* node = nodes;
+      nodes = nodes->next;
+      const Key& in = *other.GetKey<Key>(it.node_);
+      Key* out = GetKey<Key>(node);
+      if (!internal::InitializeMapKey(out, in, this->arena_)) {
+        Arena::CreateInArenaStorage(out, this->arena_, in);
+      }
+
+      static_cast<KeyMapBase<Key>*>(this)->InsertOrReplaceNode(
+          static_cast<typename KeyMapBase<Key>::KeyNode*>(node));
+    }
+  });
+}
+
+void UntypedMapBase::UntypedSwap(UntypedMapBase& other) {
+  if (arena() == other.arena()) {
+    InternalSwap(&other);
+  } else {
+    UntypedMapBase tmp(arena_, type_info_);
+    InternalSwap(&tmp);
+
+    ABSL_DCHECK(empty());
+    UntypedMergeFrom(other);
+
+    other.ClearTable(true, nullptr);
+    other.UntypedMergeFrom(tmp);
+
+    if (arena_ == nullptr) tmp.ClearTable(false, nullptr);
+  }
+}
+
 void UntypedMapBase::DeleteNode(NodeBase* node) {
   const auto destroy = absl::Overload{
       [](std::string* str) { str->~basic_string(); },

src/google/protobuf/map.h

@@ -379,6 +379,9 @@ class PROTOBUF_EXPORT UntypedMapBase {
     std::swap(arena_, other->arena_);
   }
 
+  void UntypedMergeFrom(const UntypedMapBase& other);
+  void UntypedSwap(UntypedMapBase& other);
+
   static size_type max_size() {
     return std::numeric_limits<map_index_t>::max();
   }
@@ -674,6 +677,7 @@ class KeyMapBase : public UntypedMapBase {
   using KeyNode = internal::KeyNode<Key>;
 
  protected:
+  friend UntypedMapBase;
   friend class MapFieldBase;
   friend class TcParser;
   friend struct MapTestPeer;

src/google/protobuf/map_field.cc

@@ -36,6 +36,11 @@ MapFieldBase::~MapFieldBase() {
   delete maybe_payload();
 }
 
+void MapFieldBase::SwapImpl(MapFieldBase& lhs, MapFieldBase& rhs) {
+  MapFieldBase::SwapPayload(lhs, rhs);
+  lhs.GetMapRaw().UntypedSwap(rhs.GetMapRaw());
+}
+
 template <typename Map, typename F>
 auto VisitMapKey(const MapKey& map_key, Map& map, F f) {
   switch (map_key.type()) {

src/google/protobuf/map_field.h

@@ -513,6 +513,7 @@ class PROTOBUF_EXPORT MapFieldBase : public MapFieldBaseForParse {
                                                const MapKey& map_key,
                                                MapValueRef* val);
   static bool DeleteMapValueImpl(MapFieldBase& self, const MapKey& map_key);
+  static void SwapImpl(MapFieldBase& lhs, MapFieldBase& rhs);
 
  private:
   friend class ContendedMapCleanTest;
@@ -640,7 +641,6 @@ class TypeDefinedMapFieldBase : public MapFieldBase {
   using Iter = typename Map<Key, T>::const_iterator;
 
   static void MergeFromImpl(MapFieldBase& base, const MapFieldBase& other);
-  static void SwapImpl(MapFieldBase& lhs, MapFieldBase& rhs);
 
   // map_ is inside an anonymous union so we can explicitly control its
   // destruction

src/google/protobuf/map_field_inl.h

@@ -34,14 +34,6 @@ namespace google {
 namespace protobuf {
 namespace internal {
 // ------------------------TypeDefinedMapFieldBase---------------
-template <typename Key, typename T>
-void TypeDefinedMapFieldBase<Key, T>::SwapImpl(MapFieldBase& lhs,
-                                               MapFieldBase& rhs) {
-  MapFieldBase::SwapPayload(lhs, rhs);
-  static_cast<TypeDefinedMapFieldBase&>(lhs).map_.swap(
-      static_cast<TypeDefinedMapFieldBase&>(rhs).map_);
-}
-
 template <typename Key, typename T>
 void TypeDefinedMapFieldBase<Key, T>::MergeFromImpl(MapFieldBase& base,
                                                     const MapFieldBase& other) {