1 files changed, 18 insertions, 19 deletions

diff --git a/objectivec/GPBExtensionRegistry.m b/objectivec/GPBExtensionRegistry.m
index 65534b67..b056a52d 100644
--- a/objectivec/GPBExtensionRegistry.m
+++ b/objectivec/GPBExtensionRegistry.m
@@ -57,14 +57,16 @@
 
 - (instancetype)copyWithZone:(NSZone *)zone {
   GPBExtensionRegistry *result = [[[self class] allocWithZone:zone] init];
-  if (result && mutableClassMap_.count) {
-    [result->mutableClassMap_ addEntriesFromDictionary:mutableClassMap_];
-  }
+  [result addExtensions:self];
   return result;
 }
 
-- (CFMutableDictionaryRef)extensionMapForContainingMessageClass:
-        (Class)containingMessageClass {
+- (void)addExtension:(GPBExtensionDescriptor *)extension {
+  if (extension == nil) {
+    return;
+  }
+
+  Class containingMessageClass = extension.containingMessageClass;
   CFMutableDictionaryRef extensionMap = (CFMutableDictionaryRef)
       [mutableClassMap_ objectForKey:containingMessageClass];
   if (extensionMap == nil) {
@@ -74,18 +76,9 @@
                                              &kCFTypeDictionaryValueCallBacks);
     [mutableClassMap_ setObject:(id)extensionMap
                          forKey:(id<NSCopying>)containingMessageClass];
+    CFRelease(extensionMap);
   }
-  return extensionMap;
-}
 
-- (void)addExtension:(GPBExtensionDescriptor *)extension {
-  if (extension == nil) {
-    return;
-  }
-
-  Class containingMessageClass = extension.containingMessageClass;
-  CFMutableDictionaryRef extensionMap =
-      [self extensionMapForContainingMessageClass:containingMessageClass];
   ssize_t key = extension.fieldNumber;
   CFDictionarySetValue(extensionMap, (const void *)key, extension);
 }
@@ -119,10 +112,16 @@ static void CopyKeyValue(const void *key, const void *value, void *context) {
     Class containingMessageClass = key;
     CFMutableDictionaryRef otherExtensionMap = (CFMutableDictionaryRef)value;
 
-    CFMutableDictionaryRef extensionMap =
-        [self extensionMapForContainingMessageClass:containingMessageClass];
-
-    CFDictionaryApplyFunction(otherExtensionMap, CopyKeyValue, extensionMap);
+    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);
+    }
   }];
 }