Bug on UIKit about swift-composable-architecture HOT 5 CLOSED

import SwiftUI
import ComposableArchitecture
import Combine

@Reducer
struct Feature {
    @ObservableState
    struct State: Equatable {
        let id: Int
        var child: ChildFeature.State?

        var isSelected: Bool = false
        var isPending: Bool = false
    }

    enum Action {
        case task
        case selectionReceived(Int)
        case pendingSelectionReceived(Set<Int>)
        case child(ChildFeature.Action)
    }

    let child: ChildFeature
    @Dependency(\.selectionClient) var selectionClient

    var body: some Reducer<State, Action> {
        Reduce { state, action in
            switch action {
            case .task:
                return .run { send in
                    try await withThrowingTaskGroup(of: Void.self) { group in
                        group.addTask {
                            for await selection in selectionClient.selection.values {
                                await send(.selectionReceived(selection))
                            }
                        }
                        group.addTask {
                            for await selections in selectionClient.pendingSelection.values {
                                await send(.pendingSelectionReceived(selections))
                            }
                        }
                        try await group.waitForAll()
                    }
                }
            case let .selectionReceived(id):
                let stateID = state.id
                state.isSelected = id == stateID
                return updateEffect(state: &state)
            case let .pendingSelectionReceived(ids):
                let stateID = state.id
                state.isPending = ids.contains(stateID)
                return updateEffect(state: &state)
            case .child:
                return .none
            }
        }
        .ifLet(\.child, action: \.child) {
            child
        }
    }

    func updateEffect(state: inout State) -> Effect<Action> {
        let isPending = state.isPending
        let isSelected = state.isSelected
        let shouldSelect = isPending || isSelected
        if shouldSelect != (state.child != nil) {
            state.child = shouldSelect ? .init(id: state.id) : nil
        }
        state.child?.isSelected = isSelected
        state.child?.isPending = isPending
        return .none
    }
}

@Reducer
struct ChildFeature {
    @ObservableState
    struct State: Equatable {
        let id: Int
        var seconds: Int = 0
        var isSelected: Bool = false
        var isPending: Bool = false
    }

    enum Action {
        case task
        case timeReceived
    }

    var body: some Reducer<State, Action> {
        Reduce { state, action in
            switch action {
            case .task:
                return .run { [id = state.id] send in
                    await withTaskCancellationHandler(operation: {
                        for await _ in Timer.publish(every: 1, on: .main, in: .common).autoconnect().values {
                            await send(.timeReceived)
                        }
                    }, onCancel: {
                        if id == 0 {
                            print("CANCEL \(id)")
                        }
                    })
                }
            case .timeReceived:
                state.seconds += 1
                return .none
            }
        }
    }
}

struct ContentView: View {
    let store: StoreOf<Feature>

    var body: some View {
        ZStack {
            if let childStore = store.scope(state: \.child, action: \.child) {
                ChildView(store: childStore)
            }
        }
        .task {
            await store.send(.task).finish()
        }
    }
}

struct ChildView: View {
    let store: StoreOf<ChildFeature>

    var body: some View {
        ZStack {
            if store.isSelected {
                Color.green
            } else if store.isPending {
                Color.yellow
            }
            Text(store.seconds, format: .number)
                .font(.title)
        }
        .task {
            await store.send(.task).finish()
        }
    }
}

struct SelectionClient: DependencyKey {
    let selection: AnyPublisher<Int, Never>
    let select: (Int) -> Void

    let pendingSelection: AnyPublisher<Set<Int>, Never>
    let setPendingSelection: (Set<Int>) -> Void

    static let liveValue: SelectionClient = {
        let subject = CurrentValueSubject<Int, Never>(0)
        let pendingSelectionSubject = CurrentValueSubject<Set<Int>, Never>([])
        return .init(
            selection: subject.eraseToAnyPublisher(),
            select: { subject.value = $0 },
            pendingSelection: pendingSelectionSubject.eraseToAnyPublisher(),
            setPendingSelection: { pendingSelectionSubject.value = $0 }
        )
    }()
}

extension DependencyValues {
    var selectionClient: SelectionClient {
        self[SelectionClient.self]
    }
}

@main
struct TCABugApp: App {
    @State var idSelected = false

    var body: some Scene {
        WindowGroup {
            VStack {
                Button(idSelected ? "ID" : "No ID") {
                    idSelected.toggle()
                }
                HostingView(idView: idSelected)
                    .id(idSelected)
            }
        }
    }
}

