Completely remove the VTable from MapFieldBase.

The sync callback is replaced with a single global injection point. We don't need it to exist per instance. It still needs to be injected to prevent a direct dependency from the parser into the reflection runtime.

The GetPrototype function is replaced with the prototype itself. We don't need a function to get the prototype when we can much more cheaply just store the prototype.

Also, move the check for GetMap/MutableMap to be inlined to avoid a regression in the table-driven parser.

PiperOrigin-RevId: 722656293

src/google/protobuf/compiler/cpp/message.cc

@@ -1514,9 +1514,8 @@ void MessageGenerator::GenerateMapEntryClassDefinition(io::Printer* p) {
           template <typename = void>
           explicit PROTOBUF_CONSTEXPR $classname$($pbi$::ConstantInitialized);
           explicit $classname$($pb$::Arena* $nullable$ arena);
-          static const $classname$* $nonnull$ internal_default_instance() {
-            return reinterpret_cast<const $classname$*>(
-                &_$classname$_default_instance_);
+          static constexpr const void* $nonnull$ internal_default_instance() {
+            return &_$classname$_default_instance_;
           }
 
           $decl_verify_func$;

src/google/protobuf/dynamic_message.cc

@@ -103,12 +103,6 @@ class DynamicMapField final : public MapFieldBase {
 
   // Must be first for GetMapRaw to work.
   UntypedMapBase map_;
-
-  const Message* default_entry_;
-
-  static const VTable kVTable;
-
-  static const Message* GetPrototypeImpl(const MapFieldBase& map);
 };
 
 static UntypedMapBase::TypeKind CppTypeToTypeKind(
@@ -153,23 +147,20 @@ static auto DefaultEntryToTypeInfo(
 DynamicMapField::DynamicMapField(const Message* default_entry,
                                  const Message* mapped_default_entry_if_message,
                                  Arena* arena)
-    : MapFieldBase(&kVTable, arena),
+    : MapFieldBase(default_entry, arena),
       map_(arena, DefaultEntryToTypeInfo(default_entry,
-                                         mapped_default_entry_if_message)),
-      default_entry_(default_entry) {}
-
-constexpr DynamicMapField::VTable DynamicMapField::kVTable =
-    MakeVTable<DynamicMapField>();
+                                         mapped_default_entry_if_message)) {
+  // This invariant is required by `GetMapRaw` to easily access the map
+  // member without paying for dynamic dispatch.
+  static_assert(MapFieldBaseForParse::MapOffset() ==
+                PROTOBUF_FIELD_OFFSET(DynamicMapField, map_));
+}
 
 DynamicMapField::~DynamicMapField() {
   ABSL_DCHECK_EQ(arena(), nullptr);
   map_.ClearTable(false);
 }
 
-const Message* DynamicMapField::GetPrototypeImpl(const MapFieldBase& map) {
-  return static_cast<const DynamicMapField&>(map).default_entry_;
-}
-
 }  // namespace internal
 
 using internal::DynamicMapField;

src/google/protobuf/map.cc

@@ -8,6 +8,7 @@
 #include "google/protobuf/map.h"
 
 #include <algorithm>
+#include <atomic>
 #include <cstddef>
 #include <cstdint>
 #include <string>
@@ -27,6 +28,9 @@ namespace google {
 namespace protobuf {
 namespace internal {
 
+std::atomic<MapFieldBaseForParse::SyncFunc>
+    MapFieldBaseForParse::sync_map_with_repeated{};
+
 NodeBase* const kGlobalEmptyTable[kGlobalEmptyTableSize] = {};
 
 void UntypedMapBase::UntypedMergeFrom(const UntypedMapBase& other) {

src/google/protobuf/map.h

@@ -15,6 +15,7 @@
 #define GOOGLE_PROTOBUF_MAP_H__
 
 #include <algorithm>
+#include <atomic>
 #include <cstddef>
 #include <cstdint>
 #include <cstring>
@@ -594,22 +595,65 @@ inline void UntypedMapIterator::PlusPlus() {
 class MapFieldBaseForParse {
  public:
   const UntypedMapBase& GetMap() const {
-    return vtable_->get_map(*this, false);
+    const auto p = payload_.load(std::memory_order_acquire);
+    // If this instance has a payload, then it might need sync'n.
+    if (ABSL_PREDICT_FALSE(IsPayload(p))) {
+      sync_map_with_repeated.load(std::memory_order_relaxed)(*this, false);
+    }
+    return GetMapRaw();
   }
+
   UntypedMapBase* MutableMap() {
-    return &const_cast<UntypedMapBase&>(vtable_->get_map(*this, true));
+    const auto p = payload_.load(std::memory_order_acquire);
+    // If this instance has a payload, then it might need sync'n.
+    if (ABSL_PREDICT_FALSE(IsPayload(p))) {
+      sync_map_with_repeated.load(std::memory_order_relaxed)(*this, true);
+    }
+    return &GetMapRaw();
   }
 
  protected:
-  struct VTable {
-    const UntypedMapBase& (*get_map)(const MapFieldBaseForParse&,
-                                     bool is_mutable);
-  };
-  explicit constexpr MapFieldBaseForParse(const VTable* vtable)
-      : vtable_(vtable) {}
+  static constexpr size_t MapOffset() { return sizeof(MapFieldBaseForParse); }
+
+  // See assertion in TypeDefinedMapFieldBase::TypeDefinedMapFieldBase()
+  const UntypedMapBase& GetMapRaw() const {
+    return *reinterpret_cast<const UntypedMapBase*>(
+        reinterpret_cast<const char*>(this) + MapOffset());
+  }
+  UntypedMapBase& GetMapRaw() {
+    return *reinterpret_cast<UntypedMapBase*>(reinterpret_cast<char*>(this) +
+                                              MapOffset());
+  }
+
+  // Injected from map_field.cc once we need to use it.
+  // We can't have a strong dep on it because it would cause protobuf_lite to
+  // depend on reflection.
+  using SyncFunc = void (*)(const MapFieldBaseForParse&, bool is_mutable);
+  static std::atomic<SyncFunc> sync_map_with_repeated;
+
+  // The prototype is a `Message`, but due to restrictions on constexpr in the
+  // codegen we are receiving it as `void` during constant evaluation.
+  explicit constexpr MapFieldBaseForParse(const void* prototype_as_void)
+      : prototype_as_void_(prototype_as_void) {}
+
+  enum class TaggedPtr : uintptr_t {};
+  explicit MapFieldBaseForParse(const Message* prototype, TaggedPtr ptr)
+      : payload_(ptr), prototype_as_void_(prototype) {
+    // We should not have a payload on construction.
+    ABSL_DCHECK(!IsPayload(ptr));
+  }
+
   ~MapFieldBaseForParse() = default;
 
-  const VTable* vtable_;
+  static constexpr uintptr_t kHasPayloadBit = 1;
+
+  static bool IsPayload(TaggedPtr p) {
+    return static_cast<uintptr_t>(p) & kHasPayloadBit;
+  }
+
+  mutable std::atomic<TaggedPtr> payload_{};
+  const void* prototype_as_void_;
+  PROTOBUF_TSAN_DECLARE_MEMBER;
 };
 
 // The value might be of different signedness, so use memcpy to extract it.

src/google/protobuf/map_field.cc

@@ -134,14 +134,6 @@ bool MapFieldBase::LookupMapValueNoSync(const MapKey& map_key,
   });
 }
 
-const UntypedMapBase& MapFieldBase::GetMapImpl(const MapFieldBaseForParse& map,
-                                               bool is_mutable) {
-  const auto& self = static_cast<const MapFieldBase&>(map);
-  self.SyncMapWithRepeatedField();
-  if (is_mutable) const_cast<MapFieldBase&>(self).SetMapDirty();
-  return self.GetMapRaw();
-}
-
 void MapFieldBase::MapBegin(MapIterator* map_iter) const {
   map_iter->iter_ = GetMap().begin();
   SetMapIteratorValue(map_iter);
@@ -195,8 +187,18 @@ static void SwapRelaxed(std::atomic<T>& a, std::atomic<T>& b) {
 MapFieldBase::ReflectionPayload& MapFieldBase::PayloadSlow() const {
   auto p = payload_.load(std::memory_order_acquire);
   if (!IsPayload(p)) {
+    // Inject the sync callback.
+    sync_map_with_repeated.store(
+        [](auto& map, bool is_mutable) {
+          const auto& self = static_cast<const MapFieldBase&>(map);
+          self.SyncMapWithRepeatedField();
+          if (is_mutable) const_cast<MapFieldBase&>(self).SetMapDirty();
+        },
+        std::memory_order_relaxed);
+
     auto* arena = ToArena(p);
     auto* payload = Arena::Create<ReflectionPayload>(arena, arena);
+
     auto new_p = ToTaggedPtr(payload);
     if (payload_.compare_exchange_strong(p, new_p, std::memory_order_acq_rel)) {
       // We were able to store it.

src/google/protobuf/map_field.h

@@ -297,10 +297,10 @@ inline const Message& GetMapEntryValuePrototype(const Message& default_entry) {
 // reflection implementation only. Users should never use this directly.
 class PROTOBUF_EXPORT MapFieldBase : public MapFieldBaseForParse {
  public:
-  explicit constexpr MapFieldBase(const VTable* vtable)
-      : MapFieldBaseForParse(vtable) {}
-  explicit MapFieldBase(const VTable* vtable, Arena* arena)
-      : MapFieldBaseForParse(vtable), payload_{ToTaggedPtr(arena)} {}
+  explicit constexpr MapFieldBase(const void* prototype_as_void)
+      : MapFieldBaseForParse(prototype_as_void) {}
+  explicit MapFieldBase(const Message* prototype, Arena* arena)
+      : MapFieldBaseForParse(prototype, ToTaggedPtr(arena)) {}
   MapFieldBase(const MapFieldBase&) = delete;
   MapFieldBase& operator=(const MapFieldBase&) = delete;
 
@@ -308,17 +308,6 @@ class PROTOBUF_EXPORT MapFieldBase : public MapFieldBaseForParse {
   // "protected" stops users from deleting a `MapFieldBase *`
   ~MapFieldBase();
 
-  struct VTable : MapFieldBaseForParse::VTable {
-    const Message* (*get_prototype)(const MapFieldBase& map);
-  };
-  template <typename T>
-  static constexpr VTable MakeVTable() {
-    VTable out{};
-    out.get_map = &T::GetMapImpl;
-    out.get_prototype = &T::GetPrototypeImpl;
-    return out;
-  }
-
  public:
   // Returns reference to internal repeated field. Data written using
   // Map's api prior to calling this function is guarantted to be
@@ -328,8 +317,6 @@ class PROTOBUF_EXPORT MapFieldBase : public MapFieldBaseForParse {
   // Like above. Returns mutable pointer to the internal repeated field.
   RepeatedPtrFieldBase* MutableRepeatedField();
 
-  const VTable* vtable() const { return static_cast<const VTable*>(vtable_); }
-
   bool ContainsMapKey(const MapKey& map_key) const {
     return LookupMapValue(map_key, static_cast<MapValueConstRef*>(nullptr));
   }
@@ -371,7 +358,9 @@ class PROTOBUF_EXPORT MapFieldBase : public MapFieldBaseForParse {
   }
 
  protected:
-  const Message* GetPrototype() const { return vtable()->get_prototype(*this); }
+  const Message* GetPrototype() const {
+    return reinterpret_cast<const Message*>(prototype_as_void_);
+  }
   void ClearMapNoSync();
 
   // Synchronizes the content in Map to RepeatedPtrField if there is any change
@@ -411,15 +400,16 @@ class PROTOBUF_EXPORT MapFieldBase : public MapFieldBaseForParse {
   // MapFieldBase-derived object, and there is no synchronization going
   // on between them, tsan will alert.
 #if defined(ABSL_HAVE_THREAD_SANITIZER)
-  void ConstAccess() const { ABSL_CHECK_EQ(seq1_, seq2_); }
+  // Using prototype_as_void_ as an arbitrary member that we can read/write.
+  void ConstAccess() const {
+    auto* p = prototype_as_void_;
+    asm volatile("" : "+r"(p));
+  }
   void MutableAccess() {
-    if (seq1_ & 1) {
-      seq2_ = ++seq1_;
-    } else {
-      seq1_ = ++seq2_;
-    }
+    auto* p = prototype_as_void_;
+    asm volatile("" : "+r"(p));
+    prototype_as_void_ = p;
   }
-  unsigned int seq1_ = 0, seq2_ = 0;
 #else
   void ConstAccess() const {}
   void MutableAccess() {}
@@ -466,9 +456,6 @@ class PROTOBUF_EXPORT MapFieldBase : public MapFieldBaseForParse {
                         : STATE_MODIFIED_MAP;
   }
 
-  static const UntypedMapBase& GetMapImpl(const MapFieldBaseForParse& map,
-                                          bool is_mutable);
-
  private:
   friend class ContendedMapCleanTest;
   friend class GeneratedMessageReflection;
@@ -491,14 +478,6 @@ class PROTOBUF_EXPORT MapFieldBase : public MapFieldBaseForParse {
         ->PlacementNew(msg, arena());
   }
 
-  // See assertion in TypeDefinedMapFieldBase::TypeDefinedMapFieldBase()
-  const UntypedMapBase& GetMapRaw() const {
-    return *reinterpret_cast<const UntypedMapBase*>(this + 1);
-  }
-  UntypedMapBase& GetMapRaw() {
-    return *reinterpret_cast<UntypedMapBase*>(this + 1);
-  }
-
   // Virtual helper methods for MapIterator. MapIterator doesn't have the
   // type helper for key and value. Call these help methods to deal with
   // different types. Real helper methods are implemented in
@@ -514,10 +493,6 @@ class PROTOBUF_EXPORT MapFieldBase : public MapFieldBaseForParse {
   void IncreaseIterator(MapIterator* map_iter) const;
 
   bool LookupMapValueNoSync(const MapKey& map_key, MapValueConstRef* val) const;
-
-  enum class TaggedPtr : uintptr_t {};
-  static constexpr uintptr_t kHasPayloadBit = 1;
-
   static ReflectionPayload* ToPayload(TaggedPtr p) {
     ABSL_DCHECK(IsPayload(p));
     auto* res = reinterpret_cast<ReflectionPayload*>(static_cast<uintptr_t>(p) -
@@ -536,11 +511,6 @@ class PROTOBUF_EXPORT MapFieldBase : public MapFieldBaseForParse {
   static TaggedPtr ToTaggedPtr(Arena* p) {
     return static_cast<TaggedPtr>(reinterpret_cast<uintptr_t>(p));
   }
-  static bool IsPayload(TaggedPtr p) {
-    return static_cast<uintptr_t>(p) & kHasPayloadBit;
-  }
-
-  mutable std::atomic<TaggedPtr> payload_{};
 };
 
 // This class provides common Map Reflection implementations for generated
@@ -548,23 +518,22 @@ class PROTOBUF_EXPORT MapFieldBase : public MapFieldBaseForParse {
 template <typename Key, typename T>
 class TypeDefinedMapFieldBase : public MapFieldBase {
  public:
-  explicit constexpr TypeDefinedMapFieldBase(const VTable* vtable)
-      : MapFieldBase(vtable), map_() {
-    // This invariant is required by MapFieldBase to easily access the map
-    // member without paying for dynamic dispatch. It reduces code size.
-    static_assert(PROTOBUF_FIELD_OFFSET(TypeDefinedMapFieldBase, map_) ==
-                      sizeof(MapFieldBase),
-                  "");
+  explicit constexpr TypeDefinedMapFieldBase(const void* prototype_as_void)
+      : MapFieldBase(prototype_as_void), map_() {
+    // This invariant is required by `GetMapRaw` to easily access the map
+    // member without paying for dynamic dispatch.
+    static_assert(MapFieldBaseForParse::MapOffset() ==
+                  PROTOBUF_FIELD_OFFSET(TypeDefinedMapFieldBase, map_));
   }
   TypeDefinedMapFieldBase(const TypeDefinedMapFieldBase&) = delete;
   TypeDefinedMapFieldBase& operator=(const TypeDefinedMapFieldBase&) = delete;
 
-  TypeDefinedMapFieldBase(const VTable* vtable, Arena* arena)
-      : MapFieldBase(vtable, arena), map_(arena) {}
+  TypeDefinedMapFieldBase(const Message* prototype, Arena* arena)
+      : MapFieldBase(prototype, arena), map_(arena) {}
 
-  TypeDefinedMapFieldBase(const VTable* vtable, Arena* arena,
+  TypeDefinedMapFieldBase(const Message* prototype, Arena* arena,
                           const TypeDefinedMapFieldBase& from)
-      : MapFieldBase(vtable, arena), map_(arena, from.GetMap()) {}
+      : MapFieldBase(prototype, arena), map_(arena, from.GetMap()) {}
 
  protected:
   ~TypeDefinedMapFieldBase() { map_.~Map(); }
@@ -623,38 +592,33 @@ class MapField final : public TypeDefinedMapFieldBase<Key, T> {
   static constexpr WireFormatLite::FieldType kKeyFieldType = kKeyFieldType_;
   static constexpr WireFormatLite::FieldType kValueFieldType = kValueFieldType_;
 
-  constexpr MapField() : MapField::TypeDefinedMapFieldBase(&kVTable) {}
+  constexpr MapField()
+      : MapField::TypeDefinedMapFieldBase(
+            Derived::internal_default_instance()) {}
   MapField(const MapField&) = delete;
   MapField& operator=(const MapField&) = delete;
   ~MapField() = default;
 
   explicit MapField(Arena* arena)
-      : TypeDefinedMapFieldBase<Key, T>(&kVTable, arena) {}
+      : TypeDefinedMapFieldBase<Key, T>(
+            static_cast<const Message*>(Derived::internal_default_instance()),
+            arena) {}
   MapField(ArenaInitialized, Arena* arena) : MapField(arena) {}
   MapField(InternalVisibility, Arena* arena) : MapField(arena) {}
   MapField(InternalVisibility, Arena* arena, const MapField& from)
-      : TypeDefinedMapFieldBase<Key, T>(&kVTable, arena, from) {}
+      : TypeDefinedMapFieldBase<Key, T>(
+            static_cast<const Message*>(Derived::internal_default_instance()),
+            arena, from) {}
 
  private:
   typedef void InternalArenaConstructable_;
   typedef void DestructorSkippable_;
 
-  static const Message* GetPrototypeImpl(const MapFieldBase& map);
-
-  static const MapFieldBase::VTable kVTable;
-
   friend class google::protobuf::Arena;
   friend class MapFieldBase;
   friend class MapFieldStateTest;  // For testing, it needs raw access to impl_
 };
 
-template <typename Derived, typename Key, typename T,
-          WireFormatLite::FieldType kKeyFieldType_,
-          WireFormatLite::FieldType kValueFieldType_>
-PROTOBUF_CONSTINIT const MapFieldBase::VTable
-    MapField<Derived, Key, T, kKeyFieldType_, kValueFieldType_>::kVTable =
-        MapField::template MakeVTable<MapField>();
-
 template <typename Key, typename T>
 bool AllAreInitialized(const TypeDefinedMapFieldBase<Key, T>& field) {
   for (const auto& p : field.GetMap()) {

src/google/protobuf/map_field_inl.h

@@ -23,28 +23,6 @@
 #include "google/protobuf/message.h"
 #include "google/protobuf/port.h"
 
-// must be last
-#include "google/protobuf/port_def.inc"
-
-#ifdef SWIG
-#error "You cannot SWIG proto headers"
-#endif
-
-namespace google {
-namespace protobuf {
-namespace internal {
-template <typename Derived, typename Key, typename T,
-          WireFormatLite::FieldType kKeyFieldType,
-          WireFormatLite::FieldType kValueFieldType>
-const Message*
-MapField<Derived, Key, T, kKeyFieldType, kValueFieldType>::GetPrototypeImpl(
-    const MapFieldBase&) {
-  return Derived::internal_default_instance();
-}
-}  // namespace internal
-}  // namespace protobuf
-}  // namespace google
-
-#include "google/protobuf/port_undef.inc"
+// TODO: Remove this file now that it is empty
 
 #endif  // GOOGLE_PROTOBUF_MAP_FIELD_INL_H__