Make WireFormat::ByteSize aware of LazyField. (roll forward after TGP).

PiperOrigin-RevId: 672719183
protocolbuffers · Sep 9, 2024 · 9c835af · 9c835af
1 parent 960b5ba
commit 9c835af
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 1 deletion.
diff --git a/src/google/protobuf/extension_set.cc b/src/google/protobuf/extension_set.cc
@@ -513,6 +513,14 @@ 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();
+}
+
 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
@@ -576,6 +576,8 @@ 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;
+
   // 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
@@ -1705,7 +1705,26 @@ size_t WireFormat::FieldDataOnlyByteSize(const FieldDescriptor* field,
     HANDLE_FIXED_TYPE(BOOL, Bool)
 
     HANDLE_TYPE(GROUP, Group, Message)
-    HANDLE_TYPE(MESSAGE, Message, Message)
+
+    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;
+    }
+
 #undef HANDLE_TYPE
 #undef HANDLE_FIXED_TYPE