erikdoe · dmaclach · Nov 25, 2020
diff --git a/Source/OCMock.xcodeproj/project.pbxproj b/Source/OCMock.xcodeproj/project.pbxproj
@@ -7,8 +7,6 @@
 	objects = {
 
 /* Begin PBXBuildFile section */
-		031E50581BB4A56300E257C3 /* OCMBoxedReturnValueProviderTests.m in Sources */ = {isa = PBXBuildFile; fileRef = 031E50571BB4A56300E257C3 /* OCMBoxedReturnValueProviderTests.m */; };
-		031E50591BB4A56300E257C3 /* OCMBoxedReturnValueProviderTests.m in Sources */ = {isa = PBXBuildFile; fileRef = 031E50571BB4A56300E257C3 /* OCMBoxedReturnValueProviderTests.m */; };
 		0322DA65191188D100CACAF1 /* OCMockObjectVerifyAfterRunTests.m in Sources */ = {isa = PBXBuildFile; fileRef = 0322DA64191188D100CACAF1 /* OCMockObjectVerifyAfterRunTests.m */; };
 		0322DA66191188D100CACAF1 /* OCMockObjectVerifyAfterRunTests.m in Sources */ = {isa = PBXBuildFile; fileRef = 0322DA64191188D100CACAF1 /* OCMockObjectVerifyAfterRunTests.m */; };
 		0322DA6919118B4600CACAF1 /* OCMVerifier.h in Headers */ = {isa = PBXBuildFile; fileRef = 0322DA6719118B4600CACAF1 /* OCMVerifier.h */; settings = {ATTRIBUTES = (Public, ); }; };
@@ -477,7 +475,6 @@
 		030EF0B814632FD000B04273 /* OCMock.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = OCMock.h; sourceTree = "<group>"; };
 		030EF0DC14632FF700B04273 /* libOCMock.a */ = {isa = PBXFileReference; explicitFileType = archive.ar; includeInIndex = 0; path = libOCMock.a; sourceTree = BUILT_PRODUCTS_DIR; };
 		030EF0E114632FF700B04273 /* OCMockLib-Prefix.pch */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = "OCMockLib-Prefix.pch"; sourceTree = "<group>"; };
-		031E50571BB4A56300E257C3 /* OCMBoxedReturnValueProviderTests.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = OCMBoxedReturnValueProviderTests.m; sourceTree = "<group>"; };
 		0322DA64191188D100CACAF1 /* OCMockObjectVerifyAfterRunTests.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = OCMockObjectVerifyAfterRunTests.m; sourceTree = "<group>"; };
 		0322DA6719118B4600CACAF1 /* OCMVerifier.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = OCMVerifier.h; sourceTree = "<group>"; };
 		0322DA6819118B4600CACAF1 /* OCMVerifier.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = OCMVerifier.m; sourceTree = "<group>"; };
@@ -765,7 +762,6 @@
 				8BF73E52246CA75E00B9A52C /* OCMNoEscapeBlockTests.m */,
 				03B316271463350E0052CD09 /* OCMStubRecorderTests.m */,
 				037ECD5318FAD84100AF0E4C /* OCMInvocationMatcherTests.m */,
-				031E50571BB4A56300E257C3 /* OCMBoxedReturnValueProviderTests.m */,
 				03B316211463350E0052CD09 /* OCMConstraintTests.m */,
 				8B11D4B62448E2E900247BE2 /* OCMCPlusPlus98Tests.mm */,
 				8B11D4B92448E53600247BE2 /* OCMCPlusPlus11Tests.mm */,
@@ -1544,7 +1540,6 @@
 				8BF73E53246CA75E00B9A52C /* OCMNoEscapeBlockTests.m in Sources */,
 				8B11D4B72448E2E900247BE2 /* OCMCPlusPlus98Tests.mm in Sources */,
 				2FA28AB33F01A7D980F2C705 /* OCMockObjectDynamicPropertyMockingTests.m in Sources */,
