experimental example of a reactive µframework with automatic differentiation

AutoDiff.swift

import Foundation

private func get(keyPaths: [AnyKeyPath], from object: Any) -> Any {
  return keyPaths.reduce(object) { acc, kp in acc[keyPath: kp] }
}
extension Collection where Element: Equatable { //might not work for unordered collections
  func hasPrefix(_ prefix: Self) -> Bool {
    guard let firstFromPrefix = prefix.first else { return true }
    guard let firstFromSelf = self.first else { return false }
    return firstFromPrefix == firstFromSelf && dropFirst().hasPrefix(prefix.dropFirst())
  }
}

public final class StateContainer<State> {
  fileprivate var state: State // TODO: id love to somehow get rid of the stored property, as it means state duplication (should be doable with another closure)
  private let onStateChanged: (State, [AnyKeyPath]) -> Void
  fileprivate init(state: State, onStateChanged: @escaping (State, [AnyKeyPath]) -> Void) {
    self.state = state
    self.onStateChanged = onStateChanged
  }

  public subscript<A>(_ keyPath: WritableKeyPath<State, A>) -> A {
    get { return state[keyPath: keyPath] }
    set {
      state[keyPath: keyPath] = newValue
      onStateChanged(state, [keyPath])
    }
  }
  public subscript<B>(_ keyPath: WritableKeyPath<State, B>) -> StateContainer<B> {
    get { return map(keyPath) }
  }
  private func map<B>(_ keyPath: WritableKeyPath<State, B>) -> StateContainer<B> {
    return StateContainer<B>(state: state[keyPath: keyPath]) { newChildState, dirtyChildKeyPaths in
      self.state[keyPath: keyPath] = newChildState
      self.onStateChanged(self.state, [keyPath] + dirtyChildKeyPaths)
    }
  }
}

public final class Store<State, Action> {
  private var stateContainer: StateContainer<State>!
  private let reduce: (StateContainer<State>, Action) -> Void
  public init(state: State, reduce: @escaping (StateContainer<State>, Action) -> Void) {
    self.reduce = reduce
    stateContainer = StateContainer(state: state) { [weak self] _, dirtyKeyPaths in
      self?.dirtyKeyPaths.append(dirtyKeyPaths)
    }
  }
  private var dirtyKeyPaths: [[AnyKeyPath]] = []

  public func dispatch(_ action: Action) {
    reduce(stateContainer, action)
    flush()
  }

  private var observers: [[AnyKeyPath]: [(Any) -> Void]] = [:]
  public func observe<A>(_ keyPath: KeyPath<State, A>, with: @escaping (A) -> Void) {
    observe([keyPath]) { with($0 as! A) }
  }
  // atm both the observe methods and the subscript accessors have to be overloaded for any number of keyPath elements
  // keyPaths longer than 1 element arent directly supported and have to be passed as this tuple
  // adding `KeyPath.hasPrefix()` (+ possibly others) to Swift would probably simplify the implementation as well as make it safer
  public func observe<A,B>(_ keyPaths: (KeyPath<State, A>, KeyPath<A,B>), with: @escaping (B) -> Void) {
    observe([keyPaths.0, keyPaths.1]) { with($0 as! B) }
  }
  private func observe(_ keyPaths: [AnyKeyPath], with: @escaping (Any) -> Void) {
    observers[keyPaths] = (observers[keyPaths] ?? []) + [with]
  }

  private func flush() {
    dirtyKeyPaths.forEach { kps in
      // at this point we want to notify:
      // - observers of the keyPath that changed
      // - observers of any child keyPaths (in case their value also changed)
      let thisAndChildObservers = observers.filter { $0.key.hasPrefix(kps) }
      thisAndChildObservers.forEach { key, observers in
        // TODO: probably notify in order of subscription, currently its undefined order cause of the Dictionary
        // TODO: make sure we dont notify twice in case both parent and child changed in one reduce pass (or multiple synchronous passes if we support Effects)
        // TODO: still somehow diff to make sure it actually changed? would have to be done here (as changing from and back to a value still marks the keyPath as dirty)
        observers.forEach { observer in observer(get(keyPaths: key, from: stateContainer.state)) }
      }
    }
    dirtyKeyPaths = []
  }
}

struct AppState {
  var counter = Counter(value: 0)
  struct Counter {
    var value: Int
  }
}
enum Action {
  case increment
  case changeWholeCounter
}
func reduce(state: StateContainer<AppState>, action: Action) {
  switch action {
  case .increment:
//    state[(\.counter.value)] = state[(\.counter.value)] + 1 //dont use keyPaths longer than 1 for writing else it breaks
    state[\.counter][\.value] = state[(\.counter.value)] + 1 //using longer keyPaths for reading is fine
  case .changeWholeCounter:
    state[(\.counter)] = AppState.Counter(value: 42)
  }
}



let store = Store<AppState,Action>(state: AppState(), reduce: reduce)

//we can get notified of changes at any keyPath (or any of its parents)
store.observe(\.counter) { print("counter changed: \($0)") }
store.observe((\.counter, \.value)) { print("counter value changed: \($0)") }

store.dispatch(.increment)
store.dispatch(.increment)
store.dispatch(.changeWholeCounter)
store.dispatch(.increment)

// its unclear at this point if this is usable in practice
// I feel it might break with collections (+ possibly other State structures)
// On the other hand, I feel like this may not be the only usecase for `KeyPath.hasPrefix()`, considering a Swift Evolution thread
// TODO: try more examples
print("✅")