From 5911099659dc1d15dd6b1967a6b5afbdc8563db8 Mon Sep 17 00:00:00 2001
From: Thomas Van Lenten <thomasvl@google.com>
Date: Fri, 19 Feb 2021 09:10:47 -0500
Subject: [PATCH] Move the class map to a CFDictionary.

Since the keys are `Class`-s, there's no need to hash/copy/etc. them. This
avoids causing `+initialize` on the classes just when building up a registry.
---
 objectivec/GPBExtensionRegistry.m | 49 ++++++++++++++++---------------
 1 file changed, 25 insertions(+), 24 deletions(-)

diff --git a/objectivec/GPBExtensionRegistry.m b/objectivec/GPBExtensionRegistry.m
index b056a52d95be..e3ff7c4059e8 100644
--- a/objectivec/GPBExtensionRegistry.m
+++ b/objectivec/GPBExtensionRegistry.m
@@ -34,18 +34,20 @@
 #import "GPBDescriptor.h"
 
 @implementation GPBExtensionRegistry {
-  NSMutableDictionary *mutableClassMap_;
+  CFMutableDictionaryRef mutableClassMap_;
 }
 
 - (instancetype)init {
   if ((self = [super init])) {
-    mutableClassMap_ = [[NSMutableDictionary alloc] init];
+    // The keys are ObjC classes, so straight up ptr comparisons are fine.
+    mutableClassMap_ = CFDictionaryCreateMutable(kCFAllocatorDefault, 0, NULL,
+                                             &kCFTypeDictionaryValueCallBacks);
   }
   return self;
 }
 
 - (void)dealloc {
-  [mutableClassMap_ release];
+  CFRelease(mutableClassMap_);
   [super dealloc];
 }
 
@@ -68,14 +70,13 @@ - (void)addExtension:(GPBExtensionDescriptor *)extension {
 
   Class containingMessageClass = extension.containingMessageClass;
   CFMutableDictionaryRef extensionMap = (CFMutableDictionaryRef)
-      [mutableClassMap_ objectForKey:containingMessageClass];
+      CFDictionaryGetValue(mutableClassMap_, containingMessageClass);
   if (extensionMap == nil) {
     // Use a custom dictionary here because the keys are numbers and conversion
     // back and forth from NSNumber isn't worth the cost.
     extensionMap = CFDictionaryCreateMutable(kCFAllocatorDefault, 0, NULL,
                                              &kCFTypeDictionaryValueCallBacks);
-    [mutableClassMap_ setObject:(id)extensionMap
-                         forKey:(id<NSCopying>)containingMessageClass];
+    CFDictionarySetValue(mutableClassMap_, containingMessageClass, extensionMap);
     CFRelease(extensionMap);
   }
 
@@ -87,7 +88,7 @@ - (GPBExtensionDescriptor *)extensionForDescriptor:(GPBDescriptor *)descriptor
                                        fieldNumber:(NSInteger)fieldNumber {
   Class messageClass = descriptor.messageClass;
   CFMutableDictionaryRef extensionMap = (CFMutableDictionaryRef)
-      [mutableClassMap_ objectForKey:messageClass];
+      CFDictionaryGetValue(mutableClassMap_, messageClass);
   ssize_t key = fieldNumber;
   GPBExtensionDescriptor *result =
       (extensionMap
@@ -101,28 +102,28 @@ static void CopyKeyValue(const void *key, const void *value, void *context) {
   CFDictionarySetValue(extensionMap, key, value);
 }
 
+static void CopySubDictionary(const void *key, const void *value, void *context) {
+  CFMutableDictionaryRef mutableClassMap = (CFMutableDictionaryRef)context;
+  Class containingMessageClass = key;
+  CFMutableDictionaryRef otherExtensionMap = (CFMutableDictionaryRef)value;
+
+  CFMutableDictionaryRef extensionMap = (CFMutableDictionaryRef)
+      CFDictionaryGetValue(mutableClassMap, containingMessageClass);
+  if (extensionMap == nil) {
+    extensionMap = CFDictionaryCreateMutableCopy(kCFAllocatorDefault, 0, otherExtensionMap);
+    CFDictionarySetValue(mutableClassMap, containingMessageClass, extensionMap);
+    CFRelease(extensionMap);
+  } else {
+    CFDictionaryApplyFunction(otherExtensionMap, CopyKeyValue, extensionMap);
+  }
+}
+
 - (void)addExtensions:(GPBExtensionRegistry *)registry {
   if (registry == nil) {
     // In the case where there are no extensions just ignore.
     return;
   }
-  NSMutableDictionary *otherClassMap = registry->mutableClassMap_;
-  [otherClassMap enumerateKeysAndObjectsUsingBlock:^(id key, id value, BOOL * stop) {
-#pragma unused(stop)
-    Class containingMessageClass = key;
-    CFMutableDictionaryRef otherExtensionMap = (CFMutableDictionaryRef)value;
-
-    CFMutableDictionaryRef extensionMap = (CFMutableDictionaryRef)
-        [mutableClassMap_ objectForKey:containingMessageClass];
-    if (extensionMap == nil) {
-      extensionMap = CFDictionaryCreateMutableCopy(kCFAllocatorDefault, 0, otherExtensionMap);
-      [mutableClassMap_ setObject:(id)extensionMap
-                           forKey:(id<NSCopying>)containingMessageClass];
-      CFRelease(extensionMap);
-    } else {
-      CFDictionaryApplyFunction(otherExtensionMap, CopyKeyValue, extensionMap);
-    }
-  }];
+  CFDictionaryApplyFunction(registry->mutableClassMap_, CopySubDictionary, mutableClassMap_);
 }
 
 #pragma clang diagnostic pop