From bfe46c5163ff00f4749f6ff45beef4093e556ff0 Mon Sep 17 00:00:00 2001
From: Protobuf Team Bot <protobuf-github-bot@google.com>
Date: Thu, 5 Dec 2024 13:16:56 -0800
Subject: [PATCH] Speed up node initialization in MpMap.

PiperOrigin-RevId: 703225930
---
 .../generated_message_tctable_lite.cc         | 47 ++++++++++---------
 src/google/protobuf/map.h                     | 13 +++++
 2 files changed, 38 insertions(+), 22 deletions(-)

diff --git a/src/google/protobuf/generated_message_tctable_lite.cc b/src/google/protobuf/generated_message_tctable_lite.cc
index ed01ff9b7a88..207c4b8f34e8 100644
--- a/src/google/protobuf/generated_message_tctable_lite.cc
+++ b/src/google/protobuf/generated_message_tctable_lite.cc
@@ -2826,28 +2826,30 @@ PROTOBUF_NOINLINE const char* TcParser::MpMap(PROTOBUF_TC_PARAM_DECL) {
 
   while (true) {
     NodeBase* node = map.AllocNode();
+    char* const node_end =
+        reinterpret_cast<char*>(node) + map.type_info().node_size;
+    void* const node_key = node->GetVoidKey();
+
+    // Due to node alignment we can guarantee that we have at least 8 writable
+    // bytes from the key position to the end of the node.
+    // We can initialize the first and last 8 bytes, which takes care of all the
+    // scalar value types. This makes the VisitXXX calls below faster because
+    // the switch is much smaller.
+    // Assert this in debug mode, just in case.
+    ABSL_DCHECK_GE(node_end - static_cast<char*>(node_key), sizeof(uint64_t));
+    memset(node_key, 0, sizeof(uint64_t));
+    memset(node_end - sizeof(uint64_t), 0, sizeof(uint64_t));
+
+    map.VisitKey(  //
+        node, absl::Overload{
+                  [&](std::string* str) {
+                    Arena::CreateInArenaStorage(str, map.arena());
+                  },
+                  // Already initialized above. Do nothing here.
+                  [](void*) {},
+              });
 
-    map.VisitKey(node,  //
-                 absl::Overload{
-                     [&](std::string* str) {
-                       Arena::CreateInArenaStorage(str, map.arena());
-                     },
-                     [&](void* scalar) {
-                       // Due to node alignment we can guarantee that we have at
-                       // least 8 writable bytes at the key position (as long as
-                       // we do it before we initialize the value). We can
-                       // unconditionally write here.
-                       // Assert this in debug mode, just in case.
-                       ABSL_DCHECK_GE(
-                           reinterpret_cast<char*>(node) +
-                               map.type_info().node_size -
-                               reinterpret_cast<char*>(node->GetVoidKey()),
-                           sizeof(uint64_t));
-                       memset(node->GetVoidKey(), 0, sizeof(uint64_t));
-                     },
-                 });
-
-    map.VisitValue(
+    map.VisitValue(  //
         node, absl::Overload{
                   [&](std::string* str) {
                     Arena::CreateInArenaStorage(str, map.arena());
@@ -2855,7 +2857,8 @@ PROTOBUF_NOINLINE const char* TcParser::MpMap(PROTOBUF_TC_PARAM_DECL) {
                   [&](MessageLite* msg) {
                     aux[1].table->class_data->PlacementNew(msg, map.arena());
                   },
-                  [](auto* scalar) { memset(scalar, 0, sizeof(*scalar)); },
+                  // Already initialized above. Do nothing here.
+                  [](void*) {},
               });
 
     ptr = ctx->ParseLengthDelimitedInlined(ptr, [&](const char* ptr) {
diff --git a/src/google/protobuf/map.h b/src/google/protobuf/map.h
index 26283f4c71a7..f9a8cccb8902 100644
--- a/src/google/protobuf/map.h
+++ b/src/google/protobuf/map.h
@@ -1013,6 +1013,19 @@ class Map : private internal::KeyMapBase<internal::KeyForBase<Key>> {
     static_assert(
         sizeof(Map) == sizeof(internal::UntypedMapBase),
         "Map must not have any data members beyond what is in UntypedMapBase.");
+
+    // Check for MpMap optimizations.
+    if constexpr (std::is_scalar_v<key_type>) {
+      static_assert(sizeof(key_type) <= sizeof(uint64_t),
+                    "Scalar must be <= than uint64_t");
+    }
+    if constexpr (std::is_scalar_v<mapped_type>) {
+      static_assert(sizeof(mapped_type) <= sizeof(uint64_t),
+                    "Scalar must be <= than uint64_t");
+    }
+    static_assert(internal::kMaxMessageAlignment >= sizeof(uint64_t));
+    static_assert(sizeof(Node) - sizeof(internal::NodeBase) >= sizeof(uint64_t),
+                  "We must have at least this bytes for MpMap initialization");
   }
 
   template <typename P>