Add support for unretained and unsafeunretained arguments for stubs.

Allows marking an argument in a stub as being either 'unsafe' or 'unsafeunretained'.
An unsafe object argument is retained by the stub, but not by invocations on the mock.
An unsafe unretained object argument is not retained by the stub or by invocations on the
mock.

This allows for mocking of methods that do not retain their arguments. This should
simplify testing of methods that are called in dealloc, and give people a way out of
other retain-loop problems.

Source/OCMock.xcodeproj/project.pbxproj

@@ -279,6 +279,16 @@
 		817EB15C1BD765130047E85A /* OCMBlockArgCaller.h in Headers */ = {isa = PBXBuildFile; fileRef = 2FA2891034E7B73AA3511D17 /* OCMBlockArgCaller.h */; };
 		817EB15D1BD765130047E85A /* OCMArgAction.h in Headers */ = {isa = PBXBuildFile; fileRef = 2FA2833B48908EAD36444671 /* OCMArgAction.h */; };
 		817EB1661BD7674D0047E85A /* OCMFunctionsPrivate.h in Headers */ = {isa = PBXBuildFile; fileRef = 03F370CA1BAA1DE800CAD3E8 /* OCMFunctionsPrivate.h */; };
+		8BF740142476E4B400B9A52C /* OCMUnretainedArgument.h in Headers */ = {isa = PBXBuildFile; fileRef = 8BF740122476E4B300B9A52C /* OCMUnretainedArgument.h */; };
+		8BF740152476E4B400B9A52C /* OCMUnretainedArgument.m in Sources */ = {isa = PBXBuildFile; fileRef = 8BF740132476E4B400B9A52C /* OCMUnretainedArgument.m */; };
+		8BF740162476E59A00B9A52C /* OCMUnretainedArgument.m in Sources */ = {isa = PBXBuildFile; fileRef = 8BF740132476E4B400B9A52C /* OCMUnretainedArgument.m */; };
+		8BF740172476E59A00B9A52C /* OCMUnretainedArgument.m in Sources */ = {isa = PBXBuildFile; fileRef = 8BF740132476E4B400B9A52C /* OCMUnretainedArgument.m */; };
+		8BF740182476E59B00B9A52C /* OCMUnretainedArgument.m in Sources */ = {isa = PBXBuildFile; fileRef = 8BF740132476E4B400B9A52C /* OCMUnretainedArgument.m */; };
+		8BF740192476E59C00B9A52C /* OCMUnretainedArgument.m in Sources */ = {isa = PBXBuildFile; fileRef = 8BF740132476E4B400B9A52C /* OCMUnretainedArgument.m */; };
+		8BF7401A24771FD600B9A52C /* OCMUnretainedArgument.h in Headers */ = {isa = PBXBuildFile; fileRef = 8BF740122476E4B300B9A52C /* OCMUnretainedArgument.h */; };
+		8BF7401B24771FD700B9A52C /* OCMUnretainedArgument.h in Headers */ = {isa = PBXBuildFile; fileRef = 8BF740122476E4B300B9A52C /* OCMUnretainedArgument.h */; };
+		8BF7401C24771FD700B9A52C /* OCMUnretainedArgument.h in Headers */ = {isa = PBXBuildFile; fileRef = 8BF740122476E4B300B9A52C /* OCMUnretainedArgument.h */; };
+		8BF7401D24771FD800B9A52C /* OCMUnretainedArgument.h in Headers */ = {isa = PBXBuildFile; fileRef = 8BF740122476E4B300B9A52C /* OCMUnretainedArgument.h */; };
 		8DE97C5522B43EE60098C63F /* OCMockObject.m in Sources */ = {isa = PBXBuildFile; fileRef = 03B3159E146333BF0052CD09 /* OCMockObject.m */; };
 		8DE97C5622B43EE60098C63F /* OCClassMockObject.m in Sources */ = {isa = PBXBuildFile; fileRef = 03B3158C146333BF0052CD09 /* OCClassMockObject.m */; };
 		8DE97C5722B43EE60098C63F /* OCPartialMockObject.m in Sources */ = {isa = PBXBuildFile; fileRef = 03B315AA146333BF0052CD09 /* OCPartialMockObject.m */; };
