Caching whether methods needs structure return.

Author: erikdoe

Source/OCMock/NSObject+OCMAdditions.m

@@ -20,14 +20,33 @@
 
 @implementation NSObject(OCMAdditions)
 
+static NSMutableDictionary *_OCMReturnTypeCache;
+
 + (IMP)instanceMethodForwarderForSelector:(SEL)aSelector
 {
-    // use NSSelectorFromString and not @selector to avoid warning
-    SEL selectorWithNoImplementation = NSSelectorFromString(@"methodWhichMustNotExist::::");
+    // use sel_registerName() and not @selector to avoid warning
+    SEL selectorWithNoImplementation = sel_registerName("methodWhichMustNotExist::::");
 
 #ifndef __arm64__
-    NSMethodSignature *sig = [self instanceMethodSignatureForSelector:aSelector];
-    if([sig usesSpecialStructureReturn])
+    if(_OCMReturnTypeCache == nil)
+        _OCMReturnTypeCache = [[NSMutableDictionary alloc] init];
+
+    BOOL needsStructureReturn;
+    NSString *cacheKey = [NSString stringWithFormat:@"%p-%p", (void *)self, aSelector];
+    NSNumber *cachedValue = [_OCMReturnTypeCache objectForKey:cacheKey];
+
+    if(cachedValue == nil)
+    {
+        NSMethodSignature *sig = [self instanceMethodSignatureForSelector:aSelector];
+        needsStructureReturn = [sig usesSpecialStructureReturn];
+        [_OCMReturnTypeCache setObject:[NSNumber numberWithBool:needsStructureReturn] forKey:cacheKey];
+    }
+    else
+    {
+        needsStructureReturn = [cachedValue boolValue];
+    }
+
+    if(needsStructureReturn)
         return class_getMethodImplementation_stret([NSObject class], selectorWithNoImplementation);
 #endif
 

Source/OCMock/OCClassMockObject.m

@@ -142,7 +142,7 @@ - (void)setupForwarderForClassMethodSelector:(SEL)selector
     const char *types = method_getTypeEncoding(originalMethod);
 
     Class metaClass = object_getClass(mockedClass);
-    IMP forwarderIMP = [metaClass instanceMethodForwarderForSelector:selector];
+    IMP forwarderIMP = [originalMetaClass instanceMethodForwarderForSelector:selector];
     class_replaceMethod(metaClass, selector, forwarderIMP, types);
     class_addMethod(metaClass, aliasSelector, originalIMP, types);
 }

Source/OCMock/OCPartialMockObject.m

@@ -162,7 +162,7 @@ - (void)setupForwarderForSelector:(SEL)sel
         types = ([[mockedClass instanceMethodSignatureForSelector:sel] fullObjCTypes]);
 
     Class subclass = object_getClass([self realObject]);
-    IMP forwarderIMP = [subclass instanceMethodForwarderForSelector:sel];
+    IMP forwarderIMP = [mockedClass instanceMethodForwarderForSelector:sel];
     class_replaceMethod(subclass, sel, forwarderIMP, types);
 	class_addMethod(subclass, aliasSelector, originalIMP, types);
 }