From 224573d66a0cc958c76cb43d8b2eb3aa7cdb89f2 Mon Sep 17 00:00:00 2001
From: Thomas Van Lenten <thomasvl@google.com>
Date: Tue, 23 Jul 2024 09:55:21 -0700
Subject: [PATCH] [ObjC] Internal helper for getting the unknown field data
 from a message

PiperOrigin-RevId: 655202619
---
 objectivec/GPBMessage.m                | 19 +++++++++++++++++++
 objectivec/GPBMessage_PackagePrivate.h |  3 +++
 objectivec/GPBUnknownFields.m          | 19 +++----------------
 3 files changed, 25 insertions(+), 16 deletions(-)

diff --git a/objectivec/GPBMessage.m b/objectivec/GPBMessage.m
index 1761eab485..b0367699dc 100644
--- a/objectivec/GPBMessage.m
+++ b/objectivec/GPBMessage.m
@@ -3684,4 +3684,23 @@ id GPBGetObjectIvarWithField(GPBMessage *self, GPBFieldDescriptor *field) {
   return expected;
 }
 
+NSData *GPBMessageUnknownFieldsData(GPBMessage *self) {
+  NSData *result = nil;
+  GPBUnknownFieldSet *unknownFields = self->unknownFields_;
+  if (unknownFields) {
+    if (self.descriptor.isWireFormat) {
+      NSMutableData *mutableData =
+          [NSMutableData dataWithLength:unknownFields.serializedSizeAsMessageSet];
+      GPBCodedOutputStream *output = [[GPBCodedOutputStream alloc] initWithData:mutableData];
+      [unknownFields writeAsMessageSetTo:output];
+      [output flush];
+      [output release];
+      result = mutableData;
+    } else {
+      result = [unknownFields data];
+    }
+  }
+  return result;
+}
+
 #pragma clang diagnostic pop
diff --git a/objectivec/GPBMessage_PackagePrivate.h b/objectivec/GPBMessage_PackagePrivate.h
index 65a1a7c5ba..58649c44f7 100644
--- a/objectivec/GPBMessage_PackagePrivate.h
+++ b/objectivec/GPBMessage_PackagePrivate.h
@@ -81,4 +81,7 @@ void GPBAutocreatedDictionaryModified(GPBMessage *self, id dictionary);
 // autocreated reference to this message.
 void GPBClearMessageAutocreator(GPBMessage *self);
 
+// The data (or null) from the unknown fields of a message;
+NSData *GPBMessageUnknownFieldsData(GPBMessage *self);
+
 CF_EXTERN_C_END
diff --git a/objectivec/GPBUnknownFields.m b/objectivec/GPBUnknownFields.m
index f3073ad16d..836e249f67 100644
--- a/objectivec/GPBUnknownFields.m
+++ b/objectivec/GPBUnknownFields.m
@@ -14,6 +14,7 @@
 #import "GPBCodedOutputStream_PackagePrivate.h"
 #import "GPBDescriptor.h"
 #import "GPBMessage.h"
+#import "GPBMessage_PackagePrivate.h"
 #import "GPBUnknownField.h"
 #import "GPBUnknownFieldSet_PackagePrivate.h"
 #import "GPBUnknownField_PackagePrivate.h"
@@ -196,22 +197,8 @@ static BOOL MergeFromInputStream(GPBUnknownFields *self, GPBCodedInputStream *in
   self = [super init];
   if (self) {
     fields_ = [[NSMutableArray alloc] init];
-    // TODO: b/349146447 - Move off the legacy class and directly to the data once Message is
-    // updated.
-    GPBUnknownFieldSet *legacyUnknownFields = [message unknownFields];
-    if (legacyUnknownFields) {
-      NSData *data;
-      if (message.descriptor.isWireFormat) {
-        NSMutableData *mutableData =
-            [NSMutableData dataWithLength:legacyUnknownFields.serializedSizeAsMessageSet];
-        GPBCodedOutputStream *output = [[GPBCodedOutputStream alloc] initWithData:mutableData];
-        [legacyUnknownFields writeAsMessageSetTo:output];
-        [output flush];
-        [output release];
-        data = mutableData;
-      } else {
-        data = [legacyUnknownFields data];
-      }
+    NSData *data = GPBMessageUnknownFieldsData(message);
+    if (data) {
       GPBCodedInputStream *input = [[GPBCodedInputStream alloc] initWithData:data];
       // Parse until the end of the data (tag will be zero).
       if (!MergeFromInputStream(self, input, 0)) {