@@ -569,6 +579,8 @@
 		3CFBDD751BB3DB200050D9C5 /* TestClassWithCustomReferenceCounting.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = TestClassWithCustomReferenceCounting.h; sourceTree = "<group>"; };
 		3CFBDD761BB3DB200050D9C5 /* TestClassWithCustomReferenceCounting.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = TestClassWithCustomReferenceCounting.m; sourceTree = "<group>"; };
 		817EB1621BD765130047E85A /* OCMock.framework */ = {isa = PBXFileReference; explicitFileType = wrapper.framework; includeInIndex = 0; path = OCMock.framework; sourceTree = BUILT_PRODUCTS_DIR; };
+		8BF740122476E4B300B9A52C /* OCMUnretainedArgument.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = OCMUnretainedArgument.h; sourceTree = "<group>"; };
+		8BF740132476E4B400B9A52C /* OCMUnretainedArgument.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = OCMUnretainedArgument.m; sourceTree = "<group>"; };
 		8DE97CA022B43EE60098C63F /* OCMock.framework */ = {isa = PBXFileReference; explicitFileType = wrapper.framework; includeInIndex = 0; path = OCMock.framework; sourceTree = BUILT_PRODUCTS_DIR; };
 		A02926811CA0725A00594AAF /* TestObjects.xcdatamodel */ = {isa = PBXFileReference; lastKnownFileType = wrapper.xcdatamodel; path = TestObjects.xcdatamodel; sourceTree = "<group>"; };
 		D31108AD1828DB8700737925 /* OCMockLibTests.xctest */ = {isa = PBXFileReference; explicitFileType = wrapper.cfbundle; includeInIndex = 0; path = OCMockLibTests.xctest; sourceTree = BUILT_PRODUCTS_DIR; };
@@ -867,6 +879,8 @@
 				03B315A2146333BF0052CD09 /* OCMPassByRefSetter.m */,
 				2FA2891034E7B73AA3511D17 /* OCMBlockArgCaller.h */,
 				2FA283D58AA7569D8A5B0C57 /* OCMBlockArgCaller.m */,
+				8BF740122476E4B300B9A52C /* OCMUnretainedArgument.h */,
+				8BF740132476E4B400B9A52C /* OCMUnretainedArgument.m */,
 			);
 			name = "Argument Constraints and Actions";
 			sourceTree = "<group>";
@@ -957,6 +971,7 @@
 				03B315F5146333C00052CD09 /* OCMPassByRefSetter.h in Headers */,
 				03B315FA146333C00052CD09 /* OCMRealObjectForwarder.h in Headers */,
 				03B315FF146333C00052CD09 /* OCMObjectReturnValueProvider.h in Headers */,
+				8BF740142476E4B400B9A52C /* OCMUnretainedArgument.h in Headers */,
 				03B31604146333C00052CD09 /* OCObserverMockObject.h in Headers */,
 				03B31609146333C00052CD09 /* OCPartialMockObject.h in Headers */,
 				0368656D1D357317005E6BEE /* OCMQuantifier.h in Headers */,
@@ -1006,6 +1021,7 @@
 				817EB1661BD7674D0047E85A /* OCMFunctionsPrivate.h in Headers */,
 				03B31605146333C00052CD09 /* OCObserverMockObject.h in Headers */,
 				03B3160A146333C00052CD09 /* OCPartialMockObject.h in Headers */,
+				8BF7401A24771FD600B9A52C /* OCMUnretainedArgument.h in Headers */,
 				03B31614146333C00052CD09 /* OCProtocolMockObject.h in Headers */,
 				2FA28E1EB6B8536785258DF5 /* OCMInvocationMatcher.h in Headers */,
 				0322DA6A19118B4600CACAF1 /* OCMVerifier.h in Headers */,