struct HostingView: UIViewControllerRepresentable {
    let idView: Bool

    func makeUIViewController(context: Context) -> ViewController {
        .init(idView: idView)
    }

    func updateUIViewController(_ uiViewController: ViewController, context: Context) { }
}

extension HostingView {
    final class ViewController: UIPageViewController, UIPageViewControllerDataSource, UIPageViewControllerDelegate {
        private let idView: Bool
        @Dependency(\.selectionClient) private var selectionClient

        private lazy var childControllers: [UIViewController] = [
            childController(id: 0),
            childController(id: 1),
            childController(id: 2),
            childController(id: 3)
        ]

        init(idView: Bool) {
            self.idView = idView
            super.init(transitionStyle: .scroll, navigationOrientation: .vertical)
        }

        required init?(coder: NSCoder) {
            fatalError("init(coder:) has not been implemented")
        }

        override func viewDidLoad() {
            super.viewDidLoad()
            dataSource = self
            delegate = self
            selectionClient.select(0)
            setViewControllers([childControllers[0]], direction: .forward, animated: false)
        }

        func childController(id: Int) -> UIViewController {
            let view = childView(id: id)
            return UIHostingController(rootView: view)
        }

        @ViewBuilder
        func childView(id: Int) -> some View {
            if idView {
                let store = IDStoreOf<Feature>(initialState: .init(id: id)) {
                    Feature(child: .init())
                }
                IDStoreView(store: store) { store in
                    ContentView(store: store)
                }
            } else {
                ContentView(store: .init(initialState: .init(id: id), reducer: {
                    Feature(child: .init())
                }))
            }
        }

        func pageViewController(_ pageViewController: UIPageViewController, viewControllerBefore viewController: UIViewController) -> UIViewController? {
            guard let index = childControllers.firstIndex(of: viewController), childControllers.indices.contains(index - 1) else { return nil }
            return childControllers[index - 1]
        }

        func pageViewController(_ pageViewController: UIPageViewController, viewControllerAfter viewController: UIViewController) -> UIViewController? {
            guard let index = childControllers.firstIndex(of: viewController), childControllers.indices.contains(index + 1) else { return nil }
            return childControllers[index + 1]
        }

        func pageViewController(_ pageViewController: UIPageViewController, willTransitionTo pendingViewControllers: [UIViewController]) {
            selectionClient.setPendingSelection(Set(pendingViewControllers.compactMap {
                childControllers.firstIndex(of: $0)
            }))
        }

        func pageViewController(_ pageViewController: UIPageViewController, didFinishAnimating finished: Bool, previousViewControllers: [UIViewController], transitionCompleted completed: Bool) {
            if completed, let viewController = pageViewController.viewControllers?.last, let index = childControllers.firstIndex(of: viewController) {
                selectionClient.select(index)
            }
            selectionClient.setPendingSelection([])
        }
    }
}

struct IDStoreView<ChildReducer, ChildView>: View where ChildView: View, ChildReducer: Reducer {
    let store: StoreOf<IDReducer>
    let content: (StoreOf<ChildReducer>) -> ChildView

    init(
        store idStore: IDStoreOf<ChildReducer>,
        content: @escaping (StoreOf<ChildReducer>) -> ChildView
    ) {
        let id = ChildID()
        self.store = .init(
            initialState: .init(idChild: .init(id: id, child: idStore.initialState)),
            reducer: { IDReducer(child: idStore.reducer) }
        )
        self.content = content
    }

    var body: some View {
        ForEach(store.scope(state: \.elements, action: \.elements), id: \.id) { store in
            content(store.scope(state: \.child, action: \.child))
        }
    }
}

struct IDStoreOf<ChildReducer> where ChildReducer: Reducer {
    let initialState: ChildReducer.State
    let reducer: ChildReducer

    init(
        initialState: ChildReducer.State,
        @ReducerBuilder<ChildReducer.State, ChildReducer.Action> reducer: () -> ChildReducer
    ) {
        self.initialState = initialState
        self.reducer = reducer()
    }

