Identified State

[arnauddorgans]

Across many projects I built with TCA in recent years, I have experienced the same need that is lacking for a solution in the current TCA library

Let's say that you have an app with real-time features, and in an optional child feature, you have to subscribe to live events for a specified id.

Today, what we can do to properly have a child feature identified by a key and have all its running effects be cancelled properly when the child state ID changes, is to use an IdentifiedArray even if the array will ever contain one or zero items.

What I would like to have is either an IdentifiedState (propertyWrapper) or to have something that do it automaticaly if the child feature state is Identifiable.

This is the unit tests I expect to pass

@MainActor
@Suite
struct IfLetReducerTests {
    @Test func scope() async {
        let store = Store(initialState: .init(), reducer: { Feature() })
        let storeNoChild: StoreOf<ChildFeature>? = store.scope(state: \.child, action: \.child)
        #expect(storeNoChild == nil)
        store.send(.setChild(.init(id: 1)))
        let storeChild: StoreOf<ChildFeature>? = store.scope(state: \.child, action: \.child)
        #expect(storeChild != nil)
        store.send(.setChild(.init(id: 1)))
        let storeChildUpdated: StoreOf<ChildFeature>? = store.scope(state: \.child, action: \.child)
        #expect(storeChild === storeChildUpdated)
        store.send(.setChild(.init(id: 2)))
        let storeChild2: StoreOf<ChildFeature>? = store.scope(state: \.child, action: \.child)
        #expect(storeChild !== storeChild2)
    }

    @Test func reducer() async {
        var runningCount = 0
        var isRunning: [Int: Bool] = [:]
        let store = TestStore(initialState: .init(), reducer: { Feature() }, withDependencies: {
            $0[FakeAPI.self] = .init(fetch: { @MainActor in
                runningCount += 1
                let id = runningCount
                isRunning[id] = true
                try? await Task.never()
                isRunning[id] = false
            })
        })
        store.exhaustivity = .off
        await store.send(.setChild(.init(id: 1)))
        #expect(isRunning[1] == true)
        #expect(runningCount == 1)
        await store.send(.setChild(nil))
        #expect(isRunning[1] == false)
        #expect(runningCount == 1)
        await store.send(.setChild(.init(id: 1)))
        #expect(isRunning[2] == true)
        #expect(runningCount == 2)
        await store.send(.setChild(.init(id: 1)))
        #expect(isRunning[2] == true)
        #expect(isRunning[3] == true)
        #expect(runningCount == 3)
        await store.send(.setChild(.init(id: 2)))
        #expect(isRunning[2] == false)
        #expect(isRunning[3] == false)
        #expect(runningCount == 4)
    }
}

@Reducer
private struct Feature {
    @ObservableState
    struct State: Equatable {
        var child: ChildFeature.State?
    }

    enum Action {
        case setChild(ChildFeature.State?)
        case child(ChildFeature.Action)
    }

    var body: some ReducerOf<Self> {
        Reduce { state, action in
            switch action {
            case let .setChild(value):
                state.child = value
                if value != nil {
                    return reduce(into: &state, action: .child(.lifeTask))
                }
                return .none
            case .child:
                return .none
            }
        }
        .ifLet(\.child, action: \.child) {
            ChildFeature()
        }
    }
}

@Reducer
private struct ChildFeature {
    @ObservableState
    struct State: Equatable, Identifiable {
        var id: Int
    }

    enum Action {
        case lifeTask
    }

    @Dependency(FakeAPI.self) private var api

    var body: some ReducerOf<Self> {
        Reduce { state, action in
            switch action {
            case .lifeTask:
                return .run { _ in
                    try await api.fetch()
                }
            }
        }
    }
}

private struct FakeAPI: TestDependencyKey {
    var fetch: @Sendable () async throws -> Void

    static let testValue: FakeAPI = {
        return .init(fetch: {
            try await Task.never()
        })
    }()
}

And here is what I did in my project to make the tests pass