@@ -1057,6 +1073,7 @@
 				817EB1591BD765130047E85A /* NSObject+OCMAdditions.h in Headers */,
 				817EB15A1BD765130047E85A /* NSValue+OCMAdditions.h in Headers */,
 				817EB15B1BD765130047E85A /* OCMFunctions.h in Headers */,
+				8BF7401C24771FD700B9A52C /* OCMUnretainedArgument.h in Headers */,
 				817EB15C1BD765130047E85A /* OCMBlockArgCaller.h in Headers */,
 				817EB15D1BD765130047E85A /* OCMArgAction.h in Headers */,
 				2FA28806443827E286F12F6F /* OCMNonRetainingObjectReturnValueProvider.h in Headers */,
@@ -1089,6 +1106,7 @@
 				8DE97C8C22B43EE60098C63F /* OCMBoxedReturnValueProvider.h in Headers */,
 				8DE97C8D22B43EE60098C63F /* OCMExceptionReturnValueProvider.h in Headers */,
 				8DE97C8E22B43EE60098C63F /* OCMIndirectReturnValueProvider.h in Headers */,
+				8BF7401D24771FD800B9A52C /* OCMUnretainedArgument.h in Headers */,
 				8DE97C8F22B43EE60098C63F /* OCMNotificationPoster.h in Headers */,
 				8DE97C9022B43EE60098C63F /* OCMObjectReturnValueProvider.h in Headers */,
 				8DE97C9122B43EE60098C63F /* OCMFunctionsPrivate.h in Headers */,
@@ -1146,6 +1164,7 @@
 				F0B951481B00810C00942C38 /* NSObject+OCMAdditions.h in Headers */,
 				F0B951491B00810C00942C38 /* NSValue+OCMAdditions.h in Headers */,
 				F0B9514A1B00810C00942C38 /* OCMFunctions.h in Headers */,
+				8BF7401B24771FD700B9A52C /* OCMUnretainedArgument.h in Headers */,
 				2FA28B7BDB3319A499E90525 /* OCMBlockArgCaller.h in Headers */,
 				2FA280E60213BA09F007C173 /* OCMArgAction.h in Headers */,
 				2FA28AFBD67EAB9DD1F23BF5 /* OCMNonRetainingObjectReturnValueProvider.h in Headers */,
@@ -1409,6 +1428,7 @@
 				03B315CA146333BF0052CD09 /* OCMBlockCaller.m in Sources */,
 				036865681D3572ED005E6BEE /* OCMQuantifier.m in Sources */,
 				03B315CF146333BF0052CD09 /* OCMBoxedReturnValueProvider.m in Sources */,
+				8BF740152476E4B400B9A52C /* OCMUnretainedArgument.m in Sources */,
 				03B315D4146333BF0052CD09 /* OCMConstraint.m in Sources */,
 				03B315D9146333BF0052CD09 /* OCMExceptionReturnValueProvider.m in Sources */,
 				03B315DE146333BF0052CD09 /* OCMIndirectReturnValueProvider.m in Sources */,
@@ -1451,6 +1471,7 @@
 				03B315CC146333BF0052CD09 /* OCMBlockCaller.m in Sources */,
 				036865691D3572ED005E6BEE /* OCMQuantifier.m in Sources */,
 				03B315D1146333BF0052CD09 /* OCMBoxedReturnValueProvider.m in Sources */,
+				8BF740162476E59A00B9A52C /* OCMUnretainedArgument.m in Sources */,
 				03DCED6D183406BC0059089E /* NSObject+OCMAdditions.m in Sources */,
 				03B315D6146333BF0052CD09 /* OCMConstraint.m in Sources */,
 				03B315DB146333BF0052CD09 /* OCMExceptionReturnValueProvider.m in Sources */,
@@ -1521,6 +1542,7 @@
 				817EB11F1BD765130047E85A /* OCMInvocationMatcher.m in Sources */,
 				0368656B1D3572ED005E6BEE /* OCMQuantifier.m in Sources */,
 				817EB1201BD765130047E85A /* OCMInvocationStub.m in Sources */,
