From c5c9c89441e6d04bc1246aa78f44a33f2a407251 Mon Sep 17 00:00:00 2001
From: Thomas Van Lenten <thomasvl@google.com>
Date: Thu, 1 Aug 2024 14:25:07 -0700
Subject: [PATCH] [ObjC] Add apis for removing things from `GPBUnknownFields`.

PiperOrigin-RevId: 658538490
---
 objectivec/GPBUnknownFields.h           | 20 +++++++
 objectivec/GPBUnknownFields.m           | 28 +++++++++
 objectivec/Tests/GPBUnknownFieldsTest.m | 76 +++++++++++++++++++++++++
 3 files changed, 124 insertions(+)

diff --git a/objectivec/GPBUnknownFields.h b/objectivec/GPBUnknownFields.h
index 1ee50576a4..4f05c2b972 100644
--- a/objectivec/GPBUnknownFields.h
+++ b/objectivec/GPBUnknownFields.h
@@ -109,6 +109,26 @@ __attribute__((objc_subclassing_restricted))
  **/
 - (nonnull GPBUnknownFields *)addGroupWithFieldNumber:(int32_t)fieldNumber;
 
+/**
+ * Removes the given field from the set.
+ *
+ * It is a programming error to attempt to remove a field that is not in this collection.
+ *
+ * Reminder: it is not save to mutate the collection while also using fast enumeration on it.
+ *
+ * @param field The field to remove.
+ **/
+- (void)removeField:(nonnull GPBUnknownField *)field;
+
+/**
+ * Removes all of the fields from the collection that have the given field number.
+ *
+ * If there are no fields with the given field number, this is a no-op.
+ *
+ * @param fieldNumber The field number to remove.
+ **/
+- (void)clearFieldNumber:(int32_t)fieldNumber;
+
 @end
 
 @interface GPBUnknownFields (AccessHelpers)
diff --git a/objectivec/GPBUnknownFields.m b/objectivec/GPBUnknownFields.m
index 836e249f67..2672d10b15 100644
--- a/objectivec/GPBUnknownFields.m
+++ b/objectivec/GPBUnknownFields.m
@@ -318,6 +318,34 @@ static BOOL MergeFromInputStream(GPBUnknownFields *self, GPBCodedInputStream *in
   return [group autorelease];
 }
 
+- (void)removeField:(nonnull GPBUnknownField *)field {
+  NSUInteger count = fields_.count;
+  [fields_ removeObjectIdenticalTo:field];
+  if (count == fields_.count) {
+    [NSException raise:NSInvalidArgumentException format:@"The field was not present."];
+  }
+}
+
+- (void)clearFieldNumber:(int32_t)fieldNumber {
+  CHECK_FIELD_NUMBER(fieldNumber);
+  NSMutableIndexSet *toRemove = nil;
+  NSUInteger idx = 0;
+  for (GPBUnknownField *field in fields_) {
+    if (field->number_ == fieldNumber) {
+      if (toRemove == nil) {
+        toRemove = [[NSMutableIndexSet alloc] initWithIndex:idx];
+      } else {
+        [toRemove addIndex:idx];
+      }
+    }
+    ++idx;
+  }
+  if (toRemove) {
+    [fields_ removeObjectsAtIndexes:toRemove];
+    [toRemove release];
+  }
+}
+
 #pragma mark - NSFastEnumeration protocol
 
 - (NSUInteger)countByEnumeratingWithState:(NSFastEnumerationState *)state
diff --git a/objectivec/Tests/GPBUnknownFieldsTest.m b/objectivec/Tests/GPBUnknownFieldsTest.m
index 8c2adcf32c..6c6481d480 100644
--- a/objectivec/Tests/GPBUnknownFieldsTest.m
+++ b/objectivec/Tests/GPBUnknownFieldsTest.m
@@ -510,6 +510,82 @@
   XCTAssertNil([ufs fields:99]);  // Not present
 }
 