    init<BaseReducer>(
        initialState: @autoclosure () -> ChildReducer.State,
        @ReducerBuilder<ChildReducer.State, ChildReducer.Action> reducer: () -> ChildReducer,
        withDependencies prepareDependencies: (inout DependencyValues) -> Void
    ) where ChildReducer == _DependencyKeyWritingReducer<BaseReducer> {
        let (initialState, reducer, dependencies) = withDependencies(prepareDependencies) {
            @Dependency(\.self) var dependencies
            return (initialState(), reducer(), dependencies)
        }
        self.init(
            initialState: initialState,
            reducer: { reducer.dependency(\.self, dependencies) }
        )
    }
}

extension IDStoreView {
    typealias ChildID = UUID

    @Reducer
    struct IDReducer {
        @ObservableState
        struct State {
            var elements: IdentifiedArrayOf<IDChild.State>

            init(idChild: IDChild.State) {
                self.elements = .init(uniqueElements: [idChild])
            }
        }

        enum Action {
            case elements(IdentifiedActionOf<IDChild>)
        }

        let child: ChildReducer

        var body: some Reducer<State, Action> {
            Reduce { state, action in
                switch action {
                case .elements:
                    return .none
                }
            }
            .forEach(\.elements, action: \.elements) {
                IDChild(child: child)
            }
        }
    }
}

extension IDStoreView.IDReducer {
    @Reducer
    struct IDChild {
        @ObservableState
        struct State: Identifiable {
            var id: IDStoreView.ChildID
            var child: ChildReducer.State
        }

        enum Action {
            case child(ChildReducer.Action)
        }

        let child: ChildReducer

        var body: some Reducer<State, Action> {
            Scope(state: \.child, action:  \.child) {
                child
            }
        }
    }
}

import SwiftUI
import ComposableArchitecture
import Combine

@Reducer
struct ParentFeature {
  @ObservableState
  struct State {
    var features: IdentifiedArrayOf<Feature.State> = []
  }
  enum Action {
    case features(IdentifiedActionOf<Feature>)
  }
  var body: some ReducerOf<Self> {
    Reduce { state, action in
      switch action {
      case .features:
        return .none
      }
    }
    .forEach(\.features, action: \.features) {
      Feature()
    }
  }
}

@Reducer
struct Feature {
  @ObservableState
  struct State: Equatable, Identifiable {
    let id: Int
    var child: ChildFeature.State?

    var isSelected: Bool = false
    var isPending: Bool = false
  }

  enum Action {
    case task
    case selectionReceived(Int)
    case pendingSelectionReceived(Set<Int>)
    case child(ChildFeature.Action)
  }

  @Dependency(\.selectionClient) var selectionClient

  var body: some Reducer<State, Action> {
    Reduce { state, action in
      switch action {
      case .task:
        return .run { send in
          try await withThrowingTaskGroup(of: Void.self) { group in
            group.addTask {
              for await selection in selectionClient.selection.values {
                await send(.selectionReceived(selection))
              }
            }
            group.addTask {
              for await selections in selectionClient.pendingSelection.values {
                await send(.pendingSelectionReceived(selections))
              }
            }
            try await group.waitForAll()
          }
        }
      case let .selectionReceived(id):
        let stateID = state.id
        state.isSelected = id == stateID
        return updateEffect(state: &state)
      case let .pendingSelectionReceived(ids):
        let stateID = state.id
        state.isPending = ids.contains(stateID)
        return updateEffect(state: &state)
      case .child:
        return .none
      }
    }
    .ifLet(\.child, action: \.child) {
      ChildFeature()
    }
  }

  func updateEffect(state: inout State) -> Effect<Action> {
    let isPending = state.isPending
    let isSelected = state.isSelected
    let shouldSelect = isPending || isSelected
    if shouldSelect != (state.child != nil) {
      state.child = shouldSelect ? .init(id: state.id) : nil
    }
    state.child?.isSelected = isSelected
    state.child?.isPending = isPending
    return .none
  }
}

@Reducer
struct ChildFeature {
  @ObservableState
  struct State: Equatable {
    let id: Int
    var seconds: Int = 0
    var isSelected: Bool = false
    var isPending: Bool = false
  }

  enum Action {
    case task
    case timeReceived
  }

  var body: some Reducer<State, Action> {
    Reduce { state, action in
      switch action {
      case .task:
        return .run { [id = state.id] send in
          await withTaskCancellationHandler(operation: {
            for await _ in Timer.publish(every: 1, on: .main, in: .common).autoconnect().values {
              await send(.timeReceived)
            }
          }, onCancel: {
            if id == 0 {
              print("CANCEL \(id)")
            }
          })
        }
      case .timeReceived:
        state.seconds += 1
        return .none
      }
    }
  }
}

