TCA @sharedstate

SharedState.swift

import ComposableArchitecture
import SwiftUI

/*

 This is a proof of concept for "shared state" in The Composable Architecture.

 Goals:

 * An ergonomic way for child domains to access data provided by a parent
 * The parent doesn't need to know which children need the data
 * The child doesn't need to know who's providing the data
 * The child is functional in isolation
 * Modifications to shared state only propagate down to children
 * The parent can decide which children receive the data, even sending different data to each child

 Non-Goals:

 * Children can't modify shared state.
 * Child `State` is not mutated outside of a reducer. Doing so would introduce reference semantics that break how TCA is designed.

 It works like this:

 1. Define a `SharedStateKey`
 2. In the parent domain:
    * Use the `@SharedState<Key>` property wrapper in `State` to read and write the value.
    * Wrap child reducers in `WithSharedState` to propagate the value to that subtree of reducers.
 3. In the child domain:
    * Use the `@SharedStateValue<Key>` property wrapper to read the shared value
    * Use the `observeSharedState` higher-order reducer to automatically update state when the value changes.

 */

/// Define a key into the shared value, with a default value.
struct CounterKey: SharedStateKey {
    static var defaultValue: Int = 4
}

struct ParentFeature: ReducerProtocol {
    struct State: Equatable {
        @SharedState<CounterKey> var counter = 10
        var child1 = ChildFeature.State()
        var child2 = ChildFeature.State()
        var child3 = ChildFeature.State()
    }
    enum Action: Equatable {
        case increment
        case child1(ChildFeature.Action)
        case child2(ChildFeature.Action)
        case child3(ChildFeature.Action)
    }
    init() {}
    var body: some ReducerProtocol<State, Action> {
        Reduce { state, action in
            switch action {
            case .increment:
                state.counter += 1
                return .none
            case .child1, .child2, .child3:
                return .none
            }
        }
        // These two children use the value of `counter`
        WithSharedState(\.$counter) {
            Scope(state: \.child1, action: /Action.child1) {
                ChildFeature()
            }
            Scope(state: \.child2, action: /Action.child2) {
                ChildFeature()
            }
        }
        // This child is not affected
        Scope(state: \.child3, action: /Action.child3) {
            ChildFeature()
        }
    }
}

struct ChildFeature: ReducerProtocol {
    struct State: Equatable {
        var sum: Int = 0
        var localCount: Int = 0

        // Putting a `SharedStateValue` in `State` lets us read like other state,
        // and use the `observeSharedState` to respond to changes.
        @SharedStateValue<CounterKey> var sharedCount

        init() {
            // Putting a `SharedStateValue` in `init` lets us grab the
            // current value upon initialization only.
            @SharedStateValue<CounterKey> var counter
            print("ChildFeature.init", counter, self.sharedCount)
        }
    }
    enum Action: Equatable {
        case sharedCount(SharedStateAction<CounterKey>)
        case sum
        case task
    }
    // Putting a `SharedStateValue` in the Reducer lets us treat it like a dependency,
    // accessing it only when needed and not modifying State when changed.
    @SharedStateValue<CounterKey> var counter
    init() {}
    var body: some ReducerProtocol<State, Action> {
        Reduce { state, action in
            switch action {
            case .sharedCount(.willChange(let newValue)):
                print("ChildFeature.willChange", state.sharedCount, "=>", newValue)
                return .none
            case .sum:
                state.sum = state.localCount + state.sharedCount
                return .none
            case .task:
                state.localCount = .random(in: 0..<10)
                return .none
            }
        }
        .observeSharedState(\.$sharedCount, action: /Action.sharedCount)
    }
}

struct ParentView: View {
    let store: StoreOf<ParentFeature>
    var body: some View {
        List {
            WithViewStore(store) { viewStore in
                HStack {
                    Button(action: { viewStore.send(.increment) }) {
                        Text("Increment")
                    }
                    Spacer()
                    Text(viewStore.counter.formatted())
                }
            }
            Section {
                ChildView(store: store.scope(state: \.child1, action: ParentFeature.Action.child1))
            }
            Section {
                ChildView(store: store.scope(state: \.child2, action: ParentFeature.Action.child2))
            }
            Section {
                ChildView(store: store.scope(state: \.child3, action: ParentFeature.Action.child3))
            }
        }
    }
}