-				031E50581BB4A56300E257C3 /* OCMBoxedReturnValueProviderTests.m in Sources */,
 			);
 			runOnlyForDeploymentPostprocessing = 0;
 		};
@@ -1637,7 +1632,6 @@
 			buildActionMask = 2147483647;
 			files = (
 				3C76716C1BB3EBC500FDC9F4 /* TestClassWithCustomReferenceCounting.m in Sources */,
-				031E50591BB4A56300E257C3 /* OCMBoxedReturnValueProviderTests.m in Sources */,
 				8B3786A924E5BD6400FD1B5B /* OCMFunctionsTests.m in Sources */,
 				03C9CA1E18F05A84006DF94D /* OCMArgTests.m in Sources */,
 				036865651D3571A8005E6BEE /* OCMQuantifierTests.m in Sources */,

diff --git a/Source/OCMock/NSMethodSignature+OCMAdditions.h b/Source/OCMock/NSMethodSignature+OCMAdditions.h
@@ -26,4 +26,8 @@
 - (NSString *)fullTypeString;
 - (const char *)fullObjCTypes;
 
+// True if the return type of the method is "compatible" with the valueType and value.
+// Compatible is defined the same as `OCMIsObjCTypeCompatibleWithValueType`.
+- (BOOL)isMethodReturnTypeCompatibleWithValueType:(const char *)valueType value:(const void *)value valueSize:(size_t)valueSize;
+
 @end
diff --git a/Source/OCMock/NSMethodSignature+OCMAdditions.m b/Source/OCMock/NSMethodSignature+OCMAdditions.m
@@ -178,4 +178,9 @@ - (const char *)fullObjCTypes
     return [[self fullTypeString] UTF8String];
 }
 
+- (BOOL)isMethodReturnTypeCompatibleWithValueType:(const char *)valueType value:(const void *)value valueSize:(size_t)valueSize
+{
+    return OCMIsObjCTypeCompatibleWithValueType([self methodReturnType], valueType, value, valueSize);
+}
+
 @end
diff --git a/Source/OCMock/OCMBlockCaller.h b/Source/OCMock/OCMBlockCaller.h
@@ -18,10 +18,24 @@
 
 @interface OCMBlockCaller : NSObject
 {
-    void (^block)(NSInvocation *);
+    id block;
 }
 
-- (id)initWithCallBlock:(void (^)(NSInvocation *))theBlock;
+/*
+ * Call blocks can have one of four types:
+ * a) A simple block ^{ NSLog(@"hi"); }.
+ * b) The new style ^(id localSelf, type0 arg0, type1 arg1) { ... }
+ *    where types and args match the arguments passed to the selector we are
+ *    stubbing.
+ * c) The deprecated style ^(NSInvocation *anInvocation) { ... }. This case
+ *    cannot have a return value. If a return value is desired it must be set
+ *    on `anInvocation`.
+ * d) nil
+ *
+ * If it is (a) or (b) and there is a return value it must match the return type
+ * of the selector. If there is no return value then the method will return 0.
+ */
+- (id)initWithCallBlock:(id)theBlock;
 
 - (void)handleInvocation:(NSInvocation *)anInvocation;
 

diff --git a/Source/OCMock/OCMBlockCaller.m b/Source/OCMock/OCMBlockCaller.m
@@ -15,11 +15,13 @@
  */
 
 #import "OCMBlockCaller.h"
-
+#import "NSMethodSignature+OCMAdditions.h"
+#import "OCMFunctionsPrivate.h"
+#import "NSInvocation+OCMAdditions.h"
 
 @implementation OCMBlockCaller
 
