Move SpaceUsed calculation to UntypedMapBase to reduce code size.

src/google/protobuf/BUILD.bazel

@@ -569,6 +569,7 @@ cc_library(
         "@com_google_absl//absl/base:prefetch",
         "@com_google_absl//absl/container:btree",
         "@com_google_absl//absl/container:flat_hash_set",
+        "@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/map.cc

@@ -14,6 +14,7 @@
 #include <string>
 
 #include "absl/base/optimization.h"
+#include "absl/functional/overload.h"
 #include "absl/log/absl_check.h"
 #include "google/protobuf/arena.h"
 #include "google/protobuf/message_lite.h"
@@ -86,12 +87,27 @@ void UntypedMapBase::ClearTableImpl(bool reset, void (*destroy)(NodeBase*)) {
   }
 }
 
-size_t UntypedMapBase::SpaceUsedInTable(size_t sizeof_node) const {
+size_t UntypedMapBase::SpaceUsedExcludingSelfLong() const {
   size_t size = 0;
   // The size of the table.
   size += sizeof(void*) * num_buckets_;
   // All the nodes.
-  size += sizeof_node * num_elements_;
+  size += type_info_.node_size * num_elements_;
+  VisitAllNodes([&](auto* key, auto* value) {
+    const auto space_used = absl::Overload{
+        [](const std::string* str) -> size_t {
+          return StringSpaceUsedExcludingSelfLong(*str);
+        },
+        [&](const MessageLite* msg) -> size_t {
+          const auto* class_data = GetClassData(*msg);
+          if (class_data->is_lite) return 0;
+          return class_data->full().descriptor_methods->space_used_long(*msg) -
+                 class_data->allocation_size();
+        },
+        [](const void*) -> size_t { return 0; }};
+    size += space_used(key);
+    size += space_used(value);
+  });
   return size;
 }
 

src/google/protobuf/map.h

@@ -284,35 +284,9 @@ struct NodeBase {
   }
 };
 
-// This captures all numeric types.
-inline size_t MapValueSpaceUsedExcludingSelfLong(bool) { return 0; }
-inline size_t MapValueSpaceUsedExcludingSelfLong(const std::string& str) {
-  return StringSpaceUsedExcludingSelfLong(str);
-}
-template <typename T,
-          typename = decltype(std::declval<const T&>().SpaceUsedLong())>
-size_t MapValueSpaceUsedExcludingSelfLong(const T& message) {
-  return message.SpaceUsedLong() - sizeof(T);
-}
-
 constexpr size_t kGlobalEmptyTableSize = 1;
 PROTOBUF_EXPORT extern NodeBase* const kGlobalEmptyTable[kGlobalEmptyTableSize];
 
-template <typename Map,
-          typename = typename std::enable_if<
-              !std::is_scalar<typename Map::key_type>::value ||
-              !std::is_scalar<typename Map::mapped_type>::value>::type>
-size_t SpaceUsedInValues(const Map* map) {
-  size_t size = 0;
-  for (const auto& v : *map) {
-    size += internal::MapValueSpaceUsedExcludingSelfLong(v.first) +
-            internal::MapValueSpaceUsedExcludingSelfLong(v.second);
-  }
-  return size;
-}
-
-inline size_t SpaceUsedInValues(const void*) { return 0; }
-
 class UntypedMapBase;
 
 class UntypedMapIterator {
@@ -457,8 +431,19 @@ class PROTOBUF_EXPORT UntypedMapBase {
   UntypedMapBase(const UntypedMapBase&) = delete;
   UntypedMapBase& operator=(const UntypedMapBase&) = delete;
 
-  template <typename T = void>
+  template <typename T>
+  T* GetKey(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_.key_type));
+    return reinterpret_cast<T*>(node->GetVoidKey());
+  }
+
+  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);
   }
@@ -468,6 +453,9 @@ class PROTOBUF_EXPORT UntypedMapBase {
     ClearTableImpl(reset, destroy);
   }
 
+  // Space used for the table and nodes.
+  size_t SpaceUsedExcludingSelfLong() const;
+
  protected:
   // 16 bytes is the minimum useful size for the array cache in the arena.
   enum : map_index_t { kMinTableSize = 16 / sizeof(void*) };
