Add support for controlling retain/copy semantics for arguments to stubs.

Allows marking an argument in a stub as having various semantics:
  - is not retained by invocations
    Object arguments are retained by default in OCMock. In some cases to avoid retain
    loops you need to mark an argument as unretained.
  - is not retained by stub
    Stub arguments are retained by default in OCMock. In some specialized cases you
    do not want the stub arguments retained
  - is copied by invocation
    Some arguments have copy semantics and we need the invocation to copy the argument
    instead of retain it.

Author: dmaclach

Source/OCMock/NSInvocation+OCMAdditions.h

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

Source/OCMock/NSInvocation+OCMAdditions.m

@@ -53,9 +53,20 @@ + (NSInvocation *)invocationForBlock:(id)block withArguments:(NSArray *)argument
 }
 
 
+- (OCMConstraintOptions)getArgumentContraintOptionsForArgumentAtIndex:(NSUInteger)index
+{
+    id argument;
+    [self getArgument:&argument atIndex:index];
+    if(![argument isProxy] && [argument isKindOfClass:[OCMConstraint class]])
+    {
+        return [(OCMConstraint *)argument constraintOptions];
+    }
+    return OCMConstraintDefaultOptions;
+}
+
 static NSString *const OCMRetainedObjectArgumentsKey = @"OCMRetainedObjectArgumentsKey";
 
-- (void)retainObjectArgumentsExcludingObject:(id)objectToExclude excludingObjectsAtIndexes:(NSIndexSet *)indexes;
+- (void)applyConstraintOptionsFromStubInvocation:(NSInvocation *)stubInvocation excludingObject:(id)objectToExclude
 {
     if(objc_getAssociatedObject(self, OCMRetainedObjectArgumentsKey) != nil)
     {
@@ -80,16 +91,7 @@ - (void)retainObjectArgumentsExcludingObject:(id)objectToExclude excludingObject
     for(NSUInteger index = 2; index < numberOfArguments; index++)
     {
         const char *argumentType = [[self methodSignature] getArgumentTypeAtIndex:index];
-        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)
+        if (OCMIsObjectType(argumentType))
         {
             id argument;
             [self getArgument:&argument atIndex:index];
@@ -118,7 +120,21 @@ - (void)retainObjectArgumentsExcludingObject:(id)objectToExclude excludingObject
                 }
                 else
                 {
-                	[retainedArguments addObject:argument];
+                    // Conform to the constraintOptions in the stub (if any).
+                    OCMConstraintOptions constraintOptions = [stubInvocation getArgumentContraintOptionsForArgumentAtIndex:index];
+                    if((constraintOptions & OCMConstraintCopyInvocationArg))
+                    {
+                        // Copy not only retains the copy in our array
+                        // but updates the arg in the invocation that we store.
+                        id argCopy = [argument copy];
+                        [retainedArguments addObject:argCopy];
+                        [self setArgument:&argCopy atIndex:index];
+                        [argCopy release];
+                    }
+                    else if(!(constraintOptions & OCMConstraintDoNotRetainInvocationArg))
+                    {
+                        [retainedArguments addObject:argument];
+                    }
                 }
             }
         }

Source/OCMock/OCMArg.h

@@ -16,10 +16,37 @@
 
 #import <Foundation/Foundation.h>
 
+// Options for controlling how OCMArgs function.
+typedef NS_OPTIONS(NSUInteger, OCMArgOptions) {
+  // The OCMArg will retain/release the value passed to it, and invocations on a stub that has 
+  // arguments that the OCMArg is constraining will retain the values passed to them for the
+  // arguments being constrained by the OCMArg.
+  OCMArgDefaultOptions = 0UL,
+
+  // The OCMArg will not retain/release the value passed to it. Is only applicable for 
+  // `isEqual:options:` and `isNotEqual:options`. The caller is responsible for making sure that the 
+  // arg is valid for the required lifetime. Note that unless `OCMArgDoNotRetainInvocationArg` is
+  // also specified, invocations of the stub that the OCMArg arg is constraining will retain values 
+  // passed to them for the arguments being constrained by the OCMArg. `OCMArgNeverRetainArg` is
+  // usually what you want to use.
+  OCMArgDoNotRetainStubArg = (1UL << 0),
+
+  // Invocations on a stub that has arguments that the OCMArg is constraining will retain/release
+  // the values passed to them for the arguments being constrained by the OCMArg. 
+  OCMArgDoNotRetainInvocationArg = (1UL << 1),
+  
+  // Invocations on a stub that has arguments that the OCMArg is constraining will copy/release
+  // the values passed to them for the arguments being constrained by the OCMArg. 
+  OCMArgCopyInvocationArg = (1UL << 2),
+  
+  OCMArgNeverRetainArg = OCMArgDoNotRetainStubArg | OCMArgDoNotRetainInvocationArg,
+};
+
 @interface OCMArg : NSObject 
 
 // constraining arguments
 
