From c5c509001dd5675a8d5d8ff6fb739233af474649 Mon Sep 17 00:00:00 2001
From: Protobuf Team Bot <protobuf-github-bot@google.com>
Date: Mon, 16 Sep 2024 13:19:22 -0700
Subject: [PATCH] Make WireFormat::(ByteSize|_InternalSerialize) LazyField
 compatible (take #3 but disabled)

PiperOrigin-RevId: 675266565
---
 src/google/protobuf/extension_set.cc | 25 +++++++++++++++++++++++++
 src/google/protobuf/extension_set.h  |  5 +++++
 src/google/protobuf/wire_format.cc   | 23 +++++++++++++++++++++++
 3 files changed, 53 insertions(+)

diff --git a/src/google/protobuf/extension_set.cc b/src/google/protobuf/extension_set.cc
index 12a730ce1ca3..44f4b1ec2925 100644
--- a/src/google/protobuf/extension_set.cc
+++ b/src/google/protobuf/extension_set.cc
@@ -513,6 +513,31 @@ const int& ExtensionSet::GetRefRepeatedEnum(int number, int index) const {
   return extension->ptr.repeated_enum_value->Get(index);
 }
 
+size_t ExtensionSet::GetMessageByteSizeLong(int number) const {
+  const Extension* extension = FindOrNull(number);
+  ABSL_CHECK(extension != nullptr) << "not present";
+  ABSL_DCHECK_TYPE(*extension, OPTIONAL_FIELD, MESSAGE);
+  return extension->is_lazy ? extension->ptr.lazymessage_value->ByteSizeLong()
+                            : extension->ptr.message_value->ByteSizeLong();
+}
+
+uint8_t* ExtensionSet::InternalSerializeMessage(
+    int number, const MessageLite* prototype, uint8_t* target,
+    io::EpsCopyOutputStream* stream) const {
+  const Extension* extension = FindOrNull(number);
+  ABSL_CHECK(extension != nullptr) << "not present";
+  ABSL_DCHECK_TYPE(*extension, OPTIONAL_FIELD, MESSAGE);
+
+  if (extension->is_lazy) {
+    return extension->ptr.lazymessage_value->WriteMessageToArray(
+        prototype, number, target, stream);
+  }
+
+  const auto* msg = extension->ptr.message_value;
+  return WireFormatLite::InternalWriteMessage(
+      number, *msg, msg->GetCachedSize(), target, stream);
+}
+
 void ExtensionSet::SetRepeatedEnum(int number, int index, int value) {
   Extension* extension = FindOrNull(number);
   ABSL_CHECK(extension != nullptr) << "Index out-of-bounds (field is empty).";
diff --git a/src/google/protobuf/extension_set.h b/src/google/protobuf/extension_set.h
index 5a8d9abff302..0daa4bbfc1b9 100644
--- a/src/google/protobuf/extension_set.h
+++ b/src/google/protobuf/extension_set.h
@@ -576,6 +576,11 @@ class PROTOBUF_EXPORT ExtensionSet {
   const bool& GetRefRepeatedBool(int number, int index) const;
   const int& GetRefRepeatedEnum(int number, int index) const;
 
+  size_t GetMessageByteSizeLong(int number) const;
+  uint8_t* InternalSerializeMessage(int number, const MessageLite* prototype,
+                                    uint8_t* target,
+                                    io::EpsCopyOutputStream* stream) const;
+
   // Implementation of _InternalSerialize for non-empty map_.
   uint8_t* _InternalSerializeImpl(const MessageLite* extendee,
                                   int start_field_number, int end_field_number,
diff --git a/src/google/protobuf/wire_format.cc b/src/google/protobuf/wire_format.cc
index 302155ce285d..4bdc949c8224 100644
--- a/src/google/protobuf/wire_format.cc
+++ b/src/google/protobuf/wire_format.cc
@@ -1232,6 +1232,7 @@ uint8_t* WireFormat::InternalSerializeField(const FieldDescriptor* field,
     return InternalSerializeMessageSetItem(field, message, target, stream);
   }
 
+
   // For map fields, we can use either repeated field reflection or map
   // reflection.  Our choice has some subtle effects.  If we use repeated field
   // reflection here, then the repeated field representation becomes
@@ -1707,7 +1708,29 @@ size_t WireFormat::FieldDataOnlyByteSize(const FieldDescriptor* field,
     HANDLE_FIXED_TYPE(BOOL, Bool)
 
     HANDLE_TYPE(GROUP, Group, Message)
+#ifndef PROTOBUF_ENABLE_FIX_CODE_SIZE_LAZY
     HANDLE_TYPE(MESSAGE, Message, Message)
+#else  // !PROTOBUF_ENABLE_FIX_CODE_SIZE_LAZY
+    case FieldDescriptor::TYPE_MESSAGE: {
+      if (field->is_repeated()) {
+        for (size_t j = 0; j < count; ++j) {
+          data_size += WireFormatLite::MessageSize(
+              message_reflection->GetRepeatedMessage(message, field, j));
+        }
+        break;
+      }
+      if (field->is_extension()) {
+        data_size += WireFormatLite::LengthDelimitedSize(
+            message_reflection->GetExtensionSet(message).GetMessageByteSizeLong(
+                field->number()));
+        break;
+      }
+      data_size += WireFormatLite::MessageSize(
+          message_reflection->GetMessage(message, field));
+      break;
+    }
+#endif  // PROTOBUF_ENABLE_FIX_CODE_SIZE_LAZY
+
 #undef HANDLE_TYPE
 #undef HANDLE_FIXED_TYPE