@@ -495,6 +483,21 @@ class PROTOBUF_EXPORT UntypedMapBase {
   // All the end iterators are singletons anyway.
   static UntypedMapIterator EndIterator() { return {nullptr, nullptr, 0}; }
 
+  // Calls `f(k)` with the key of the node, where `k` is the appropriate type
+  // according to the stored TypeInfo.
+  template <typename F>
+  auto VisitKey(NodeBase* node, F f) const;
+
+  // Calls `f(v)` with the value of the node, where `v` is the appropriate type
+  // according to the stored TypeInfo.
+  // Messages are visited as `MessageLite`, and enums are visited as int32.
+  template <typename F>
+  auto VisitValue(NodeBase* node, F f) const;
+
+  // As above, but calls `f(k, v)` for every node in the map.
+  template <typename F>
+  void VisitAllNodes(F f) const;
+
  protected:
   friend class TcParser;
   friend struct MapTestPeer;
@@ -503,6 +506,13 @@ class PROTOBUF_EXPORT UntypedMapBase {
   friend class RustMapHelper;
   friend class v2::TableDriven;
 
+  // Calls `f(type_t)` where `type_t` is an unspecified type that has a `::type`
+  // typedef in it representing the dynamic type of key/value of the node.
+  template <typename F>
+  auto VisitKeyType(F f) const;
+  template <typename F>
+  auto VisitValueType(F f) const;
+
   struct NodeAndBucket {
     NodeBase* node;
     map_index_t bucket;
@@ -583,10 +593,6 @@ class PROTOBUF_EXPORT UntypedMapBase {
                : nullptr;
   }
 
-  // Space used for the table, trees, and nodes.
-  // Does not include the indirect space used. Eg the data of a std::string.
-  size_t SpaceUsedInTable(size_t sizeof_node) const;
-
   map_index_t num_elements_;
   map_index_t num_buckets_;
   map_index_t index_of_first_non_null_;
@@ -595,6 +601,77 @@ class PROTOBUF_EXPORT UntypedMapBase {
   Allocator alloc_;
 };
 
+template <typename F>
+auto UntypedMapBase::VisitKeyType(F f) const {
+  switch (type_info_.key_type) {
+    case TypeKind::kBool:
+      return f(std::enable_if<true, bool>{});
+    case TypeKind::kU32:
+      return f(std::enable_if<true, uint32_t>{});
+    case TypeKind::kU64:
+      return f(std::enable_if<true, uint64_t>{});
+    case TypeKind::kString:
+      return f(std::enable_if<true, std::string>{});
+
+    case TypeKind::kFloat:
+    case TypeKind::kDouble:
+    case TypeKind::kMessage:
+    case TypeKind::kUnknown:
+    default:
+      Unreachable();
+  }
+}
+
+template <typename F>
+auto UntypedMapBase::VisitValueType(F f) const {
+  switch (type_info_.value_type) {
+    case TypeKind::kBool:
+      return f(std::enable_if<true, bool>{});
+    case TypeKind::kU32:
+      return f(std::enable_if<true, uint32_t>{});
+    case TypeKind::kU64:
+      return f(std::enable_if<true, uint64_t>{});
+    case TypeKind::kFloat:
+      return f(std::enable_if<true, float>{});
+    case TypeKind::kDouble:
+      return f(std::enable_if<true, double>{});
+    case TypeKind::kString:
+      return f(std::enable_if<true, std::string>{});
+    case TypeKind::kMessage:
+      return f(std::enable_if<true, MessageLite>{});
+
+    case TypeKind::kUnknown:
+    default:
+      Unreachable();
+  }
+}
+
+template <typename F>
+void UntypedMapBase::VisitAllNodes(F f) const {
+  VisitKeyType([&](auto key_type) {
+    VisitValueType([&](auto value_type) {
+      for (auto it = begin(); !it.Equals(EndIterator()); it.PlusPlus()) {
+        f(GetKey<typename decltype(key_type)::type>(it.node_),
+          GetValue<typename decltype(value_type)::type>(it.node_));
+      }
+    });
+  });
+}
+
+template <typename F>
+auto UntypedMapBase::VisitKey(NodeBase* node, F f) const {
+  return VisitKeyType([&](auto key_type) {
+    return f(GetKey<typename decltype(key_type)::type>(node));
+  });
+}
+
+template <typename F>
+auto UntypedMapBase::VisitValue(NodeBase* node, F f) const {
+  return VisitValueType([&](auto value_type) {
+    return f(GetValue<typename decltype(value_type)::type>(node));
+  });
+}
+
 inline UntypedMapIterator UntypedMapBase::begin() const {
   map_index_t bucket_index;
   NodeBase* node;
@@ -1389,7 +1466,7 @@ class Map : private internal::KeyMapBase<internal::KeyForBase<Key>> {
 
   size_t SpaceUsedExcludingSelfLong() const {
     if (empty()) return 0;
-    return SpaceUsedInternal() + internal::SpaceUsedInValues(this);
+    return internal::UntypedMapBase::SpaceUsedExcludingSelfLong();
   }
 
   static constexpr size_t InternalGetArenaOffset(internal::InternalVisibility) {
@@ -1425,10 +1502,6 @@ class Map : private internal::KeyMapBase<internal::KeyForBase<Key>> {
     }
   }
 
-  size_t SpaceUsedInternal() const {
-    return this->SpaceUsedInTable(sizeof(Node));
-  }
-
   template <typename K, typename... Args>
   std::pair<iterator, bool> TryEmplaceInternal(K&& k, Args&&... args) {
     auto p = this->FindHelper(TS::ToView(k));

src/google/protobuf/map_field.cc

@@ -34,6 +34,10 @@ MapFieldBase::~MapFieldBase() {
   delete maybe_payload();
 }
 
+size_t MapFieldBase::SpaceUsedExcludingSelfNoLockImpl(const MapFieldBase& map) {
+  return map.GetMapRaw().SpaceUsedExcludingSelfLong();
+}
+
 const UntypedMapBase& MapFieldBase::GetMapImpl(const MapFieldBaseForParse& map,
                                                bool is_mutable) {
   const auto& self = static_cast<const MapFieldBase&>(map);

src/google/protobuf/map_field.h

@@ -417,6 +417,7 @@ class PROTOBUF_EXPORT MapFieldBase : public MapFieldBaseForParse {
 
   static void SwapImpl(MapFieldBase& lhs, MapFieldBase& rhs);
   static void UnsafeShallowSwapImpl(MapFieldBase& lhs, MapFieldBase& rhs);
+  static size_t SpaceUsedExcludingSelfNoLockImpl(const MapFieldBase& map);
 
   // Tells MapFieldBase that there is new change to Map.
   void SetMapDirty() {
@@ -629,8 +630,6 @@ class TypeDefinedMapFieldBase : public MapFieldBase {
   static void SwapImpl(MapFieldBase& lhs, MapFieldBase& rhs);
   static void UnsafeShallowSwapImpl(MapFieldBase& lhs, MapFieldBase& rhs);
 
-  static size_t SpaceUsedExcludingSelfNoLockImpl(const MapFieldBase& map);
-
   // map_ is inside an anonymous union so we can explicitly control its
   // destruction
   union {

src/google/protobuf/map_field_inl.h

@@ -152,18 +152,6 @@ void TypeDefinedMapFieldBase<Key, T>::MergeFromImpl(MapFieldBase& base,
   self.SetMapDirty();
 }
 
-template <typename Key, typename T>
-size_t TypeDefinedMapFieldBase<Key, T>::SpaceUsedExcludingSelfNoLockImpl(
-    const MapFieldBase& map) {
-  // We can't compile this expression for DynamicMapField even though it is
-  // never used at runtime, so disable it at compile time.
-  if constexpr (!std::is_same<Map<Key, T>, Map<MapKey, MapValueRef>>::value) {
-    return static_cast<const TypeDefinedMapFieldBase&>(map)
-        .map_.SpaceUsedExcludingSelfLong();
-  }
-  return 0;
-}
-
 template <typename Key, typename T>
 void TypeDefinedMapFieldBase<Key, T>::UnsafeShallowSwapImpl(MapFieldBase& lhs,
                                                             MapFieldBase& rhs) {