January 9, 2013 10:52 · Jan 9, 2013
diff --git a/MWeakRefCollection.h b/MWeakRefCollection.h
@@ -0,0 +1,18 @@
+#import <Foundation/Foundation.h>
+
+@interface NSObject (MDeallocObserver)
+
+- (void)addDeallocObserverBlock:(void (^)(void))block;
+- (void)addDeallocObserverWithKey:(id<NSCopying>)key block:(void (^)(void))block;
+
+- (void)removeDeallocObserverForKey:(id<NSCopying>)key;
+
+@end
+
+@interface MWeakRefMutableDictionary : NSMutableDictionary
+
+@end
+
+@interface MWeakRefMutableSet : NSMutableSet
+
+@end
diff --git a/MWeakRefCollection.m b/MWeakRefCollection.m
@@ -0,0 +1,148 @@
+#import "MWeakRefCollection.h"
+
+#import <objc/runtime.h>
+
+@interface MDeallocObserver : NSObject
+
+@property (strong, nonatomic) NSMutableDictionary *blocks;
+
+@end
+
+@implementation MDeallocObserver
+
+- (id)init
+{
+    self = [super init];
+    if (self) {
+        _blocks = [NSMutableDictionary dictionary];
+    }
+    return self;
+}
+
+- (void)dealloc
+{
+    for (id block in [_blocks allValues]) {
+        ((void (^)(void))block)();
+    }
+}
+
+@end
+
+static char observerKey;
+
+@implementation NSObject (MDeallocObserver)
+
+- (void)addDeallocObserverBlock:(void (^)(void))block {
+    [self addDeallocObserverWithKey:block block:block];
+}
+
+- (void)addDeallocObserverWithKey:(id<NSCopying>)key block:(void (^)(void))block {
+    MDeallocObserver *observer = objc_getAssociatedObject(self, &observerKey);
+    if (!observer) {
+        observer = [[MDeallocObserver alloc] init];
+        objc_setAssociatedObject(self, &observerKey, observer, OBJC_ASSOCIATION_RETAIN);
+    }
+    [observer.blocks setObject:block forKey:key];
+}
+
+- (void)removeDeallocObserverForKey:(id<NSCopying>)key {
+    MDeallocObserver *observer = objc_getAssociatedObject(self, &observerKey);
+    [observer.blocks removeObjectForKey:key];
+}
+
+@end
+
+#pragma mark -
+
+@implementation MWeakRefMutableDictionary {
+    NSMutableDictionary *_nonretainedDict;
+}
+
+- (id)initWithCapacity:(NSUInteger)numItems {
+    self = [super init];
+    if (self) {
+        const CFDictionaryValueCallBacks valueCallbacks = {0, NULL, NULL, CFCopyDescription, CFEqual};
+        _nonretainedDict = CFBridgingRelease(CFDictionaryCreateMutable(NULL, numItems, &kCFTypeDictionaryKeyCallBacks, &valueCallbacks));
+    }
+    return self;
+}
+
+- (NSUInteger)count {
+    return [_nonretainedDict count];
+}
+
+- (id)objectForKey:(id)aKey {
+    return [_nonretainedDict objectForKey:aKey];
+}
+
+- (NSEnumerator *)keyEnumerator {
+    return [_nonretainedDict keyEnumerator];
+}
+
+- (NSUInteger)countByEnumeratingWithState:(NSFastEnumerationState *)state objects:(__unsafe_unretained id [])buffer count:(NSUInteger)len {
+    return [_nonretainedDict countByEnumeratingWithState:state objects:buffer count:len];
+}
+
+- (void)removeObjectForKey:(id)aKey {
+    id obj = [_nonretainedDict objectForKey:aKey];
+    [obj removeDeallocObserverForKey:aKey];
+
+    [_nonretainedDict removeObjectForKey:aKey];
+}
+
+- (void)setObject:(id)anObject forKey:(id <NSCopying>)aKey {
+    [_nonretainedDict setObject:anObject forKey:aKey];
+    id copiedKey = [(id)aKey copy];
+    __weak NSMutableDictionary *dict = _nonretainedDict;
+    [anObject addDeallocObserverWithKey:aKey block:^{
+        [dict removeObjectForKey:copiedKey];
+    }];
+}
+
+@end
+
+@implementation MWeakRefMutableSet {
+    NSMutableSet *_nonretainedSet;
+}
+
+- (id)initWithCapacity:(NSUInteger)numItems {
+    self = [super init];
+    if (self) {
+        const CFSetCallBacks callbacks = {0, NULL, NULL, CFCopyDescription, CFEqual};
+        _nonretainedSet = (id)CFBridgingRelease(CFSetCreateMutable(NULL, numItems, &callbacks));
+    }
+    return self;
+}
+
+- (NSUInteger)count {
+    return [_nonretainedSet count];
+}
+
+- (id)member:(id)object {
+    return [_nonretainedSet member:object];
+}
+
+- (NSEnumerator *)objectEnumerator {
+    return [_nonretainedSet objectEnumerator];
+}
+
+- (NSUInteger)countByEnumeratingWithState:(NSFastEnumerationState *)state objects:(__unsafe_unretained id [])buffer count:(NSUInteger)len {
+    return [_nonretainedSet countByEnumeratingWithState:state objects:buffer count:len];
+}
+
+- (void)addObject:(id)object {
+    [_nonretainedSet addObject:object];
+    __unsafe_unretained id unretainedobj = object;
+    __weak NSMutableSet *set = _nonretainedSet;
+    [object addDeallocObserverWithKey:[NSValue valueWithNonretainedObject:object] block:^{
+        [set removeObject:unretainedobj];
+    }];
+}
+
+- (void)removeObject:(id)object {
+    [_nonretainedSet removeObject:object];
+
+    [object removeDeallocObserverForKey:[NSValue valueWithNonretainedObject:object]];
+}
+
+@end