struct ChildView: View {
    let store: StoreOf<ChildFeature>
    var body: some View {
        WithViewStore(store) { viewStore in
            HStack {
                Text("Local Count")
                Spacer()
                Text(viewStore.localCount.formatted())
            }
            HStack {
                Text("Shared Count")
                Spacer()
                Text(viewStore.sharedCount.formatted())
            }
            HStack {
                Button(action: { viewStore.send(.sum) }) {
                    Text("Sum Counts")
                }
                Spacer()
                Text(viewStore.sum.formatted())
            }
        }
        .task { await ViewStore(store.stateless).send(.task).finish() }
    }
}

struct Parent_Previews: PreviewProvider {
    static var previews: some View {
        ParentView(
            store: Store(initialState: ParentFeature.State()) {
                ParentFeature()
            } withDependencies: {
                // This default value will be used where a parent doesn't provide one.
                $0.sharedState(CounterKey.self, 10)
            }
        )
    }
}


// ========================================================

public protocol SharedStateKey: Sendable, Equatable {
    associatedtype Value: Sendable
    static var defaultValue: Value { get }
}

/// A property wrapper that can share its value.
@propertyWrapper
public struct SharedState<Key: SharedStateKey> where Key.Value: Equatable {
    public var wrappedValue: Key.Value
    public var projectedValue: Self { self }
    public init(wrappedValue: Key.Value) {
        self.wrappedValue = wrappedValue
    }
}

extension SharedState: Equatable where Key.Value: Equatable {}
extension SharedState: Sendable where Key.Value: Sendable {}

/// A reducer that propagates shared state to its child reducer.
///
/// The shared state is only available to children of this reducer. Other reducers accessing
/// the same `SharedValueKey` key are unaffected.
public struct WithSharedState<Key: SharedStateKey, ParentState, ParentAction, Child: ReducerProtocol>: ReducerProtocol
where Key.Value: Equatable, ParentState == Child.State, ParentAction == Child.Action
{
    public init(
        _ toSharedState: KeyPath<ParentState, SharedState<Key>>,
        file: StaticString = #fileID,
        line: UInt8 = #line,
        @ReducerBuilder<Child.State, Child.Action> base: () -> Child
    ) {
        self.id = Identifier(file: "\(file)", line: line)
        self.toSharedState = toSharedState
        self.base = base()
    }
    private struct Identifier: Hashable {
        let file: String
        let line: UInt8
    }
    private let id: Identifier
    private let toSharedState: KeyPath<Child.State, SharedState<Key>>
    private let base: Child
    public func reduce(into state: inout Child.State, action: Child.Action) -> EffectTask<Child.Action> {
        let value = state[keyPath: self.toSharedState].wrappedValue
        return self.base
            .transformDependency(\._sharedValues) {
                $0 = $0.child(
                    id: self.id,
                    key: Key.self,
                    value: value
                )
            }
            .reduce(into: &state, action: action)
    }
}

/// A property wrapper that reads a value from shared state.
@propertyWrapper
public struct SharedStateValue<Key: SharedStateKey> where Key.Value: Equatable {
    fileprivate var isObserving: Bool = false
    fileprivate var _wrappedValue: Key.Value
    public fileprivate(set) var projectedValue: Self {
        get { self }
        set { self = newValue }
    }
    public init() {
        @Dependency(\._sharedValues) var sharedValues
        self._wrappedValue = sharedValues[Key.self]
    }
    public var wrappedValue: Key.Value {
        self._wrappedValue
    }
}

/// Actions sent when shared state changes.
enum SharedStateAction<Key: SharedStateKey> {
    /// Received by the reducer just before the value changes. You
    /// may compare the value in `State` to this value.
    case willChange(Key.Value)
}

extension SharedStateValue: Equatable where Key.Value: Equatable {}
extension SharedStateValue: Sendable where Key.Value: Sendable {}

extension SharedStateAction: Equatable where Key.Value: Equatable {}
extension SharedStateAction: Sendable where Key.Value: Sendable {}

