DependencyBox.swift

import Foundation

private enum AssociatedKeys {
    static var DeinitializationObserver = "DeinitializationObserver"
}

public class DependencyBox<T: AnyObject> {
    public weak var object: T? {
        didSet {
            guard object !== oldValue else { return }
            if let oldVal = oldValue {
                unregisterDeinitObserver(for: oldVal)
            }
            if let object = object {
                registerDeinitObserver(for: object)
            }
            listenerTokens.allObjects.forEach { $0.onChange(object) }
        }
    }

    private var listenerTokens = NSHashTable<DependencyBoxSubscriptionToken>.weakObjects()

    public init(_ object: T? = nil) {
        self.object = object
        if let object = object {
            registerDeinitObserver(for: object)
        }
    }

    public func subscribe(onChange: @escaping (T?) -> Void) -> DependencyBoxSubscriptionToken {
        let token = DependencyBoxSubscriptionToken(self, onChange: onChange)
        listenerTokens.add(token)
        return token
    }

    private func unsubscribe(_ token: DependencyBoxSubscriptionToken) {
        listenerTokens.remove(token)
    }

    private func registerDeinitObserver(for object: AnyObject) {
        getDeinitObserver(for: object, create: true)?.boxes.add(self)
    }

    private func unregisterDeinitObserver(for object: AnyObject) {
        getDeinitObserver(for: object, create: false)?.boxes.remove(self)
    }

    private func getDeinitObserver(for object: AnyObject, create: Bool) -> DeinitializationObserver? {
        if let observer = objc_getAssociatedObject(object, &AssociatedKeys.DeinitializationObserver) as? DeinitializationObserver {
            return observer
        } else if create {
            let observer = DeinitializationObserver()
            objc_setAssociatedObject(
                object,
                &AssociatedKeys.DeinitializationObserver,
                observer,
                objc_AssociationPolicy.OBJC_ASSOCIATION_RETAIN_NONATOMIC
            )
            return observer
        } else {
            return nil
        }
    }

    private func onObjectDeinit() {
        listenerTokens.allObjects.forEach { $0.onChange(object) }
    }

    private class DeinitializationObserver {
        let boxes = NSHashTable<DependencyBox<T>>.weakObjects()

        deinit {
            boxes.allObjects.forEach { $0.onObjectDeinit() }
        }
    }

    public final class DependencyBoxSubscriptionToken {
        private weak var box: DependencyBox?

        fileprivate let onChange: (T?) -> Void

        init(_ box: DependencyBox, onChange: @escaping (T?) -> Void) {
            self.box = box
            self.onChange = onChange
        }

        public func invalidate() {
            box?.unsubscribe(self)
        }

        deinit {
            self.box?.unsubscribe(self)
        }
    }
}