WeakSet

WeakSet.swift

import Foundation

class WeakSet<T> {
    
    var count: Int {
        return weakStorage.count
    }
    
    private let weakStorage = NSHashTable<AnyObject>.weakObjects()
    
    func add(_ object: T) {
        weakStorage.add(object as AnyObject)
    }
    
    func remove(_ object: T) {
        weakStorage.remove(object as AnyObject)
    }
    
    func removeAllObjects() {
        weakStorage.removeAllObjects()
    }
    
    var allObjects: [T] {
        let enumerator = weakStorage.objectEnumerator()
        let iterator = AnyIterator {
            return enumerator.nextObject() as! T?
        }
        return Array(iterator)
    }
}