From b32ee7787d8a624e96952bbd8f34259c19e34427 Mon Sep 17 00:00:00 2001
From: Frederik Seiffert <frederik@algoriddim.com>
Date: Sat, 4 Dec 2021 20:51:29 +0100
Subject: [PATCH] Fix weak ref handling during dealloc (#215)

Add test for setting weak ref during dealloc

Fixes #214
---
 Test/ARCTest_arc.m | 16 ++++++++++++++++
 arc.mm             | 12 ++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/Test/ARCTest_arc.m b/Test/ARCTest_arc.m
index af26898..4fd2a65 100644
--- a/Test/ARCTest_arc.m
+++ b/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;
 }
diff --git a/arc.mm b/arc.mm
index 8d23375..0619b2b 100644
--- a/arc.mm
+++ b/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);