-- (id)initWithCallBlock:(void (^)(NSInvocation *))theBlock
+-(id)initWithCallBlock:(id)theBlock
 {
     if((self = [super init]))
     {
@@ -37,9 +39,79 @@ - (void)dealloc
 
 - (void)handleInvocation:(NSInvocation *)anInvocation
 {
-    if(block != nil)
+    if(!block)
+    {
+        return;
+    }
+    NSMethodSignature *blockMethodSignature = [NSMethodSignature signatureForBlock:block];
+    NSUInteger blockNumberOfArguments = [blockMethodSignature numberOfArguments];
+    if(blockNumberOfArguments == 2 && strcmp([blockMethodSignature getArgumentTypeAtIndex:1], "@\"NSInvocation\"") == 0)
+    {
+        // This is the deprecated ^(NSInvocation *) {} case.
+        if([blockMethodSignature methodReturnLength] != 0)
+        {
+            [NSException raise:NSInvalidArgumentException format:@"NSInvocation style `andDo:` block for `-%@` cannot have return value.", NSStringFromSelector([anInvocation selector])];
+        }
+
+        void (^theBlock)(NSInvocation *) = block;
+        theBlock(anInvocation);
+        NSLog(@"Warning: Replace `^(NSInvocation *invocation) { ... }` with `%@`.", OCMBlockDeclarationForInvocation(anInvocation));
+        return;
+    }
+
+    // This handles both the ^{} case and the ^(SelfType *a, Arg1Type b, ...) case.
+    NSMethodSignature *invocationMethodSignature = [anInvocation methodSignature];
+    NSInvocation *blockInvocation = [NSInvocation invocationWithMethodSignature:blockMethodSignature];
+    NSUInteger invocationNumberOfArguments = [invocationMethodSignature numberOfArguments];
+    if(blockNumberOfArguments != 1 && blockNumberOfArguments != invocationNumberOfArguments)
+    {
+        [NSException raise:NSInvalidArgumentException format:@"Block style `andDo:` block signature has wrong number of arguments. %d vs %d", (int)invocationNumberOfArguments, (int)blockNumberOfArguments];
+    }
+    id target = [anInvocation target];
+
+    // In the ^{} case, blockNumberOfArguments will be 1, so we will just skip the whole for block.
+    for(NSUInteger argIndex = 1; argIndex < blockNumberOfArguments; ++argIndex)
+    {
+        // Set arg1 to be "localSelf".
+        // Note that in a standard NSInvocation this would be SEL, but blocks don't have SEL args.
+        if(argIndex == 1)
+        {
+            [blockInvocation setArgument:&target atIndex:1];
+            continue;
+        }
+        const char *blockArgType = [blockMethodSignature getArgumentTypeAtIndex:argIndex];
+        const char *invocationArgType = [invocationMethodSignature getArgumentTypeAtIndex:argIndex];
+        NSUInteger argSize;
+        NSGetSizeAndAlignment(blockArgType, &argSize, NULL);
+        NSMutableData *argSpace = [NSMutableData dataWithLength:argSize];
+        void *argBytes = [argSpace mutableBytes];
+        [anInvocation getArgument:argBytes atIndex:argIndex];
+        if(!OCMIsObjCTypeCompatibleWithValueType(invocationArgType, blockArgType, argBytes, argSize) && !OCMEqualTypesAllowingOpaqueStructs(blockArgType, invocationArgType))
+        {
+            [NSException raise:NSInvalidArgumentException format:@"Block style `andDo:` block signature does not match selector signature. Arg %d is `%@` vs `%@`.", (int)argIndex, OCMObjCTypeForArgumentType(blockArgType), OCMObjCTypeForArgumentType(invocationArgType)];
+        }
+        [blockInvocation setArgument:argBytes atIndex:argIndex];
+    }
+    [blockInvocation invokeWithTarget:block];
+    NSUInteger blockReturnLength = [blockMethodSignature methodReturnLength];
+    if(blockReturnLength > 0)
     {
-        block(anInvocation);
+      // If there is a return value, make sure that it matches the expected return type.
+      const char *blockReturnType = [blockMethodSignature methodReturnType];
+      NSUInteger invocationReturnLength = [invocationMethodSignature methodReturnLength];
+      const char *invocationReturnType = [invocationMethodSignature methodReturnType];
+      if(invocationReturnLength != blockReturnLength)
+      {
+          [NSException raise:NSInvalidArgumentException format:@"Block style `andDo:` block signature does not match selector signature. Return type is `%@` vs `%@`.", OCMObjCTypeForArgumentType(blockReturnType), OCMObjCTypeForArgumentType(invocationReturnType)];
+      }
+      NSMutableData *argSpace = [NSMutableData dataWithLength:invocationReturnLength];
+      void *argBytes = [argSpace mutableBytes];
+      [blockInvocation getReturnValue:argBytes];
+      if(!OCMIsObjCTypeCompatibleWithValueType(invocationReturnType, blockReturnType, argBytes, invocationReturnLength) && !OCMEqualTypesAllowingOpaqueStructs(blockReturnType, invocationReturnType))
+      {
+          [NSException raise:NSInvalidArgumentException format:@"Block style `andDo:` block signature does not match selector signature. Return type is `%@` vs `%@`.", OCMObjCTypeForArgumentType(blockReturnType), OCMObjCTypeForArgumentType(invocationReturnType)];
+      }
+      [anInvocation setReturnValue:argBytes];
     }
 }
 

diff --git a/Source/OCMock/OCMBoxedReturnValueProvider.m b/Source/OCMock/OCMBoxedReturnValueProvider.m
@@ -16,8 +16,7 @@
 
 #import "OCMBoxedReturnValueProvider.h"
 #import "NSValue+OCMAdditions.h"
-#import "OCMFunctionsPrivate.h"
-
+#import "NSMethodSignature+OCMAdditions.h"
 
 @implementation OCMBoxedReturnValueProvider
 
@@ -28,11 +27,11 @@ - (void)handleInvocation:(NSInvocation *)anInvocation
     NSGetSizeAndAlignment([returnValueAsNSValue objCType], &valueSize, NULL);
     char valueBuffer[valueSize];
     [returnValueAsNSValue getValue:valueBuffer];
+    NSMethodSignature *signature = [anInvocation methodSignature];
+    const char *returnType = [signature methodReturnType];
+    const char *returnValueType = [returnValueAsNSValue objCType];
 
-    const char *returnType = [[anInvocation methodSignature] methodReturnType];
-
-    if([self isMethodReturnType:returnType compatibleWithValueType:[returnValueAsNSValue objCType]
-                value:valueBuffer valueSize:valueSize])
+    if([signature isMethodReturnTypeCompatibleWithValueType:returnValueType value:valueBuffer valueSize:valueSize])
     {
         [anInvocation setReturnValue:valueBuffer];
     }
@@ -43,21 +42,8 @@ - (void)handleInvocation:(NSInvocation *)anInvocation
     else
     {
         [NSException raise:NSInvalidArgumentException
-                    format:@"Return value cannot be used for method; method signature declares '%s' but value is '%s'.", returnType, [returnValueAsNSValue objCType]];
+                    format:@"Return value cannot be used for method; method signature declares '%s' but value is '%s'.", returnType, returnValueType];
     }
 }
 
-- (BOOL)isMethodReturnType:(const char *)returnType compatibleWithValueType:(const char *)valueType value:(const void *)value valueSize:(size_t)valueSize
-{
-    /* Same types are obviously compatible */
-    if(strcmp(returnType, valueType) == 0)
-        return YES;
-
-    /* Special treatment for nil and Nil */
-    if(strcmp(returnType, @encode(id)) == 0 || strcmp(returnType, @encode(Class)) == 0)
-        return OCMIsNilValue(valueType, value, valueSize);
-
-    return OCMEqualTypesAllowingOpaqueStructs(returnType, valueType);
-}
-
 @end