Avoid field id comparisons in ExtensionSet serialization in case of s…

…ingle range.

PiperOrigin-RevId: 685733801

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

@@ -4511,12 +4511,23 @@ void MessageGenerator::GenerateSerializeOneField(io::Printer* p,
 void MessageGenerator::GenerateSerializeOneExtensionRange(io::Printer* p,
                                                           int start, int end) {
   auto v = p->WithVars(variables_);
-  p->Emit({{"start", start}, {"end", end}},
-          R"cc(
-            // Extension range [$start$, $end$)
-            target = this_.$extensions$._InternalSerialize(
-                internal_default_instance(), $start$, $end$, target, stream);
-          )cc");
+  p->Emit(  //
+      {{"start", start}, {"end", end}},
+      R"cc(
+        // Extension range [$start$, $end$)
+        target = this_.$extensions$._InternalSerialize(
+            internal_default_instance(), $start$, $end$, target, stream);
+      )cc");
+}
+
+void MessageGenerator::GenerateSerializeAllExtensions(io::Printer* p) {
+  auto v = p->WithVars(variables_);
+  p->Emit(
+      R"cc(
+        // All extensions.
+        target = this_.$extensions$._InternalSerializeAll(
+            internal_default_instance(), target, stream);
+      )cc");
 }
 
 void MessageGenerator::GenerateSerializeWithCachedSizesToArray(io::Printer* p) {
@@ -4682,16 +4693,23 @@ void MessageGenerator::GenerateSerializeWithCachedSizesBody(io::Printer* p) {
       }
     }
 
-    void Flush() {
-      if (has_current_range_) {
-        mg_->GenerateSerializeOneExtensionRange(p_, min_start_, max_end_);
+    void Flush(bool is_last_range) {
+      if (!has_current_range_) {
+        return;
       }
       has_current_range_ = false;
+      ++range_count_;
+      if (is_last_range && range_count_ == 1) {
+        mg_->GenerateSerializeAllExtensions(p_);
+      } else {
+        mg_->GenerateSerializeOneExtensionRange(p_, min_start_, max_end_);
+      }
     }
 
    private:
     MessageGenerator* mg_;
     io::Printer* p_;
+    int range_count_ = 0;
     bool has_current_range_ = false;
     int min_start_ = 0;
     int max_end_ = 0;
@@ -4749,12 +4767,13 @@ void MessageGenerator::GenerateSerializeWithCachedSizesBody(io::Printer* p) {
              size_t i, j;
              for (i = 0, j = 0;
                   i < ordered_fields.size() || j < sorted_extensions.size();) {
-               if ((j == sorted_extensions.size()) ||
+               bool no_more_extensions = j == sorted_extensions.size();
+               if (no_more_extensions ||
                    (i < static_cast<size_t>(descriptor_->field_count()) &&
                     ordered_fields[i]->number() <
                         sorted_extensions[j]->start_number())) {
                  const FieldDescriptor* field = ordered_fields[i++];
-                 re.Flush();
+                 re.Flush(no_more_extensions);
                  if (field->options().weak()) {
                    largest_weak_field.ReplaceIfLarger(field);
                    PrintFieldComment(Formatter{p}, field, options_);
@@ -4768,7 +4787,7 @@ void MessageGenerator::GenerateSerializeWithCachedSizesBody(io::Printer* p) {
                  re.AddToRange(sorted_extensions[j++]);
                }
              }
-             re.Flush();
+             re.Flush(/*is_last_range=*/true);
              e.EmitIfNotNull(largest_weak_field.Release());
            }},
           {"handle_unknown_fields",

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

@@ -157,6 +157,7 @@ class MessageGenerator {
   void GenerateSerializeOneofFields(
       io::Printer* p, const std::vector<const FieldDescriptor*>& fields);
   void GenerateSerializeOneExtensionRange(io::Printer* p, int start, int end);
+  void GenerateSerializeAllExtensions(io::Printer* p);
 
   // Generates has_foo() functions and variables for singular field has-bits.
   void GenerateSingularFieldHasBits(const FieldDescriptor* field,

src/google/protobuf/extension_set.cc

@@ -13,6 +13,7 @@
 
 #include <algorithm>
 #include <atomic>
+#include <cassert>
 #include <cstddef>
 #include <cstdint>
 #include <string>
@@ -1279,13 +1280,8 @@ uint8_t* ExtensionSet::_InternalSerializeImpl(
     const MessageLite* extendee, int start_field_number, int end_field_number,
     uint8_t* target, io::EpsCopyOutputStream* stream) const {
   if (PROTOBUF_PREDICT_FALSE(is_large())) {
-    const auto& end = map_.large->end();
-    for (auto it = map_.large->lower_bound(start_field_number);
-         it != end && it->first < end_field_number; ++it) {
-      target = it->second.InternalSerializeFieldWithCachedSizesToArray(
-          extendee, this, it->first, target, stream);
-    }
-    return target;
+    return _InternalSerializeImplLarge(extendee, start_field_number,
+                                       end_field_number, target, stream);
   }
   const KeyValue* end = flat_end();
   const KeyValue* it = flat_begin();
@@ -1297,6 +1293,31 @@ uint8_t* ExtensionSet::_InternalSerializeImpl(
   return target;
 }
 
+uint8_t* ExtensionSet::_InternalSerializeAllImpl(
+    const MessageLite* extendee, uint8_t* target,
+    io::EpsCopyOutputStream* stream) const {
+  ForEach(
+      [&target, extendee, stream, this](int number, const Extension& ext) {
+        target = ext.InternalSerializeFieldWithCachedSizesToArray(
+            extendee, this, number, target, stream);
+      },
+      Prefetch{});
+  return target;
+}
+
+uint8_t* ExtensionSet::_InternalSerializeImplLarge(
+    const MessageLite* extendee, int start_field_number, int end_field_number,
+    uint8_t* target, io::EpsCopyOutputStream* stream) const {
+  assert(is_large());
+  const auto& end = map_.large->end();
+  for (auto it = map_.large->lower_bound(start_field_number);
+       it != end && it->first < end_field_number; ++it) {
+    target = it->second.InternalSerializeFieldWithCachedSizesToArray(
+        extendee, this, it->first, target, stream);
+  }
+  return target;
+}
+
 uint8_t* ExtensionSet::InternalSerializeMessageSetWithCachedSizesToArray(
     const MessageLite* extendee, uint8_t* target,
     io::EpsCopyOutputStream* stream) const {

src/google/protobuf/extension_set.h

@@ -479,7 +479,6 @@ class PROTOBUF_EXPORT ExtensionSet {
   // serialized extensions.
   //
   // Returns a pointer past the last written byte.
-
   uint8_t* _InternalSerialize(const MessageLite* extendee,
                               int start_field_number, int end_field_number,
                               uint8_t* target,
@@ -492,6 +491,16 @@ class PROTOBUF_EXPORT ExtensionSet {
                                   end_field_number, target, stream);
   }
 
+  // Same as _InternalSerialize, but do not verify the range of field numbers.
+  uint8_t* _InternalSerializeAll(const MessageLite* extendee, uint8_t* target,
+                                 io::EpsCopyOutputStream* stream) const {
+    if (flat_size_ == 0) {
+      assert(!is_large());
+      return target;
+    }
+    return _InternalSerializeAllImpl(extendee, target, stream);
+  }
+
   // Like above but serializes in MessageSet format.
   void SerializeMessageSetWithCachedSizes(const MessageLite* extendee,
                                           io::CodedOutputStream* output) const {
@@ -586,6 +595,17 @@ class PROTOBUF_EXPORT ExtensionSet {
                                   int start_field_number, int end_field_number,
                                   uint8_t* target,
                                   io::EpsCopyOutputStream* stream) const;
+  // Implementation of _InternalSerializeAll for non-empty map_.
+  uint8_t* _InternalSerializeAllImpl(const MessageLite* extendee,
+                                     uint8_t* target,
+                                     io::EpsCopyOutputStream* stream) const;
+  // Implementation of _InternalSerialize for large map_.
+  // Extracted as a separate method to avoid inlining and to reuse in
+  // _InternalSerializeAllImpl.
+  uint8_t* _InternalSerializeImplLarge(const MessageLite* extendee,
+                                       int start_field_number,
+                                       int end_field_number, uint8_t* target,
+                                       io::EpsCopyOutputStream* stream) const;
   // Interface of a lazily parsed singular message extension.
   class PROTOBUF_EXPORT LazyMessageExtension {
    public: