Weak reference version of NSMutableArray and NSMutableSet. Use associated object to watch dealloc of objects and remove them when deallocated.

MWeakRefCollection.h

#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

MWeakRefCollection.m

#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