Avoid retaining objects that are being deallocated.

Source/OCMock/NSInvocation+OCMAdditions.m

@@ -86,7 +86,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))
                 {
@@ -122,7 +122,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
 {