store.swift

import Foundation

class DisposeToken {
    let block: (Void) -> Void
    init(block: @escaping (Void) -> Void) { self.block = block }
    deinit { block() }
}

class Store<State> {
    
    typealias Callback = (State) -> Void
    
    fileprivate var callbacks = [:] as [Int: Callback]
    fileprivate var nextID = 0
    
    var state: State {
        didSet {
            for (_, callback) in callbacks {
                callback(state)
            }
        }
    }
    
    init(initial state: State) {
        self.state = state
    }
    
    func update(_ block: @escaping (inout (State)) -> Void) {
        var state = self.state
        block(&state)
        self.state = state
    }
    
    func onUpdate(_ callback: @escaping (State) -> Void) -> DisposeToken {
        
        let id = nextID
        nextID += 1
        
        callbacks[id] = callback
        callback(state)
        
        return DisposeToken { [weak self] in self?.callbacks[id] = nil }
    }
}