+				8BF740182476E59B00B9A52C /* OCMUnretainedArgument.m in Sources */,
 				817EB1211BD765130047E85A /* OCMInvocationExpectation.m in Sources */,
 				817EB1221BD765130047E85A /* OCMRealObjectForwarder.m in Sources */,
 				817EB1231BD765130047E85A /* OCMBlockCaller.m in Sources */,
@@ -1563,6 +1585,7 @@
 				8DE97C5C22B43EE60098C63F /* OCMVerifier.m in Sources */,
 				8DE97C5D22B43EE60098C63F /* OCMInvocationMatcher.m in Sources */,
 				8DE97C5E22B43EE60098C63F /* OCMInvocationStub.m in Sources */,
+				8BF740192476E59C00B9A52C /* OCMUnretainedArgument.m in Sources */,
 				8DE97C5F22B43EE60098C63F /* OCMInvocationExpectation.m in Sources */,
 				8DE97C6022B43EE60098C63F /* OCMRealObjectForwarder.m in Sources */,
 				8DE97C6122B43EE60098C63F /* OCMBlockCaller.m in Sources */,
@@ -1633,6 +1656,7 @@
 				F0B951141B0080EC00942C38 /* OCMInvocationMatcher.m in Sources */,
 				0368656A1D3572ED005E6BEE /* OCMQuantifier.m in Sources */,
 				F0B951151B0080EC00942C38 /* OCMInvocationStub.m in Sources */,
+				8BF740172476E59A00B9A52C /* OCMUnretainedArgument.m in Sources */,
 				F0B951161B0080EC00942C38 /* OCMInvocationExpectation.m in Sources */,
 				F0B951171B0080EC00942C38 /* OCMRealObjectForwarder.m in Sources */,
 				F0B951181B0080EC00942C38 /* OCMBlockCaller.m in Sources */,

Source/OCMock/NSInvocation+OCMAdditions.h

@@ -20,7 +20,7 @@
 
 + (NSInvocation *)invocationForBlock:(id)block withArguments:(NSArray *)arguments;
 
-- (void)retainObjectArgumentsExcludingObject:(id)objectToExclude;
+- (void)retainObjectArgumentsExcludingObject:(id)objectToExclude excludingObjectsAtIndexes:(NSIndexSet *)indexes;
 
 - (id)getArgumentAtIndexAsObject:(NSInteger)argIndex;
 

Source/OCMock/NSInvocation+OCMAdditions.m

@@ -55,12 +55,11 @@ + (NSInvocation *)invocationForBlock:(id)block withArguments:(NSArray *)argument
 
 static NSString *const OCMRetainedObjectArgumentsKey = @"OCMRetainedObjectArgumentsKey";
 
