SetAlgebra for classes

ClassSet.swift

//
//  ClassSet.swift
//  audio
//
//  Created by Rob Napier on 1/25/18.
//  Copyright © 2018 Jaybird LLC. All rights reserved.
//

/// Unordered set of unique class objects. Does not require Equatable or Hashable.
struct ClassSet<Element> where Element: AnyObject {
    fileprivate var storage: [ObjectIdentifier: Element] = [:]
    init() {}
}

extension ClassSet: Equatable {
    static func == (lhs: ClassSet<Element>, rhs: ClassSet<Element>) -> Bool {
        return Set(lhs.storage.keys) == Set(rhs.storage.keys)
    }
}

extension ClassSet: SetAlgebra {
    var isEmpty: Bool { return storage.isEmpty }

    func contains(_ member: Element) -> Bool {
        return storage[ObjectIdentifier(member)] != nil
    }

    func union(_ other: ClassSet<Element>) -> ClassSet<Element> {
        var result = self
        result.formUnion(other)
        return result
    }

    func intersection(_ other: ClassSet<Element>) -> ClassSet<Element> {
        var result = self
        result.formIntersection(other)
        return result
    }

    func symmetricDifference(_ other: ClassSet<Element>) -> ClassSet<Element> {
        var result = self
        result.formSymmetricDifference(other)
        return result
    }

    @discardableResult
    mutating func insert(_ newMember: Element) -> (inserted: Bool, memberAfterInsert: Element) {
        let key = ObjectIdentifier(newMember)
        if let existing = storage[key] {
            return (false, existing)
        } else {
            storage[key] = newMember
            return (true, newMember)
        }
    }

    @discardableResult
    mutating func remove(_ member: Element) -> Element? {
        return storage.removeValue(forKey: ObjectIdentifier(member))
    }

    @discardableResult
    mutating func update(with newMember: Element) -> Element? {
        let key = ObjectIdentifier(newMember)
        let oldValue = storage[ObjectIdentifier(newMember)]
        storage[key] = newMember
        return oldValue
    }

    mutating func formUnion(_ other: ClassSet<Element>) {
        storage.merge(other.storage) { (current, _) in current }
    }

    mutating func formIntersection(_ other: ClassSet<Element>) {
        for missingKey in storage.keys where !other.storage.keys.contains(missingKey) {
            storage.removeValue(forKey: missingKey)
        }
    }

    mutating func formSymmetricDifference(_ other: ClassSet<Element>) {
        for member in other {
            if contains(member) {
                remove(member)
            } else {
                insert(member)
            }
        }
    }
}

extension ClassSet: Collection {
    typealias Storage = [ObjectIdentifier: Element]
    typealias Index = Storage.Index
    var startIndex: Index { return storage.startIndex }
    var endIndex: Index { return storage.endIndex }

    subscript(position: Index) -> Element {
        return storage[position].value
    }

    func index(after i: Index) -> Index {
        return storage.index(after: i)
    }
}