DisposeCache.swift

import RxSwift

public final class DisposeCache<Key>
  : ExpressibleByDictionaryLiteral where Key : Hashable {
  ///
  private var lock = NSRecursiveLock()

  ///
  private var storage: [Key: Disposable]

  ///
  private var isDisposed = false

  ///
  public var isEmpty: Bool {
    return storage.isEmpty
  }

  ///
  public var count: Int {
    return storage.count
  }

  ///
  public init() {
    storage = [:]
  }

  ///
  public init(dictionaryLiteral elements: (Key, Disposable)...) {
    storage = Dictionary(uniqueKeysWithValues: elements)
  }

  ///
  public init<S>(
    uniqueKeysWithDisposables keysAndDisposables: S
  ) where S : Sequence, S.Element == (Key, Disposable) {
    storage = Dictionary(uniqueKeysWithValues: keysAndDisposables)
  }

  ///
  deinit { dispose() }
}

extension DisposeCache {
  ///
  public subscript (_ key: Key) -> Disposable? {
    get { return storage[key] }
    set { assign(newValue, to: key)?.dispose() }
  }

  ///
  public func contains(_ key: Key) -> Bool {
    return storage[key] != nil
  }

  ///
  public func remove(_ key: Key) {
    self[key] = nil
  }
}

extension DisposeCache {
  ///
  private func assign(_ disposable: Disposable?, to key: Key) -> Disposable? {
    lock.lock()
    defer { lock.unlock() }
    if isDisposed {
      return disposable
    }
    let oldDisposable = storage[key]
    storage[key] = disposable
    return oldDisposable
  }

  ///
  private func dispose() {
    lock.lock()
    defer { lock.unlock() }
    let disposables = storage.map { $0.value }
    storage.removeAll(keepingCapacity: false)
    isDisposed = true
    disposables.forEach { $0.dispose() }
  }
}

extension Disposable {
  ///
  public func disposed<Key>(by cache: DisposeCache<Key>, key: Key) {
    cache[key] = self
  }
}