+// constrain using OCMArgDefaultOptions
 + (id)any;
 + (SEL)anySelector;
 + (void *)anyPointer;
@@ -32,21 +59,14 @@
 + (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;
++ (id)anyWithOptions:(OCMArgOptions)options;
++ (id)isNilWithOptions:(OCMArgOptions)options;
++ (id)isNotNilWithOptions:(OCMArgOptions)options;
++ (id)isEqual:(id)value options:(OCMArgOptions)options;
++ (id)isNotEqual:(id)value options:(OCMArgOptions)options;
++ (id)isKindOfClass:(Class)cls options:(OCMArgOptions)options;
++ (id)checkWithSelector:(SEL)selector onObject:(id)anObject options:(OCMArgOptions)options;
++ (id)checkWithOptions:(OCMArgOptions)options withBlock:(BOOL (^)(id obj))block;
 
 // manipulating arguments
 
@@ -61,9 +81,6 @@
 
 + (id)resolveSpecialValues:(NSValue *)value;
 
-// Return YES if `object` is either an unretained or an unsafe unretained object.
-+ (BOOL)isUnretained:(id)object;
-
 @end
 
 #define OCMOCK_ANY [OCMArg any]

Source/OCMock/OCMArg.m

@@ -25,7 +25,7 @@ @implementation OCMArg
 
 + (id)any
 {
-	return [[[OCMAnyConstraint alloc] init] autorelease];
+	return [self anyWithOptions:OCMArgDefaultOptions];
 }
 
 + (void *)anyPointer
@@ -45,39 +45,79 @@ + (SEL)anySelector
 
 + (id)isNil
 {
-	return [OCMIsNilConstraint constraint];
+	return [self isNilWithOptions:OCMArgDefaultOptions];
 }
 
 + (id)isNotNil
 {
-	return [OCMIsNotNilConstraint constraint];
+	return [self isNotNilWithOptions:OCMArgDefaultOptions];
 }
 
 + (id)isEqual:(id)value
 {
-	return [[[OCMIsEqualConstraint alloc] initWithTestValue:value] autorelease];
+  return [self isEqual:value options:OCMArgDefaultOptions];
 }
 
 + (id)isNotEqual:(id)value
 {
-	return [[[OCMIsNotEqualConstraint alloc] initWithTestValue:value] autorelease];
+  return [self isNotEqual:value options:OCMArgDefaultOptions];
 }
 
 + (id)isKindOfClass:(Class)cls
 {
-	return [[[OCMBlockConstraint alloc] initWithConstraintBlock:^BOOL(id obj) {
-        return [obj isKindOfClass:cls];
-    }] autorelease];
+  return [self isKindOfClass:cls options:OCMArgDefaultOptions];
 }
 
 + (id)checkWithSelector:(SEL)selector onObject:(id)anObject
 {
-	return [OCMConstraint constraintWithSelector:selector onObject:anObject];
+  return [self checkWithSelector:selector onObject:anObject options:OCMArgDefaultOptions];
 }
 
 + (id)checkWithBlock:(BOOL (^)(id))block
 {
-	return [[[OCMBlockConstraint alloc] initWithConstraintBlock:block] autorelease];
+  return [self checkWithOptions:OCMArgDefaultOptions withBlock:block];
+}
+
++ (id)anyWithOptions:(OCMArgOptions)options
+{
+  return [[[OCMAnyConstraint alloc] initWithOptions:[self constraintOptionsFromArgOptions:options]] autorelease];
+}
+
++ (id)isNilWithOptions:(OCMArgOptions)options
+{
+  return [[[OCMIsEqualConstraint alloc] initWithTestValue:nil options:[self constraintOptionsFromArgOptions:options]] autorelease];
+}
+
++ (id)isNotNilWithOptions:(OCMArgOptions)options
+{
+  return [[[OCMIsNotEqualConstraint alloc] initWithTestValue:nil options:[self constraintOptionsFromArgOptions:options]] autorelease];
+}
+
++ (id)isEqual:(id)value options:(OCMArgOptions)options
+{
+    return [[[OCMIsEqualConstraint alloc] initWithTestValue:value options:[self constraintOptionsFromArgOptions:options]] autorelease];
+}
+
++ (id)isNotEqual:(id)value options:(OCMArgOptions)options
+{
+  return [[[OCMIsNotEqualConstraint alloc] initWithTestValue:value options:[self constraintOptionsFromArgOptions:options]] autorelease];
+}
+
++ (id)isKindOfClass:(Class)cls options:(OCMArgOptions)options
+{
+  return [[[OCMBlockConstraint alloc] initWithOptions:[self constraintOptionsFromArgOptions:options] block:^BOOL(id obj) {
+        return [obj isKindOfClass:cls];
+    }] autorelease];
+}
+
++ (id)checkWithSelector:(SEL)selector onObject:(id)anObject options:(OCMArgOptions)options
+{
+  return [OCMConstraint constraintWithSelector:selector onObject:anObject options:[self constraintOptionsFromArgOptions:options]];
+}
+
++ (id)checkWithOptions:(OCMArgOptions)options withBlock:(BOOL (^)(id obj))block
+{
+  return [[[OCMBlockConstraint alloc] initWithOptions:[self constraintOptionsFromArgOptions:options] block:block] autorelease];
 }
 
 + (id)unretainedObject:(id)anObject
@@ -152,9 +192,13 @@ + (id)resolveSpecialValues:(NSValue *)value
 	return value;
 }
 
