From d61e9b3209e52de97e7c785a0402e827e0e28db0 Mon Sep 17 00:00:00 2001
From: theraven <theraven@72102866-910b-0410-8b05-ffd578937521>
Date: Mon, 4 Jul 2011 23:27:29 +0000
Subject: [PATCH] Fix __weak references to blocks.

---
 arc.m            | 42 +++++++++++++++++++++++++++++++++---------
 blocks_runtime.m |  8 ++++++++
 2 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/arc.m b/arc.m
index 415ecf2..3c9da92 100644
--- a/arc.m
+++ b/arc.m
@@ -14,7 +14,8 @@
 pthread_key_t ReturnRetained;
 #endif
 
-extern struct objc_class _NSConcreteStackBlock;
+extern void _NSConcreteStackBlock;
+extern void _NSConcreteGlobalBlock;
 
 @interface NSAutoreleasePool
 + (Class)class;
@@ -247,20 +248,26 @@ PRIVATE void init_arc(void)
 #endif
 }
 
+void* block_load_weak(void *block);
+
 id objc_storeWeak(id *addr, id obj)
 {
+	fprintf(stderr, "Storing weak value %p\n", obj);
 	id old = *addr;
 	LOCK_FOR_SCOPE(&weakRefLock);
-	WeakRef *oldRef = weak_ref_table_get(weakRefs, old);
-	while (NULL != oldRef)
+	if (nil != old)
 	{
-		for (int i=0 ; i<4 ; i++)
+		WeakRef *oldRef = weak_ref_table_get(weakRefs, old);
+		while (NULL != oldRef)
 		{
-			if (oldRef->ref[i] == addr)
+			for (int i=0 ; i<4 ; i++)
 			{
-				oldRef->ref[i] = 0;
-				oldRef = 0;
-				break;
+				if (oldRef->ref[i] == addr)
+				{
+					oldRef->ref[i] = 0;
+					oldRef = 0;
+					break;
+				}
 			}
 		}
 	}
@@ -269,7 +276,23 @@ id objc_storeWeak(id *addr, id obj)
 		*addr = obj;
 		return nil;
 	}
-	obj = _objc_weak_load(obj);
+	if (&_NSConcreteGlobalBlock == obj->isa)
+	{
+		// If this is a global block, it's never deallocated, so secretly make
+		// this a strong reference
+		// TODO: We probably also want to do the same for constant strings and
+		// classes.
+		*addr = obj;
+		return obj;
+	}
+	if (&_NSConcreteStackBlock == obj->isa)
+	{
+		block_load_weak(obj);
+	}
+	else
+	{
+		obj = _objc_weak_load(obj);
+	}
 	if (nil != obj)
 	{
 		WeakRef *ref = weak_ref_table_get(weakRefs, obj);
@@ -301,6 +324,7 @@ id objc_storeWeak(id *addr, id obj)
 			weak_ref_insert(weakRefs, newRef);
 		}
 	}
+	*addr = obj;
 	return obj;
 }
 
diff --git a/blocks_runtime.m b/blocks_runtime.m
index 128ae06..9d694af 100644
--- a/blocks_runtime.m
+++ b/blocks_runtime.m
@@ -27,6 +27,7 @@
 #import "objc/runtime.h"
 #import "objc/objc-arc.h"
 #include "gc_ops.h"
+#include "visibility.h"
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
@@ -496,7 +497,14 @@ void _Block_release(void *src)
 		{
 			if(self->flags & BLOCK_HAS_COPY_DISPOSE)
 				self->descriptor->dispose_helper(self);
+			objc_delete_weak_refs((id)self);
 			gc->free(self);
 		}
 	}
 }
+
+PRIVATE void* block_load_weak(void *block)
+{
+	struct block_literal *self = block;
+	return (self->reserved) > 0 ? block : 0;
+}