Fix a bug in InlinedStringField where donate steal is not enabled.

PiperOrigin-RevId: 625715238
protocolbuffers · Apr 17, 2024 · bb3ff34 · bb3ff34
1 parent 58fb524
commit bb3ff34
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 13 deletions.
diff --git a/src/google/protobuf/compiler/cpp/message.cc b/src/google/protobuf/compiler/cpp/message.cc
@@ -2477,7 +2477,12 @@ void MessageGenerator::GenerateImplMemberInit(io::Printer* p,
       bool dtor_on_demand = NeedsArenaDestructor() == ArenaDtorNeeds::kOnDemand;
       auto values = [&] {
         for (size_t i = 0; i < InlinedStringDonatedSize(); ++i) {
-          p->Emit(i ? ", ~0u" : dtor_on_demand ? "~0u" : "0xFFFFFFFEu");
+          if (i > 0) {
+            p->Emit(", ");
+          }
+          p->Emit(dtor_on_demand
+                      ? "::_pbi::InitDonatingStates()"
+                      : "::_pbi::InitDonatingStates() & 0xFFFFFFFEu");
         }
       };
       separator();
@@ -3029,12 +3034,24 @@ void MessageGenerator::GenerateArenaEnabledCopyConstructor(io::Printer* p) {
   };
 
   auto maybe_register_arena_dtor = [&] {
-    if (NeedsArenaDestructor() == ArenaDtorNeeds::kRequired) {
-      p->Emit(R"cc(
-        if (arena != nullptr) {
-          arena->OwnCustomDestructor(this, &$classname$::ArenaDtor);
-        }
-      )cc");
+    switch (NeedsArenaDestructor()) {
+      case ArenaDtorNeeds::kRequired: {
+        p->Emit(R"cc(
+          if (arena != nullptr) {
+            arena->OwnCustomDestructor(this, &$classname$::ArenaDtor);
+          }
+        )cc");
+        break;
+      }
+      case ArenaDtorNeeds::kOnDemand: {
+        p->Emit(R"cc(
+          ::_pbi::InternalRegisterArenaDtor(arena, this,
+                                            &$classname$::ArenaDtor);
+        )cc");
+        break;
+      }
+      case ArenaDtorNeeds::kNone:
+        break;
     }
   };
 
@@ -3071,12 +3088,24 @@ void MessageGenerator::GenerateStructors(io::Printer* p) {
            [&] {
              if (HasSimpleBaseClass(descriptor_, options_)) return;
              p->Emit(R"cc(SharedCtor(arena);)cc");
-             if (NeedsArenaDestructor() == ArenaDtorNeeds::kRequired) {
-               p->Emit(R"cc(
-                 if (arena != nullptr) {
-                   arena->OwnCustomDestructor(this, &$classname$::ArenaDtor);
-                 }
-               )cc");
+             switch (NeedsArenaDestructor()) {
+               case ArenaDtorNeeds::kRequired: {
+                 p->Emit(R"cc(
+                   if (arena != nullptr) {
+                     arena->OwnCustomDestructor(this, &$classname$::ArenaDtor);
+                   }
+                 )cc");
+                 break;
+               }
+               case ArenaDtorNeeds::kOnDemand: {
+                 p->Emit(R"cc(
+                   ::_pbi::InternalRegisterArenaDtor(arena, this,
+                                                     &$classname$::ArenaDtor);
+                 )cc");
+                 break;
+               }
+               case ArenaDtorNeeds::kNone:
+                 break;
              }
            }},
       },

diff --git a/src/google/protobuf/inlined_string_field.h b/src/google/protobuf/inlined_string_field.h
@@ -373,6 +373,19 @@ inline InlinedStringField::InlinedStringField(
 }
 
 
+#ifdef GOOGLE_PROTOBUF_INTERNAL_DONATE_STEAL_INLINE
+constexpr uint32_t InitDonatingStates() { return ~0u; }
+inline void InternalRegisterArenaDtor(Arena*, void*, void (*)(void*)) {}
+#else   // !GOOGLE_PROTOBUF_INTERNAL_DONATE_STEAL_INLINE
+constexpr uint32_t InitDonatingStates() { return 0u; }
+inline void InternalRegisterArenaDtor(Arena* arena, void* object,
+                                      void (*destruct)(void*)) {
+  if (arena != nullptr) {
+    arena->OwnCustomDestructor(object, destruct);
+  }
+}
+#endif  // GOOGLE_PROTOBUF_INTERNAL_DONATE_STEAL_INLINE
+
 inline InlinedStringField::InlinedStringField(Arena* /*arena*/) { Init(); }
 
 inline InlinedStringField::InlinedStringField(Arena* arena,