Use Arena enabled copy constructor when merging child messages

PiperOrigin-RevId: 575841937

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

@@ -263,6 +263,7 @@ void HasBitVars(const FieldDescriptor* field, const Options& opts,
                 absl::optional<uint32_t> idx, std::vector<Sub>& vars) {
   if (!idx.has_value()) {
     vars.emplace_back("set_hasbit", "");
+    vars.emplace_back("this_set_hasbit", "");
     vars.emplace_back("clear_hasbit", "");
     return;
   }
@@ -283,6 +284,9 @@ void HasBitVars(const FieldDescriptor* field, const Options& opts,
   vars.emplace_back("has_hasbit", has);
   vars.emplace_back(Sub("set_hasbit", set).WithSuffix(";"));
   vars.emplace_back(Sub("clear_hasbit", clr).WithSuffix(";"));
+
+  set = absl::StrFormat("_this->%s[%d] |= %s;", has_bits, index, mask);
+  vars.emplace_back(Sub("this_set_hasbit", set).WithSuffix(";"));
 }
 
 void InlinedStringVars(const FieldDescriptor* field, const Options& opts,

src/google/protobuf/compiler/cpp/field.h

@@ -39,6 +39,11 @@ namespace cpp {
 // matter of clean composability.
 class FieldGeneratorBase {
  public:
+  // `GeneratorFunction` defines a subset of generator functions that may have
+  // additional optimizations or requirements such as 'uses a local `arena`
+  // variable instead of calling GetArena()'
+  enum class GeneratorFunction { kMergeFrom };
+
   FieldGeneratorBase(const FieldDescriptor* descriptor, const Options& options,
                      MessageSCCAnalyzer* scc_analyzer);
 
@@ -100,6 +105,10 @@ class FieldGeneratorBase {
     return has_default_constexpr_constructor_;
   }
 
+  // Returns true if this generator requires an 'arena' parameter on the
+  // given generator function.
+  virtual bool RequiresArena(GeneratorFunction) const { return false; }
+
   virtual std::vector<io::Printer::Sub> MakeVars() const { return {}; }
 
   virtual void GeneratePrivateMembers(io::Printer* p) const = 0;
@@ -230,6 +239,8 @@ class FieldGenerator {
   }
 
  public:
+  using GeneratorFunction = FieldGeneratorBase::GeneratorFunction;
+
   FieldGenerator(const FieldGenerator&) = delete;
   FieldGenerator& operator=(const FieldGenerator&) = delete;
   FieldGenerator(FieldGenerator&&) = default;
@@ -256,6 +267,11 @@ class FieldGenerator {
     return impl_->has_default_constexpr_constructor();
   }
 
+  // Requirements: see FieldGeneratorBase for documentation
+  bool RequiresArena(GeneratorFunction function) const {
+    return impl_->RequiresArena(function);
+  }
+
   // Prints private members needed to represent this field.
   //
   // These are placed inside the class definition.

src/google/protobuf/compiler/cpp/field_generators/message_field.cc

@@ -102,6 +102,8 @@ class SingularMessage : public FieldGeneratorBase {
     )cc");
   }
 
+  bool RequiresArena(GeneratorFunction function) const override;
+
   void GenerateNonInlineAccessorDefinitions(io::Printer* p) const override {}
 
   void GenerateAccessorDeclarations(io::Printer* p) const override;
@@ -415,15 +417,32 @@ void SingularMessage::GenerateMessageClearingCode(io::Printer* p) const {
   }
 }
 
+bool SingularMessage::RequiresArena(GeneratorFunction function) const {
+  switch (function) {
+    case GeneratorFunction::kMergeFrom:
+      return !(is_weak() || is_oneof() || should_split());
+  }
+  return false;
+}
+
 void SingularMessage::GenerateMergingCode(io::Printer* p) const {
   if (is_weak()) {
     p->Emit(
         "_Internal::mutable_$name$(_this)->CheckTypeAndMergeFrom(\n"
         "    _Internal::$name$(&from));\n");
-  } else {
+  } else if (is_oneof() || should_split()) {
     p->Emit(
         "_this->_internal_mutable_$name$()->$Submsg$::MergeFrom(\n"
         "    from._internal_$name$());\n");
+  } else {
+    p->Emit(R"cc(
+      $this_set_hasbit$;
+      if (_this->$field_$ == nullptr) {
+        _this->$field_$ = CreateMaybeMessage<$Submsg$>(arena, *from.$field_$);
+      } else {
+        _this->$field_$->MergeFrom(*from.$field_$);
+      }
+    )cc");
   }
 }
 

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

@@ -3496,6 +3496,15 @@ void MessageGenerator::GenerateMergeFrom(io::Printer* p) {
   }
 }
 
+bool MessageGenerator::RequiresArena(GeneratorFunction function) const {
+  for (const FieldDescriptor* field : FieldRange(descriptor_)) {
+    if (field_generators_.get(field).RequiresArena(function)) {
+      return true;
+    }
+  }
+  return false;
+}
+
 void MessageGenerator::GenerateClassSpecificMergeImpl(io::Printer* p) {
   if (HasSimpleBaseClass(descriptor_, options_)) return;
   // Generate the class-specific MergeFrom, which avoids the ABSL_CHECK and
@@ -3515,6 +3524,11 @@ void MessageGenerator::GenerateClassSpecificMergeImpl(io::Printer* p) {
         "  auto& from = static_cast<const $classname$&>(from_msg);\n");
   }
   format.Indent();
+  if (RequiresArena(GeneratorFunction::kMergeFrom)) {
+    p->Emit(R"cc(
+      ::$proto_ns$::Arena* arena = _this->GetArena();
+    )cc");
+  }
   format(
       "$annotate_mergefrom$"
       "// @@protoc_insertion_point(class_specific_merge_from_start:"

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

@@ -80,6 +80,7 @@ class MessageGenerator {
   const Descriptor* descriptor() const { return descriptor_; }
 
  private:
+  using GeneratorFunction = FieldGeneratorBase::GeneratorFunction;
   enum class InitType { kConstexpr, kArena, kArenaCopy };
 
   // Generate declarations and definitions of accessors for fields.
@@ -148,6 +149,10 @@ class MessageGenerator {
   void GenerateFieldClear(const FieldDescriptor* field, bool is_inline,
                           io::Printer* p);
 
+  // Returns true if any of the fields needs an `arena` variable containing
+  // the current message's arena, reducing `GetArena()` call churn.
+  bool RequiresArena(GeneratorFunction function) const;
+
   // Returns whether impl_ has a copy ctor.
   bool ImplHasCopyCtor() const;