Avoid dangling weak references to deallocating objects. (#200)

The previous checks for a deallocating object with negative reference count did not work because:
    1. the sign bit has to be recreated as was happening in objc_delete_weak_refs()
    2. there was no distinction between a saturated count and a negative reference count

refcount_max now indicates when the refcount has saturated and should no longer be mutated.
An underflow to -1 still maps to refcount_mask, allowing us to detect when an object is supposed to be deallocated.
Neither objc_release_fast_no_destroy_np() nor objc_retain_fast_np() mutate the refcount when it is one of those values, so the comment in objc_delete_weak_refs() was adjusted.

Test/CMakeLists.txt

@@ -27,6 +27,7 @@ set(TESTS
 	Category.m
 	ExceptionTest.m
 	FastARC.m
+	FastRefCount.m
 	Forward.m
 	ManyManySelectors.m
 	NestedExceptions.m

Test/FastRefCount.m

@@ -0,0 +1,149 @@
+#include "Test.h"
+
+void direct_saturation_test();
+
+int main()
+{
+	id obj = [Test new];
+	assert(object_getRetainCount_np(obj) == 1);
+	
+	for (int i = 0; i < 2; i++)
+	{
+		size_t count = object_getRetainCount_np(obj);
+		id ret = objc_retain_fast_np(obj);
+		assert(ret == obj);
+		assert(object_getRetainCount_np(obj) == ++count);
+	}
+	
+	for (int i = 0; i < 2; i++)
+	{
+		size_t count = object_getRetainCount_np(obj);
+		BOOL destroy = objc_release_fast_no_destroy_np(obj);
+		assert(destroy == NO);
+		assert(object_getRetainCount_np(obj) == --count);
+	}
+	
+	{
+		// Final release should prevent further retains and releases.
+		assert(objc_release_fast_no_destroy_np(obj) == YES);
+		assert(object_getRetainCount_np(obj) == 0);
+		assert(objc_retain_fast_np(obj) == nil);
+		assert(object_getRetainCount_np(obj) == 0);
+		assert(objc_release_fast_no_destroy_np(obj) == NO);
+		assert(object_getRetainCount_np(obj) == 0);
+	}
+	
+	object_dispose(obj);
+	obj = [Test new];
+	
+	{
+		// Should not be able to delete weak refs until final release.
+		id weak;
+		assert(objc_initWeak(&weak, obj) == obj);
+		assert(weak != nil);
+		assert(objc_loadWeakRetained(&weak) == obj);
+		assert(objc_release_fast_no_destroy_np(obj) == NO);
+		// Assumes a return of NO means no effect on obj at all.
+		assert(objc_delete_weak_refs(obj) == NO);
+		assert(objc_loadWeakRetained(&weak) == obj);
+		assert(objc_release_fast_no_destroy_np(obj) == NO);
+		// This will also call objc_delete_weak_refs() and succeed.
+		assert(objc_release_fast_no_destroy_np(obj) == YES);
+		objc_destroyWeak(&weak);
+		// Check what happens when the weak refs were already deleted.
+		assert(objc_delete_weak_refs(obj) == YES);
+	}
+	
+	obj = object_dispose(obj);
+	
+	direct_saturation_test();
+	return 0;
+}
+
+
+// ----------------
+// This test has knowledge of the implementation details of the ARC
+// reference counting and may need modification if the details change.
+
+const long refcount_shift = 1;
+const size_t weak_mask = ((size_t)1)<<((sizeof(size_t)*8)-refcount_shift);
+const size_t refcount_mask = ~weak_mask;
+const size_t refcount_max = refcount_mask - 1;
+
+size_t get_refcount(id obj)
+{
+	size_t *refCount = ((size_t*)obj) - 1;
+	return *refCount & refcount_mask;
+}
+
+void set_refcount(id obj, size_t count)
+{
+	size_t *refCount = ((size_t*)obj) - 1;
+	*refCount = (*refCount & weak_mask) | (count & refcount_mask);
+}
+
+void direct_saturation_test()
+{
+	{
+		id obj = [Test new];
+		// sanity check
+		objc_retain_fast_np(obj);
+		assert(object_getRetainCount_np(obj) == 2);
+		assert(get_refcount(obj) == 1);
+		
+		// Check the behaviour close to the maximum refcount.
+		set_refcount(obj, refcount_max - 3);
+		assert(object_getRetainCount_np(obj) == refcount_max - 2);
+		
+		assert(objc_retain_fast_np(obj) == obj);
+		assert(object_getRetainCount_np(obj) == refcount_max - 1);
+		
+		id weak;
+		assert(objc_initWeak(&weak, obj) == obj);
+		assert(weak != nil);
+		assert(objc_loadWeakRetained(&weak) == obj);
+		assert(object_getRetainCount_np(obj) == refcount_max);
+		
+		// This retain should cause the count to saturate.
+		assert(objc_retain_fast_np(obj) == obj);
+		assert(object_getRetainCount_np(obj) == refcount_max + 1);
+		
+		// A saturated count is no longer affected by retains or releases.
+		assert(objc_release_fast_no_destroy_np(obj) == NO);
+		assert(object_getRetainCount_np(obj) == refcount_max + 1);
+		assert(objc_retain_fast_np(obj) == obj);
+		assert(object_getRetainCount_np(obj) == refcount_max + 1);
+		
+		// Nor should any weak refs be deleted.
+		assert(objc_delete_weak_refs(obj) == NO);
+		assert(objc_loadWeakRetained(&weak) == obj);
+		assert(object_getRetainCount_np(obj) == refcount_max + 1);
+		
+		// Cleanup (can skip this if it becomes an issue)
+		objc_destroyWeak(&weak);
+		set_refcount(obj, 0);
+		objc_release_fast_no_destroy_np(obj);
+		object_dispose(obj);
+	}
+	
+	{
+		id obj = [Test new];
+		set_refcount(obj, refcount_max - 2);
+		assert(objc_retain_fast_np(obj) == obj);
+		assert(objc_retain_fast_np(obj) == obj);
+		assert(object_getRetainCount_np(obj) == refcount_max + 1);
+		
+		// Check we can init a weak ref to an object with a saturated count.
+		id weak;
+		assert(objc_initWeak(&weak, obj) == obj);
+		assert(weak != nil);
+		assert(objc_loadWeakRetained(&weak) == obj);
+		assert(object_getRetainCount_np(obj) == refcount_max + 1);
+		
+		// Cleanup (can skip this if it becomes an issue)
+		objc_destroyWeak(&weak);
+		set_refcount(obj, 0);
+		objc_release_fast_no_destroy_np(obj);
+		object_dispose(obj);
+	}
+}

arc.mm

@@ -238,12 +238,14 @@ static const size_t weak_mask = ((size_t)1)<<((sizeof(size_t)*8)-refcount_shift)
  * All of the bits other than the top bit are the real reference count.
  */
 static const size_t refcount_mask = ~weak_mask;
+static const size_t refcount_max = refcount_mask - 1;
 
 extern "C" OBJC_PUBLIC size_t object_getRetainCount_np(id obj)
 {
 	uintptr_t *refCount = ((uintptr_t*)obj) - 1;
 	uintptr_t refCountVal = __sync_fetch_and_add(refCount, 0);
-	return (((size_t)refCountVal) & refcount_mask) + 1;
+	size_t realCount = refCountVal & refcount_mask;
+	return realCount == refcount_mask ? 0 : realCount + 1;
 }
 
 extern "C" OBJC_PUBLIC id objc_retain_fast_np(id obj)
@@ -268,12 +270,12 @@ extern "C" OBJC_PUBLIC id objc_retain_fast_np(id obj)
 		// If the serialisation happens the other way, then the locked
 		// check of the reference count will happen after we've referenced
 		// this and we don't zero the references or deallocate.
-		if (realCount < 0)
+		if (realCount == refcount_mask)
 		{
 			return nil;
 		}
 		// If the reference count is saturated, don't increment it.
-		if (realCount == refcount_mask)
+		if (realCount == refcount_max)
 		{
 			return obj;
 		}
@@ -329,8 +331,8 @@ extern "C" OBJC_PUBLIC BOOL objc_release_fast_no_destroy_np(id obj)
 	do {
 		refCountVal = newVal;
 		size_t realCount = refCountVal & refcount_mask;
-		// If the reference count is saturated, don't decrement it.
-		if (realCount == refcount_mask)
+		// If the reference count is saturated or deallocating, don't decrement it.
+		if (realCount >= refcount_max)
 		{
 			return NO;
 		}
@@ -341,8 +343,7 @@ extern "C" OBJC_PUBLIC BOOL objc_release_fast_no_destroy_np(id obj)
 		uintptr_t updated = (uintptr_t)realCount;
 		newVal = __sync_val_compare_and_swap(refCount, refCountVal, updated);
 	} while (newVal != refCountVal);
-	// We allow refcounts to run into the negative, but should only
-	// deallocate once.
+	
 	if (shouldFree)
 	{
 		if (isWeak)
@@ -761,7 +762,7 @@ static BOOL setObjectHasWeakRefs(id obj)
 				size_t realCount = refCountVal & refcount_mask;
 				// If this object has already been deallocated (or is in the
 				// process of being deallocated) then don't bother storing it.
-				if (realCount < 0)
+				if (realCount == refcount_mask)
 				{
 					obj = nil;
 					cls = Nil;
@@ -848,16 +849,11 @@ extern "C" OBJC_PUBLIC BOOL objc_delete_weak_refs(id obj)
 	LOCK_FOR_SCOPE(&weakRefLock);
 	if (objc_test_class_flag(classForObject(obj), objc_class_flag_fast_arc))
 	{
-		// If another thread has done a load of a weak reference, then it will
-		// have incremented the reference count with the lock held.  It may
-		// have done so in between this thread's decrementing the reference
-		// count and its acquiring the lock.  In this case, report failure.
+		// Don't proceed if the object isn't deallocating.
 		uintptr_t *refCount = ((uintptr_t*)obj) - 1;
-		// Reconstruct the sign bit.  We don't need to do this on any other 
-		// operations, because even on increment the overflow will be correct
-		// after truncation.
-		uintptr_t refCountVal = (__sync_fetch_and_add(refCount, 0) & refcount_mask) << refcount_shift;
-		if ((((intptr_t)refCountVal) >> refcount_shift) >= 0)
+		uintptr_t refCountVal = __sync_fetch_and_add(refCount, 0);
+		size_t realCount = refCountVal & refcount_mask;
+		if (realCount != refcount_mask)
 		{
 			return NO;
 		}