Fix weak ref handling during dealloc (#215)

Add test for setting weak ref during dealloc

Fixes #214

Test/ARCTest_arc.m

@@ -60,6 +60,17 @@ id __weak var;
 }
 @end
 
+static __weak id weakRef;
+
+@interface CheckDeallocWeakRef : Test
+@end
+@implementation CheckDeallocWeakRef
+- (void)dealloc
+{
+	weakRef = self;
+}
+@end
+
 
 int main(void)
 {
@@ -115,5 +126,10 @@ int main(void)
 		[obj setWeak: x];
 	}
 	assert([obj loadWeak] != nil);
+	// Check setting weak references during dealloc
+	{
+		[CheckDeallocWeakRef new];
+	}
+	assert(weakRef == nil);
 	return 0;
 }

arc.mm

@@ -838,6 +838,12 @@ extern "C" OBJC_PUBLIC id objc_storeWeak(id *addr, id obj)
 		*addr = obj;
 		return obj;
 	}
+	// If the object is being deallocated return nil.
+	if (object_getRetainCount_np(obj) == 0)
+	{
+		*addr = nil;
+		return nil;
+	}
 	if (nil != obj)
 	{
 		*addr = (id)incrementWeakRefCount(obj);
@@ -991,6 +997,12 @@ extern "C" OBJC_PUBLIC id objc_initWeak(id *addr, id obj)
 		*addr = obj;
 		return obj;
 	}
+	// If the object is being deallocated return nil.
+	if (object_getRetainCount_np(obj) == 0)
+	{
+		*addr = nil;
+		return nil;
+	}
 	if (nil != obj)
 	{
 		*(WeakRef**)addr = incrementWeakRefCount(obj);