+- (void)testRemoveField {
+  GPBUnknownFields* ufs = [[[GPBUnknownFields alloc] init] autorelease];
+  [ufs addFieldNumber:1 varint:1];
+  [ufs addFieldNumber:1 fixed32:1];
+  [ufs addFieldNumber:1 fixed64:1];
+  XCTAssertEqual(ufs.count, 3);
+
+  NSArray<GPBUnknownField*>* fields = [ufs fields:1];
+  XCTAssertEqual(fields.count, 3);
+  GPBUnknownField* field = fields[0];
+  XCTAssertEqual(field.number, 1);
+  XCTAssertEqual(field.type, GPBUnknownFieldTypeVarint);
+  XCTAssertEqual(field.varint, 1);
+  [ufs removeField:field];  // Remove first (varint)
+  XCTAssertEqual(ufs.count, 2);
+
+  fields = [ufs fields:1];
+  XCTAssertEqual(fields.count, 2);
+  field = fields[0];
+  XCTAssertEqual(field.number, 1);
+  XCTAssertEqual(field.type, GPBUnknownFieldTypeFixed32);
+  field = fields[1];
+  XCTAssertEqual(field.number, 1);
+  XCTAssertEqual(field.type, GPBUnknownFieldTypeFixed64);
+  [ufs removeField:field];  // Remove the second (fixed64)
+  XCTAssertEqual(ufs.count, 1);
+
+  fields = [ufs fields:1];
+  XCTAssertEqual(fields.count, 1);
+  field = fields[0];
+  XCTAssertEqual(field.number, 1);
+  XCTAssertEqual(field.type, GPBUnknownFieldTypeFixed32);
+
+  field = [[field retain] autorelease];  // Hold on to this last one.
+  [ufs removeField:field];               // Remove the last one (fixed32)
+  XCTAssertEqual(ufs.count, 0);
+
+  // Trying to remove something not in the set should fail.
+  XCTAssertThrowsSpecificNamed([ufs removeField:field], NSException, NSInvalidArgumentException);
+}
+
+- (void)testClearFieldNumber {
+  GPBUnknownFields* ufs = [[[GPBUnknownFields alloc] init] autorelease];
+  [ufs addFieldNumber:1 varint:1];
+  [ufs addFieldNumber:2 fixed32:2];
+  [ufs addFieldNumber:1 fixed64:1];
+  [ufs addFieldNumber:3 varint:3];
+  XCTAssertEqual(ufs.count, 4);
+
+  [ufs clearFieldNumber:999];  // Not present, noop.
+  XCTAssertEqual(ufs.count, 4);
+
+  [ufs clearFieldNumber:1];  // Should remove slot zero and slot two.
+  XCTAssertEqual(ufs.count, 2);
+  NSArray<GPBUnknownField*>* fields = [ufs fields:2];
+  XCTAssertEqual(fields.count, 1);
+  GPBUnknownField* field = fields[0];
+  XCTAssertEqual(field.number, 2);
+  XCTAssertEqual(field.type, GPBUnknownFieldTypeFixed32);
+  XCTAssertEqual(field.fixed32, 2);
+  fields = [ufs fields:3];
+  XCTAssertEqual(fields.count, 1);
+  field = fields[0];
+  XCTAssertEqual(field.number, 3);
+  XCTAssertEqual(field.type, GPBUnknownFieldTypeVarint);
+  XCTAssertEqual(field.varint, 3);
+
+  [ufs clearFieldNumber:2];  // Should remove slot one.
+  fields = [ufs fields:3];
+  XCTAssertEqual(fields.count, 1);
+  field = fields[0];
+  XCTAssertEqual(field.number, 3);
+  XCTAssertEqual(field.type, GPBUnknownFieldTypeVarint);
+  XCTAssertEqual(field.varint, 3);
+}
+
 - (void)testFastEnumeration {
   GPBUnknownFields* ufs = [[[GPBUnknownFields alloc] init] autorelease];
   [ufs addFieldNumber:1 varint:1];