struct ContentView: View {
  let store: StoreOf<Feature>

  var body: some View {
    ZStack {
      if let childStore = store.scope(state: \.child, action: \.child) {
        ChildView(store: childStore)
      }
    }
    .task {
      await store.send(.task).finish()
    }
  }
}

struct ChildView: View {
  let store: StoreOf<ChildFeature>

  var body: some View {
    ZStack {
      if store.isSelected {
        Color.green
      } else if store.isPending {
        Color.yellow
      }
      Text(store.seconds, format: .number)
        .font(.title)
    }
    .task {
      await store.send(.task).finish()
    }
  }
}

struct SelectionClient: DependencyKey {
  let selection: AnyPublisher<Int, Never>
  let select: (Int) -> Void

  let pendingSelection: AnyPublisher<Set<Int>, Never>
  let setPendingSelection: (Set<Int>) -> Void

  static let liveValue: SelectionClient = {
    let subject = CurrentValueSubject<Int, Never>(0)
    let pendingSelectionSubject = CurrentValueSubject<Set<Int>, Never>([])
    return .init(
      selection: subject.eraseToAnyPublisher(),
      select: { subject.value = $0 },
      pendingSelection: pendingSelectionSubject.eraseToAnyPublisher(),
      setPendingSelection: { pendingSelectionSubject.value = $0 }
    )
  }()
}

extension DependencyValues {
  var selectionClient: SelectionClient {
    self[SelectionClient.self]
  }
}

@main
struct TCABugApp: App {
  @State var idSelected = false

  var body: some Scene {
    WindowGroup {
      HostingView()
    }
  }
}

struct HostingView: UIViewControllerRepresentable {
  func makeUIViewController(context: Context) -> ViewController {
    .init()
  }
  func updateUIViewController(_ uiViewController: ViewController, context: Context) { }
}

extension HostingView {
  final class ViewController: UIPageViewController, UIPageViewControllerDataSource, UIPageViewControllerDelegate {
    @Dependency(\.selectionClient) private var selectionClient
    let store = Store(
      initialState: ParentFeature.State(
        features: [
          Feature.State(id: 0),
          Feature.State(id: 1),
          Feature.State(id: 2),
          Feature.State(id: 3),
        ]
      )
    ) {
      ParentFeature()
    }

    private var featureControllers: [UIViewController] = []
    init() {
      super.init(transitionStyle: .scroll, navigationOrientation: .vertical)
    }

    required init?(coder: NSCoder) {
      fatalError("init(coder:) has not been implemented")
    }

    override func viewDidLoad() {
      super.viewDidLoad()
      dataSource = self
      delegate = self
      observe { [weak self] in
        guard let self else { return }
        featureControllers = store.features.ids.map { id in
          UIHostingController(
            rootView: ContentView(
              store: self.store.scope(
                state: \.features[id: id]!,
                action: \.features[id: id]
              )
            )
          )
        }
      }
      selectionClient.select(0)
      setViewControllers([featureControllers[0]], direction: .forward, animated: false)
    }

    func pageViewController(_ pageViewController: UIPageViewController, viewControllerBefore viewController: UIViewController) -> UIViewController? {
      guard let index = featureControllers.firstIndex(of: viewController), featureControllers.indices.contains(index - 1) else { return nil }
      return featureControllers[index - 1]
    }

    func pageViewController(_ pageViewController: UIPageViewController, viewControllerAfter viewController: UIViewController) -> UIViewController? {
      guard let index = featureControllers.firstIndex(of: viewController), featureControllers.indices.contains(index + 1) else { return nil }
      return featureControllers[index + 1]
    }

    func pageViewController(_ pageViewController: UIPageViewController, willTransitionTo pendingViewControllers: [UIViewController]) {
      selectionClient.setPendingSelection(Set(pendingViewControllers.compactMap {
        featureControllers.firstIndex(of: $0)
      }))
    }

    func pageViewController(_ pageViewController: UIPageViewController, didFinishAnimating finished: Bool, previousViewControllers: [UIViewController], transitionCompleted completed: Bool) {
      if completed, let viewController = pageViewController.viewControllers?.last, let index = featureControllers.firstIndex(of: viewController) {
        selectionClient.select(index)
      }
      selectionClient.setPendingSelection([])
    }
  }
}