diff --git a/objectivec/GPBExtensionRegistry.m b/objectivec/GPBExtensionRegistry.m index 01eb761f3..65534b67a 100644 --- a/objectivec/GPBExtensionRegistry.m +++ b/objectivec/GPBExtensionRegistry.m @@ -34,8 +34,6 @@ #import "GPBDescriptor.h" @implementation GPBExtensionRegistry { - // TODO(dmaclach): Reimplement with CFDictionaries that don't use - // objects as keys. NSMutableDictionary *mutableClassMap_; } @@ -65,13 +63,16 @@ return result; } -- (NSMutableDictionary *)extensionMapForContainingMessageClass: +- (CFMutableDictionaryRef)extensionMapForContainingMessageClass: (Class)containingMessageClass { - NSMutableDictionary *extensionMap = + CFMutableDictionaryRef extensionMap = (CFMutableDictionaryRef) [mutableClassMap_ objectForKey:containingMessageClass]; if (extensionMap == nil) { - extensionMap = [NSMutableDictionary dictionary]; - [mutableClassMap_ setObject:extensionMap + // 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)containingMessageClass]; } return extensionMap; @@ -83,17 +84,28 @@ } Class containingMessageClass = extension.containingMessageClass; - NSMutableDictionary *extensionMap = + CFMutableDictionaryRef extensionMap = [self extensionMapForContainingMessageClass:containingMessageClass]; - [extensionMap setObject:extension forKey:@(extension.fieldNumber)]; + ssize_t key = extension.fieldNumber; + CFDictionarySetValue(extensionMap, (const void *)key, extension); } - (GPBExtensionDescriptor *)extensionForDescriptor:(GPBDescriptor *)descriptor fieldNumber:(NSInteger)fieldNumber { Class messageClass = descriptor.messageClass; - NSDictionary *extensionMap = + CFMutableDictionaryRef extensionMap = (CFMutableDictionaryRef) [mutableClassMap_ objectForKey:messageClass]; - return [extensionMap objectForKey:@(fieldNumber)]; + ssize_t key = fieldNumber; + GPBExtensionDescriptor *result = + (extensionMap + ? CFDictionaryGetValue(extensionMap, (const void *)key) + : nil); + return result; +} + +static void CopyKeyValue(const void *key, const void *value, void *context) { + CFMutableDictionaryRef extensionMap = (CFMutableDictionaryRef)context; + CFDictionarySetValue(extensionMap, key, value); } - (void)addExtensions:(GPBExtensionRegistry *)registry { @@ -102,13 +114,16 @@ return; } NSMutableDictionary *otherClassMap = registry->mutableClassMap_; - for (Class containingMessageClass in otherClassMap) { - NSMutableDictionary *extensionMap = + [otherClassMap enumerateKeysAndObjectsUsingBlock:^(id key, id value, BOOL * stop) { +#pragma unused(stop) + Class containingMessageClass = key; + CFMutableDictionaryRef otherExtensionMap = (CFMutableDictionaryRef)value; + + CFMutableDictionaryRef extensionMap = [self extensionMapForContainingMessageClass:containingMessageClass]; - NSMutableDictionary *otherExtensionMap = - [registry extensionMapForContainingMessageClass:containingMessageClass]; - [extensionMap addEntriesFromDictionary:otherExtensionMap]; - } + + CFDictionaryApplyFunction(otherExtensionMap, CopyKeyValue, extensionMap); + }]; } #pragma clang diagnostic pop