Objective-C Dictionary with Zeroing Weak References

gistfile1.m

#import <Foundation/Foundation.h>

@protocol DeallocListener <NSObject>

- (void)objectDidDeallocateForKey:(id)key object:(id)obj;

@end

@interface WeakDictionary : NSMutableDictionary <DeallocListener>

@end

#import <objc/runtime.h>
#import <objc/message.h>

@interface __ObjectPair : NSObject

@property (readonly) id first, second;

@end

@implementation __ObjectPair

- (id)initWithObject:(id)first andObject:(id)second
{
  if (!(self = [super init]))
		return nil;
	
	_first = [first retain];
	_second = [second retain];
	
	return self;
}

- (void)dealloc
{
	[_first release];
	[_second release];
	
	[super dealloc];
}

- (BOOL)isEqual:(id)object
{
	return [self.first isEqual:object];
}

@end

void proxy_addDeallocListener_forKey(id self, SEL _cmd, id<DeallocListener, NSCopying> listener, id key) {
	NSMutableArray * arr = objc_getAssociatedObject(self, class_getName([self class]));
	if (!arr)
		objc_setAssociatedObject(self, class_getName([self class]), arr = [[NSMutableArray alloc] init], OBJC_ASSOCIATION_ASSIGN);
	
	[arr addObject:[[[__ObjectPair alloc] initWithObject:listener andObject:key] autorelease]];
}

void proxy_removeDeallocListener(id self, SEL _cmd, id<DeallocListener> listener) {
	NSMutableArray * arr = objc_getAssociatedObject(self, class_getName([self class]));
	[arr removeObject:listener];
}

void proxy_dealloc(id self, SEL _cmd) {
	NSMutableArray * arr = objc_getAssociatedObject(self, class_getName([self class]));
	for (__ObjectPair * pair in arr)
		[pair.first objectDidDeallocateForKey:pair.second object:self];
	[arr release];
	
	// [super dealloc]
	objc_msgSendSuper(&((struct objc_super){self, [self superclass]}), _cmd);
}

@protocol DeallocNotifier <NSObject>

- (void)addDeallocListener:(id<DeallocListener, NSCopying>)listener forKey:(id)key;
- (void)removeDeallocListener:(id<DeallocListener, NSCopying>)listener;

@end

@implementation WeakDictionary {
	NSMutableDictionary * _backing;
}

- (void)dealloc
{
	[_backing release];
	
	[super dealloc];
}

- (id)initWithObjects:(const id [])objects forKeys:(const id<NSCopying> [])keys count:(NSUInteger)cnt
{
	if (!(self = [super init]))
		return nil;
	
	_backing = [[NSMutableDictionary alloc] initWithObjects:objects forKeys:keys count:cnt];
	
	return self;
}

- (NSUInteger)count
{
	return _backing.count;
}

- (id)objectForKey:(id)aKey
{
	return [_backing objectForKey:aKey];
}

- (NSEnumerator *)keyEnumerator
{
	return [_backing keyEnumerator];
}

- (void)setObject:(id)anObject forKey:(id<NSCopying>)aKey
{
	[self removeObjectForKey:aKey];
	
	if (!anObject)
		return;
	
	Class original = [anObject class];
	NSString * newName = [NSStringFromClass(original) stringByAppendingString:@"__WeakDictionary__DeallocNotifier"];
	
	Class newClass = NSClassFromString(newName);
	if (!newClass) {
		newClass = objc_allocateClassPair(original, [newName cStringUsingEncoding:NSASCIIStringEncoding], 0);
		class_addMethod(newClass, @selector(addDeallocListener:forKey:), (IMP)&proxy_addDeallocListener_forKey, "v@:@@");
		class_addMethod(newClass, @selector(removeDeallocListener:), (IMP)&proxy_removeDeallocListener, "v@:@");
		class_addMethod(newClass, @selector(dealloc), (IMP)&proxy_dealloc, "v@:");
	}
	
	object_setClass(anObject, newClass);
	id<DeallocNotifier> newObject = anObject;
	[newObject addDeallocListener:self forKey:aKey];
	[_backing setObject:[NSValue valueWithNonretainedObject:anObject] forKey:aKey];
}

- (void)removeObjectForKey:(id)aKey
{
	NSValue * value = _backing[aKey];
	id<DeallocNotifier> object = value.nonretainedObjectValue;
	if (!object)
		return;
	
	[_backing removeObjectForKey:aKey];
	[object removeDeallocListener:self];
	object_setClass(object, [object superclass]);
}

- (void)objectDidDeallocateForKey:(id)key object:(id)obj
{
	if ([_backing[key] nonretainedObjectValue] == obj)
		[_backing removeObjectForKey:key];
}

@end