Generic NSMapTable replacement written in Swift

Question

This is my attempt at writing a generic NSMapTable with weak keys and strong values (after your feedback I'll be trying to write Strong-Key/Weak-Value and Weak-Key/Weak-Value variants to have the complete functionality of NSMapTable with generics). To my knowledge it works correctly.

public class WeakKeyDictionary<K: AnyObject, V where K: Hashable> {

    private var dict = Dictionary<HashableWeakBox<K>, V>()
    public var block: (V)->() = { _ in }

    public init() {}
    public init(dictionary: Dictionary<K, V>) {
        for (k, v) in dictionary {
            setValue(v, forKey: k)
        }
    }

    public subscript(key: K) -> V? {
        get { return valueForKey(key) }
        set { setValue(newValue, forKey: key) }
    }

    public func valueForKey(key: K) -> V? {
        return dict[HashableWeakBox(key)]
    }
    public func setValue(newValue: V?, forKey key: K) {
        let hashableBox = HashableWeakBox(key)

        if let value = newValue {
            let watcher = DeallocWatcher { [weak self] in
                if let v = self?.dict[hashableBox] { self?.block(v) }
                self?.dict[hashableBox] = nil
            }

            objc_setAssociatedObject(key, unsafeAddressOf(self), watcher, objc_AssociationPolicy(OBJC_ASSOCIATION_RETAIN_NONATOMIC))
            dict[hashableBox] = value
        }
        else {
            objc_setAssociatedObject(key, unsafeAddressOf(self), nil, 0)
            dict[hashableBox] = nil
        }
    }
    public func removeValueForKey(key: K) -> V? {
        objc_setAssociatedObject(key, unsafeAddressOf(self), nil, 0)
        return dict.removeValueForKey(HashableWeakBox(key))
    }

    public var count: Int { return dict.count }
    public var isEmpty: Bool { return dict.isEmpty }
    public var keyValues: [(K, V)] {
        var list = [(K, V)]()
        for (k, v) in dict { list.append((k.value!, v)) }
        return list
    }
    public var keys: [K] {
        var list = [K]()
        for k in dict.keys { list.append(k.value!) }
        return list
    }
    public var values: [V] {
        return Array(dict.values)
    }

    deinit {
        for box in dict.keys {
            objc_setAssociatedObject(box.value, unsafeAddressOf(self), nil, 0)
        }
    }
}

private func == <T: Hashable>(l: T, r: T) -> Bool { return l.hashValue == r.hashValue }

private class HashableWeakBox<T: AnyObject where T: Hashable>: Hashable {
    weak var value: T?
    let hashValueWhenNil: Int

    init(_ v: T) {
        value = v
        hashValueWhenNil = v.hashValue
    }

    var hashValue: Int { return value?.hashValue ?? hashValueWhenNil }
}

private class DeallocWatcher {
    let callback: ()->()
    init(_ c: ()->()) { callback = c }
    deinit { callback() }
}

The trick needed in order to be notified when one of the keys disappears is to add a watcher object (DeallocWatcher) as an associated object of the key (this only works on Swift 1.2, if I'm not mistaken). That way, whenever the key disappears, so does the watcher object, which in turn warns the dictionary class when deiniting.

The block property should, if needed, be set by the code using the class, and provides a way to be notified whenever a key is removed from the dictionary.

I'd like to know if this is reliable and works correctly in all cases. I'd also like to have a better idea of under which conditions can this be used across multiple threads. Finally, I'd like feedback on how I can better follow Swift guidelines and implement common "Swift-isms".

Answer 1

(Not a full review, but some remarks and suggestions.)

Generally, your code looks correct to me, I could not see obvious errors or problems – for the single-threaded usage.

You cannot assume that it will work correctly when used from multiple threads simultaneously. Cocoa mutable collection types are not thread-safe (see Thread Safety Summary), and – as far as I know – the same applies to Swift collections (references: Swift: process Array in parallel using GCD, Swift mutable array thread safety, Swift Arrays are not Threadsafe).

---

In

let watcher = DeallocWatcher { [weak self] in
    if let v = self?.dict[hashableBox] { self?.block(v) }
    self?.dict[hashableBox] = nil
}

self is (conditionally) unwrapped three times. I would use optional binding instead so that it is unwrapped only once. You could also take advantage of the fact that removeValueForKey() returns the old value (or nil):

let watcher = DeallocWatcher { [weak self] in
    if let s = self {
        if let v = s.dict.removeValueForKey(hashableBox) {
            s.block(v)
        }
    }
}

---

When a value for key is removed in

    else {
        objc_setAssociatedObject(key, unsafeAddressOf(self), nil, 0)
        dict[hashableBox] = nil // <-- not needed!
    }

then the first line already causes the value to be removed from dict, because disassociating the associated object fires the DeallocWatcher callback. So the second line is not needed.

---

The keys and keyValues properties can be simplified using map():

public var keyValues: [(K, V)] {
    return map(dict) { (k, v) -> (K, V) in (k.value!, v) }
}

public var keys: [K] {
    return Array(dict.keys.map { $0.value! })
}

---

You can also easily add conformance to additional protocols to make your WeakKeyDictionary more look and feel like a real dictionary.

If you add the Printable protocol

public class WeakKeyDictionary<K: AnyObject, V where K: Hashable> : Printable {

    // ...
    public var description: String  {
        return "[ " + ", ".join(map(dict) { (k, v) in "\(k.value!) : \(v) " }) + "]"
    }
}

then println(wkdict) produces an output similar to a dictionary.

And with SequenceType

public class WeakKeyDictionary<K: AnyObject, V where K: Hashable> : Printable, SequenceType {

    // ...
    public func generate() -> GeneratorOf<(K, V)> {
        var gen = dict.generate()
        return GeneratorOf {
            if let (k, v) = gen.next() {
                return (k.value!, v)
            }
            return nil
        }
    }
}

you can enumerate keys and values like

for (k, v) in wkdict {
    println((k, v))
}

Answer 2

There's one thing I missed that would make this class more "Swift-like": implementing DictionaryLiteralConvertible. If WeakKeyDictionary conforms to that protocol, this line:

let weakDict = WeakKeyDictionary(dictionary: [key1: v1, key2: v2, key3: v3])

can be written like this:

let weakDict: WeakKeyDictionary = [key1: v1, key2: v2, key3: v3]

which is much cleaner and more in line with the other Swift types.

That protocol is very easy to implement:

public convenience required init(dictionaryLiteral elements: (K, V)...) {
    var dict = [K:V]()
    for (k, v) in elements { dict[k] = v }
    self.init(dictionary: dict)
}