-+ (BOOL)isUnretained:(id)object
++ (OCMConstraintOptions)constraintOptionsFromArgOptions:(OCMArgOptions)argOptions
 {
-  return object_getClass(object) == [OCMUnretainedArgument class];
+  OCMConstraintOptions constraintOptions = 0;
+  if(argOptions & OCMArgDoNotRetainStubArg) constraintOptions |= OCMConstraintDoNotRetainStubArg;
+  if(argOptions & OCMArgDoNotRetainInvocationArg) constraintOptions |= OCMConstraintDoNotRetainInvocationArg;
+  if(argOptions & OCMArgCopyInvocationArg) constraintOptions |= OCMConstraintCopyInvocationArg;
+  return constraintOptions;
 }
 
 @end

Source/OCMock/OCMConstraint.h

@@ -16,8 +16,21 @@
 
 #import <Foundation/Foundation.h>
 
+// See OCMArgOptions for documentation on options.
+typedef NS_OPTIONS(NSUInteger, OCMConstraintOptions) {
+  OCMConstraintDefaultOptions = 0UL,
+  OCMConstraintDoNotRetainStubArg = (1UL << 0),
+  OCMConstraintDoNotRetainInvocationArg = (1UL << 1),
+  OCMConstraintCopyInvocationArg = (1UL << 2),
+  OCMConstraintNeverRetainArg = OCMConstraintDoNotRetainStubArg | OCMConstraintDoNotRetainInvocationArg,
+};
 
-@interface OCMConstraint : NSObject 
+@interface OCMConstraint : NSObject
+
+@property (readonly) OCMConstraintOptions constraintOptions;
+
+- (instancetype)initWithOptions:(OCMConstraintOptions)options NS_DESIGNATED_INITIALIZER;
+- (instancetype)init NS_UNAVAILABLE;
 
 - (BOOL)evaluate:(id)value;
 
@@ -28,6 +41,8 @@
 + (instancetype)constraintWithSelector:(SEL)aSelector onObject:(id)anObject;
 + (instancetype)constraintWithSelector:(SEL)aSelector onObject:(id)anObject withValue:(id)aValue;
 
++ (instancetype)constraintWithSelector:(SEL)aSelector onObject:(id)anObject options:(OCMConstraintOptions)options;
++ (instancetype)constraintWithSelector:(SEL)aSelector onObject:(id)anObject withValue:(id)aValue options:(OCMConstraintOptions)options;
 
 @end
 
@@ -45,8 +60,8 @@
   id testValue;
 }
 
-- (instancetype)initWithTestValue:(id)testValue NS_DESIGNATED_INITIALIZER;
-- (instancetype)init NS_UNAVAILABLE;
+- (instancetype)initWithTestValue:(id)testValue options:(OCMConstraintOptions)options NS_DESIGNATED_INITIALIZER;
+- (instancetype)initWithOptions:(OCMConstraintOptions)options NS_UNAVAILABLE;
 
 @end
 
@@ -61,8 +76,8 @@
   NSInvocation *invocation;
 }
 
-- (instancetype)initWithInvocation:(NSInvocation *)invocation NS_DESIGNATED_INITIALIZER;
-- (instancetype)init NS_UNAVAILABLE;
+- (instancetype)initWithInvocation:(NSInvocation *)invocation options:(OCMConstraintOptions)options NS_DESIGNATED_INITIALIZER;
+- (instancetype)initWithOptions:(OCMConstraintOptions)options NS_UNAVAILABLE;
 
 @end
 
@@ -71,8 +86,8 @@
 	BOOL (^block)(id);
 }
 
-- (instancetype)initWithConstraintBlock:(BOOL (^)(id))block NS_DESIGNATED_INITIALIZER;
-- (instancetype)init NS_UNAVAILABLE;
+- (instancetype)initWithOptions:(OCMConstraintOptions)options block:(BOOL (^)(id))block NS_DESIGNATED_INITIALIZER;
+- (instancetype)initWithOptions:(OCMConstraintOptions)options NS_UNAVAILABLE;
 
 @end