Multi-dictionaries in Swift

BiMultiDictionary.swift

struct BiMultiDictionary<A: Hashable, B: Hashable> {
    private var keysToValues = MultiDictionary<A, B>()
    private var valuesToKeys = MultiDictionary<B, A>()
    
    func allValues() -> [B] {
        return keysToValues.allValues()
    }
    
    func valuesForKey(key: A) -> [B]? {
        return keysToValues.valueForKey(key)
    }
    
    func keysForValue(value: B) -> [A]? {
        return valuesToKeys.valueForKey(value)
    }
    
    mutating func setKey(key: A, value: B) {
        keysToValues.setKey(key, value: value)
        valuesToKeys.setKey(value, value: key)
    }
    
    mutating func setKey(key: A, values: [B]) {
        keysToValues.setKey(key, values: values)
        
        for value in values {
            valuesToKeys.setKey(value, value: key)
        }
    }
}

DangerousBiMultiDictionary.swift

@objc class DangerousBiMultiDictionary: NSObject {
    private var dictionary = BiMultiDictionary<NSObject, NSObject>()
    
    override init() {}

    func valuesForKey(key: NSObject) -> [NSObject]? {
        return dictionary.valuesForKey(key)
    }
    
    func allValues() -> [NSObject] {
        return dictionary.allValues()
    }
    
    func keysForValue(value: NSObject) -> [NSObject]? {
        return dictionary.keysForValue(value)
    }
    
    func setKey(key: NSObject, value: NSObject) {
        return dictionary.setKey(key, value: value)
    }
    
    func setKey(key: NSObject, values: [NSObject]) {
        dictionary.setKey(key, values: values)
    }
}

MultiDictionary.swift

// TODO: Convert from `Array` to `Set` once using Swift 1.2
struct MultiDictionary<A: Hashable, B: Equatable> {
    private var dictionary = [A: [B]]()

    func valueForKey(key: A) -> [B]? {
        return dictionary[key]
    }
    
    func allValues() -> [B] {
        return reduce(dictionary.values, [], +);
    }
    
    mutating func setKey(key: A, values: [B]) {
        for value in values {
            self.setKey(key, value: value)
        }
    }
    
    mutating func setKey(key: A, value: B) {
        dictionary[key] = {
            if var valuesForKey = self.dictionary[key] {
                if (contains(valuesForKey, value)) {
                    return valuesForKey
                }
                else {
                    return valuesForKey + [value]
                }
            }
            else {
                return [value]
            }
        }()
    }
}