extension ReducerProtocol {
    /// A higher-order reducer that monitors shared state for changes and sends an action
    /// back into the system to update state with the current value.
    func observeSharedState<Key: SharedStateKey>(
        _ toSharedState: WritableKeyPath<State, SharedStateValue<Key>>,
        action toSharedAction: CasePath<Action, SharedStateAction<Key>>
    ) -> some ReducerProtocol<State, Action>
    where Key.Value: Equatable
    {
        _SharedStateObserver(
            toSharedState: toSharedState,
            toSharedAction: toSharedAction,
            base: self
        )
    }
}

struct _SharedStateObserver<Key: SharedStateKey, ParentState, ParentAction, Base: ReducerProtocol>: ReducerProtocol
where Key.Value: Equatable, ParentState == Base.State, ParentAction == Base.Action {
    let toSharedState: WritableKeyPath<ParentState, SharedStateValue<Key>>
    let toSharedAction: CasePath<ParentAction, SharedStateAction<Key>>
    let base: Base
    @Dependency(\._sharedValues) var sharedValues
    func reduce(into state: inout ParentState, action: ParentAction) -> EffectTask<ParentAction> {
        let effects = self.base.reduce(into: &state, action: action)
        switch self.toSharedAction.extract(from: action) {
        case .willChange(let value):
            state[keyPath: toSharedState]._wrappedValue = value
        case .none:
            break
        }
        guard
            !state[keyPath: self.toSharedState].isObserving
        else {
            return effects
        }
        state[keyPath: self.toSharedState].isObserving = true
        let initialValue = state[keyPath: self.toSharedState].wrappedValue
        return .merge(
            effects,
            .run { send in
                var firstValue = true
                for await value in self.sharedValues.observe(Key.self) {
                    if !firstValue || (firstValue && value != initialValue) {
                        await send(self.toSharedAction.embed(.willChange(value)))
                        firstValue = false
                    }
                }
            }
        )
    }
}

extension DependencyValues {
    /// Set the initial value for shared state. This is intended to be used for testing or previews.
    mutating func sharedState<Key: SharedStateKey>(_ key: Key.Type, _ value: Key.Value, file: StaticString = #fileID, line: UInt8 = #line) {
        self._sharedValues = self._sharedValues.child(
            id: Identifier(file: "\(file)", line: line),
            key: key,
            value: value
        )
    }
    private struct Identifier: Hashable {
        let file: String
        let line: UInt8
    }
}

import Combine

private var _sharedValueChildren = [ AnyHashable : _SharedValues ]()

struct _SharedValues: @unchecked Sendable {

    init(id: AnyHashable, values: Storage = [:]) {
        self.id = id
        self.storage = CurrentValueSubject(values)
    }

    typealias Storage = [ ObjectIdentifier : Any ]
    private let id: AnyHashable
    private var storage: CurrentValueSubject<Storage, Never>

    /// Create a child copy with new value.
    func child<Key: SharedStateKey>(id: AnyHashable, key: Key.Type, value: Key.Value) -> _SharedValues {
        var values = self.storage.value
        values[ObjectIdentifier(key)] = value
        if let child = _sharedValueChildren[id] {
            child.storage.value = values
            return child
        } else {
            let child = _SharedValues(id: id, values: values)
            _sharedValueChildren[id] = child
            return child
        }
    }

    /// Read and write to a shared value key.
    subscript<Key: SharedStateKey>(_ key: Key.Type) -> Key.Value {
        get {
            guard
                let value = self.storage.value[ObjectIdentifier(key)],
                let value = value as? Key.Value
            else {
                return key.defaultValue
            }
            return value
        }
        set {
            self.storage.value[ObjectIdentifier(key)] = newValue
        }
    }

    /// Observe changes to a shared value key.
    func observe<Key: SharedStateKey>(_ key: Key.Type) -> AsyncStream<Key.Value> where Key.Value: Equatable {
        AsyncStream(
            self.storage
                .map { storage in
                    guard
                        let value = storage[ObjectIdentifier(key)],
                        let value = value as? Key.Value
                    else {
                        return key.defaultValue
                    }
                    return value
                }
                .removeDuplicates()
                .values
        )
    }
}

extension _SharedValues: DependencyKey {
    static let liveValue = _SharedValues(id: "root")
}

extension DependencyValues {
    var _sharedValues: _SharedValues {
        get { self[_SharedValues.self] }
        set { self[_SharedValues.self] = newValue }
    }
}