public extension Reducer {
    /// ifLet but keyed by ID, clearing the child store instance whenever the ID changes.
    func ifLet<
        DestinationState, DestinationAction, Destination: Reducer<DestinationState, DestinationAction>
    >(
        _ toWrappedState: WritableKeyPath<State, DestinationState?>,
        action toWrappedAction: CaseKeyPath<Action, DestinationAction>,
        @ReducerBuilder<DestinationState, DestinationAction> destination: () -> Destination,
        fileID: StaticString = #fileID,
        filePath: StaticString = #filePath,
        line: UInt = #line,
        column: UInt = #column
    ) -> some Reducer<State, Action> where DestinationState: Identifiable {
        ifLet(
            toWrappedState,
            id: \.id, action: toWrappedAction,
            destination: destination,
            fileID: fileID,
            filePath: filePath,
            line: line,
            column: column
        )
    }

    /// ifLet but keyed by ID, clearing the child store instance whenever the ID changes.
    func ifLet<
        ID: Hashable,
        DestinationState, DestinationAction, Destination: Reducer<DestinationState, DestinationAction>
    >(
        _ toWrappedState: WritableKeyPath<State, DestinationState?>,
        id: KeyPath<DestinationState, ID>,
        action toWrappedAction: CaseKeyPath<Action, DestinationAction>,
        @ReducerBuilder<DestinationState, DestinationAction> destination: () -> Destination,
        fileID: StaticString = #fileID,
        filePath: StaticString = #filePath,
        line: UInt = #line,
        column: UInt = #column
    ) -> some Reducer<State, Action> {
        _IfLetIDReducer(
            base: self,
            destination: destination(),
            toWrappedID: id,
            toWrappedState: toWrappedState,
            toWrappedAction: toWrappedAction,
            fileID: fileID,
            filePath: filePath,
            line: line,
            column: column
        )
    }
}

public extension Store where State: ObservableState {
    func scope<ChildState, ChildAction>(
      state stateKeyPath: KeyPath<State, ChildState?>,
      action actionKeyPath: CaseKeyPath<Action, ChildAction>,
      fileID: StaticString = #fileID,
      filePath: StaticString = #filePath,
      line: UInt = #line,
      column: UInt = #column
    ) -> Store<ChildState, ChildAction>? where ChildState: Identifiable {
        scope(
            state: stateKeyPath,
            id: \.id,
            action: actionKeyPath,
            fileID: fileID,
            filePath: filePath,
            line: line,
            column: column
        )
    }

    func scope<ChildState, ChildAction, ID>(
      state stateKeyPath: KeyPath<State, ChildState?>,
      id stateIDKeyPath: KeyPath<ChildState, ID>,
      action actionKeyPath: CaseKeyPath<Action, ChildAction>,
      fileID: StaticString = #fileID,
      filePath: StaticString = #filePath,
      line: UInt = #line,
      column: UInt = #column
    ) -> Store<ChildState, ChildAction>? where ID: Hashable {
        let id = state[keyPath: stateKeyPath]?[keyPath: stateIDKeyPath]
        let path = stateKeyPath.appending(path: \.[_id: id])
        return scope(
            state: path,
            action: actionKeyPath,
            fileID: fileID,
            filePath: filePath,
            line: line,
            column: column
        )
    }
}

private struct _IfLetIDReducer<Base: Reducer, Destination: Reducer, ID: Hashable>: Reducer {
    let base: Base
    let destination: Destination
    let toWrappedID: KeyPath<Destination.State, ID>
    let toWrappedState: WritableKeyPath<Base.State, Destination.State?>
    let toWrappedAction: CaseKeyPath<Base.Action, Destination.Action>
    let fileID: StaticString
    let filePath: StaticString
    let line: UInt
    let column: UInt

    func reduce(into state: inout Base.State, action: Base.Action) -> Effect<Base.Action> {
        let id = state[keyPath: toWrappedState]?[keyPath: toWrappedID]
        let path = toWrappedState.appending(path: \.[_id: id])
        return base.ifLet(
            path,
            action: toWrappedAction,
            then: { destination },
            fileID: fileID,
            filePath: filePath,
            line: line,
            column: column
        )
        .reduce(into: &state, action: action)
    }
}

private extension Optional {
    /// - note: Only used to update internal navigationIDPath
    subscript(_id id: (some Hashable)?) -> Wrapped? {
        get { self }
        set { self = newValue }
    }
}