-- (void)retainObjectArgumentsExcludingObject:(id)objectToExclude
+- (void)retainObjectArgumentsExcludingObject:(id)objectToExclude excludingObjectsAtIndexes:(NSIndexSet *)indexes;
 {
-    if(objc_getAssociatedObject(self, OCMRetainedObjectArgumentsKey) != nil)
+    if (objc_getAssociatedObject(self, OCMRetainedObjectArgumentsKey) != nil)
     {
-        // looks like we've retained the arguments already; do nothing else
-        return;
+        [NSException raise:NSInternalInconsistencyException format:@"Arguments have already been retained"];
     }
 
     NSMutableArray *retainedArguments = [[NSMutableArray alloc] init];
@@ -80,7 +79,16 @@ - (void)retainObjectArgumentsExcludingObject:(id)objectToExclude
     for(NSUInteger index = 2; index < numberOfArguments; index++)
     {
         const char *argumentType = [[self methodSignature] getArgumentTypeAtIndex:index];
-        if(OCMIsObjectType(argumentType))
+        BOOL isObjectType = OCMIsObjectType(argumentType);
+        if ([indexes containsIndex:index])
+        {
+            if (!isObjectType)
+            {
+                [NSException raise:NSInternalInconsistencyException format:@"Argument at %d is not an object", (int)index];
+            }
+            continue;
+        }
+        if (isObjectType)
         {
             id argument;
             [self getArgument:&argument atIndex:index];
@@ -106,18 +114,9 @@ - (void)retainObjectArgumentsExcludingObject:(id)objectToExclude
     {
         id returnValue;
         [self getReturnValue:&returnValue];
-        if((returnValue != nil) && (returnValue != objectToExclude))
+        if (returnValue != nil)
         {
-            if(OCMIsBlockType(returnType))
-            {
-                id blockReturnValue = [returnValue copy];
-                [retainedArguments addObject:blockReturnValue];
-                [blockReturnValue release];
-            }
-            else
-            {
-                [retainedArguments addObject:returnValue];
-            }
+            [NSException raise:NSInternalInconsistencyException format:@"Return value should never be set"];
         }
     }
 

Source/OCMock/OCMArg.h

@@ -32,6 +32,22 @@
 + (id)checkWithSelector:(SEL)selector onObject:(id)anObject;
 + (id)checkWithBlock:(BOOL (^)(id obj))block;
 
+// Unretained object arguments are not retained by invocations on the mock, but are retained by the
+// stub itself. A use case for this is when you are stubbing an argument to a method that does not
+// retain its argument using an `OCMArg` variant that you do not want to keep a reference to.
+// See `OCMOCK_ANY_UNRETAINED`.
++ (id)unretainedObject:(id)anObject;
+
+// Unsafe unretained object arguments are not retained by invocations on the mock or by the stub.
+// A potential use case for this is when you are stubbing methods that do not retain their
+// arguments and you want to verify dealloc conditions. An example of this would be verifying
+// KVO registration/deregistration that occurs in the init/dealloc of an object. If the object were
+// retained by the mocking system in any way you would never see the deregistration.
+// Note that you *must* keep a reference to anObject outside this call or you will crash.
+// Something like `[OCMArg unsafeUnretainedObject:[[Foo alloc] init]]` under ARC is a guaranteed
+// dangling pointer problem.
++ (id)unsafeUnretainedObject:(id)anObject;
+
 // manipulating arguments
 
 + (id *)setTo:(id)value;
@@ -49,6 +65,11 @@
 
 #define OCMOCK_ANY [OCMArg any]
 
+// See comments on [OCMArg unretainedObject] and [OCMArg unsafeUnretainedObject].
+#define OCMOCK_UNSAFE_UNRETAINED(x) [OCMArg unsafeUnretainedObject:(x)]
+#define OCMOCK_UNRETAINED(x) [OCMArg unretainedObject:(x)]
+#define OCMOCK_ANY_UNRETAINED OCMOCK_UNRETAINED(OCMOCK_ANY)
+
 #if defined(__GNUC__) && !defined(__STRICT_ANSI__)
   #define OCMOCK_VALUE(variable) \
     ({ __typeof__(variable) __v = (variable); [NSValue value:&__v withObjCType:@encode(__typeof__(__v))]; })

Source/OCMock/OCMArg.m

@@ -19,6 +19,7 @@
 #import <OCMock/OCMConstraint.h>
 #import "OCMPassByRefSetter.h"
 #import "OCMBlockArgCaller.h"
+#import "OCMUnretainedArgument.h"
 
 @implementation OCMArg
 
@@ -81,6 +82,16 @@ + (id)checkWithBlock:(BOOL (^)(id))block
 	return [[[OCMBlockConstraint alloc] initWithConstraintBlock:block] autorelease];
 }
 
++ (id)unretainedObject:(id)anObject
+{
+  return [[[OCMUnretainedArgument alloc] initWithObject:anObject safe:YES] autorelease];
+}
+
++ (id)unsafeUnretainedObject:(id)anObject
+{
+  return [[[OCMUnretainedArgument alloc] initWithObject:anObject safe:NO] autorelease];
+}
+
 + (id *)setTo:(id)value
 {
 	return (id *)[[[OCMPassByRefSetter alloc] initWithValue:value] autorelease];

Source/OCMock/OCMInvocationMatcher.h

@@ -21,6 +21,7 @@
     NSInvocation *recordedInvocation;
     BOOL         recordedAsClassMethod;
     BOOL         ignoreNonObjectArgs;
+    NSIndexSet   *unretainedArgumentIndexes;
 }
 
 - (void)setInvocation:(NSInvocation *)anInvocation;
@@ -34,4 +35,6 @@
 - (BOOL)matchesSelector:(SEL)aSelector;
 - (BOOL)matchesInvocation:(NSInvocation *)anInvocation;
 
+- (NSIndexSet *)unretainedArgumentIndexes;
+
 @end

Source/OCMock/OCMInvocationMatcher.m

@@ -21,6 +21,7 @@
 #import "NSInvocation+OCMAdditions.h"
 #import "OCMInvocationMatcher.h"
 #import "OCMFunctionsPrivate.h"
+#import "OCMUnretainedArgument.h"
 
 
 @interface NSObject(HCMatcherDummy)
@@ -33,17 +34,53 @@ @implementation OCMInvocationMatcher
 - (void)dealloc
 {
     [recordedInvocation release];
+    [unretainedArgumentIndexes release];
     [super dealloc];
 }
 
+- (NSIndexSet *)unretainedArgumentIndexes
+{
+    return unretainedArgumentIndexes;
+}
+
 - (void)setInvocation:(NSInvocation *)anInvocation
 {
-    [recordedInvocation release];
+    if (recordedInvocation != nil)
+    {
+        [NSException raise:NSInternalInconsistencyException format:@"Invocation should only be set once on a matcher"];
+    }
+
+    // Strip any "unretained arguments" from the invocation and record them.
+    NSMutableIndexSet *unretainedIndexes = [NSMutableIndexSet indexSet];
+    NSMutableIndexSet *unsafeUnretainedIndexes = [NSMutableIndexSet indexSet];
+    NSMethodSignature *signature = [anInvocation methodSignature];
+    NSUInteger n = [signature numberOfArguments];
+    for(NSUInteger i = 2; i < n; i++)
+    {
+        const char *argType = [signature getArgumentTypeAtIndex:i];
+        if (OCMIsObjectType(argType)) 
+        {
+            OCMUnretainedArgument *unretainedArg;
+            [anInvocation getArgument:&unretainedArg atIndex:i];
+            if ([unretainedArg isKindOfClass:[OCMUnretainedArgument class]])
+             {
+                [unretainedIndexes addIndex:i];
+                if (![unretainedArg isSafe]) 
+                {
+                    [unsafeUnretainedIndexes addIndex:i];
+                }
+                id realArg = [unretainedArg object];
+                [anInvocation setArgument:&realArg atIndex:i];
+            }
+        }
+    }
+    unretainedArgumentIndexes = [unretainedIndexes copy];
+
     // Don't do a regular -retainArguments on the invocation that we use for matching. NSInvocation
     // effectively does an strcpy on char* arguments which messes up matching them literally and blows
     // up with anyPointer (in strlen since it's not actually a C string). Also on the off-chance that
     // anInvocation contains self as an argument, -retainArguments would create a retain cycle.
-    [anInvocation retainObjectArgumentsExcludingObject:self];
+    [anInvocation retainObjectArgumentsExcludingObject:self excludingObjectsAtIndexes:unsafeUnretainedIndexes];
     recordedInvocation = [anInvocation retain];
 }
 

Source/OCMock/OCMUnretainedArgument.h

@@ -0,0 +1,30 @@
+/*
+ *  Copyright (c) 2015-2020 Erik Doernenburg and contributors
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"); you may
+ *  not use these files except in compliance with the License. You may obtain
+ *  a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ *  License for the specific language governing permissions and limitations
+ *  under the License.
+ */
+
+#import <Foundation/Foundation.h>
+
+// Do not use directly. See methods and comments in OCMArg.h for usage.
+@interface OCMUnretainedArgument : NSObject
+{
+    id object;
+    BOOL isSafe;
+}
+
+- (instancetype)initWithObject:(id)anObject safe:(BOOL)safe;
+- (id)object;
+- (BOOL)isSafe;
+
+@end