Avoid retaining objects that are being deallocated.

We do not want to retain objects that are being deallocated because this will cause
a crash later when the retaining object releases them.

An example of this is when a delegate (that is mocked) is called in a dealloc for an
object. The mock retains the deallocating object as part of an invocation and then
things crash later when the invocation is released as part of the mock deallocating
itself.

Author: dmaclach

Source/OCMock/NSInvocation+OCMAdditions.m

@@ -66,7 +66,7 @@ - (void)retainObjectArgumentsExcludingObject:(id)objectToExclude
     NSMutableArray *retainedArguments = [[NSMutableArray alloc] init];
 
     id target = [self target];
-    if((target != nil) && (target != objectToExclude) && !object_isClass(target))
+    if((target != nil) && (target != objectToExclude) && !object_isClass(target) && !OCMIsDeallocating(target))
     {
         // Bad things will happen if the target is a block since it's not being
         // copied. There isn't a very good way to tell if an invocation's target
@@ -84,7 +84,7 @@ - (void)retainObjectArgumentsExcludingObject:(id)objectToExclude
         {
             id argument;
             [self getArgument:&argument atIndex:index];
-            if((argument != nil) && (argument != objectToExclude))
+            if((argument != nil) && (argument != objectToExclude) && !OCMIsDeallocating(argument))
             {
                 if(OCMIsBlockType(argumentType))
                 {
@@ -120,7 +120,7 @@ - (void)retainObjectArgumentsExcludingObject:(id)objectToExclude
     {
         id returnValue;
         [self getReturnValue:&returnValue];
-        if((returnValue != nil) && (returnValue != objectToExclude))
+        if((returnValue != nil) && (returnValue != objectToExclude) && !OCMIsDeallocating(returnValue))
         {
             if(OCMIsBlockType(returnType))
             {

Source/OCMock/OCMFunctions.m

@@ -32,6 +32,14 @@ - (void)recordFailureWithDescription:(NSString *)description inFile:(NSString *)
 - (void)failWithException:(NSException *)exception;
 @end
 
+@interface NSObject(OCMKnownNSObjectMethods)
+- (BOOL)_isDeallocating;
+@end
+
+// From objc/runtime/objc-internal.h
+// Only available on macOS 10.11/iOS 9.
+extern id objc_initWeakOrNil(id *location, id newObj) __attribute__((weak_import));
+extern void objc_destroyWeak(id _Nullable * _Nonnull location) __attribute__((weak_import));
 
 #pragma mark  Functions related to ObjC type system
 
@@ -466,3 +474,9 @@ void OCMReportFailure(OCMLocation *loc, NSString *description)
     }
 
 }
+
+BOOL OCMIsDeallocating(id anObject)
+{
+    return [anObject _isDeallocating];
+}
+

Source/OCMock/OCMFunctionsPrivate.h

@@ -48,6 +48,8 @@ OCPartialMockObject *OCMGetAssociatedMockForObject(id anObject);
 
 void OCMReportFailure(OCMLocation *loc, NSString *description);
 
+BOOL OCMIsDeallocating(id anObject);
+
 BOOL OCMIsNonEscapingBlock(id block);
 
 

Source/OCMock/OCMStubRecorder.m

@@ -24,6 +24,7 @@
 #import "OCMBlockCaller.h"
 #import "OCMRealObjectForwarder.h"
 #import "OCMInvocationStub.h"
+#import "OCMFunctionsPrivate.h"
 
 
 @implementation OCMStubRecorder
@@ -51,7 +52,7 @@ - (OCMInvocationStub *)stub
 - (id)andReturn:(id)anObject
 {
     id action;
-    if(anObject == mockObject)
+    if(anObject == mockObject || OCMIsDeallocating(anObject))
     {
         action = [[[OCMNonRetainingObjectReturnValueProvider alloc] initWithValue:anObject] autorelease];
     }

Source/OCMockTests/OCMockObjectRuntimeTests.m

@@ -54,14 +54,15 @@ - (NSString *)stringForTypedef:(TypedefString *)string
 
 @interface TestDelegate : NSObject
 
-- (void)go;
+- (id)go:(id)sender;
 
 @end
 
 @implementation TestDelegate
 
-- (void)go
+- (id)go:(id)sender
 {
+    return sender;
 }
 
 @end
@@ -77,7 +78,24 @@ @implementation TestClassWithDelegate
 - (void)run
 {
     TestDelegate *delegate = self.delegate;
-    [delegate go];
+    [delegate go:nil];
+}
+
+@end
+
+
+@interface TestClassThatCallsDelegateOnDealloc : NSObject
+
+@property (nonatomic, weak) TestDelegate *delegate;
+
+@end
+
+@implementation TestClassThatCallsDelegateOnDealloc
+
+- (void)dealloc
+{
+    TestDelegate *delegate = self.delegate;
+    [delegate go:self];
 }
 
 @end
@@ -313,7 +331,7 @@ - (void)testWeakReferencesShouldStayAround
 
     [object run];
 
-    OCMVerify([mockDelegate go]);
+    OCMVerify([mockDelegate go:nil]);
     XCTAssertNotNil(object.delegate, @"Should still have delegate");
 }
 
@@ -329,6 +347,16 @@ - (void)testDynamicSubclassesShouldBeDisposed
     XCTAssertEqual(numClassesBefore, numClassesAfter, @"Should have disposed dynamically generated classes.");
 }
 
+- (void)testHandlesCallingMockWithSelfAsArgumentInDealloc
+{
+    // Note that this test will crash on failure.
+    id mock = [OCMockObject mockForClass:[TestDelegate class]];
+    [[mock expect] go:OCMOCK_ANY];
+    TestClassThatCallsDelegateOnDealloc *foo = [[TestClassThatCallsDelegateOnDealloc alloc] init];
+    foo.delegate = mock;
+    foo = nil;
+    [mock verify];
+}
 
 - (void)testClassesWithResolveMethodsCanBeMocked
 {