swift-composable-architecture-mirror/Sources/ComposableArchitecture/RootStore.swift

import Combine
import Foundation

@_spi(Internals)
@MainActor
public final class RootStore {
  private var bufferedActions: [Any] = []
  let didSet = CurrentValueRelay(())
  @_spi(Internals) public var effectCancellables: [UUID: AnyCancellable] = [:]
  private var isSending = false
  private let reducer: any Reducer
  private(set) var state: Any {
    didSet {
      didSet.send(())
    }
  }

  init<State, Action>(
    initialState: State,
    reducer: some Reducer<State, Action>
  ) {
    self.state = initialState
    self.reducer = reducer
  }

  func send(_ action: Any, originatingFrom originatingAction: Any? = nil) -> Task<Void, Never>? {
    func open<State, Action>(reducer: some Reducer<State, Action>) -> Task<Void, Never>? {
      self.bufferedActions.append(action)
      guard !self.isSending else { return nil }

      self.isSending = true
      var currentState = self.state as! State
      let tasks = LockIsolated<[Task<Void, Never>]>([])
      defer {
        withExtendedLifetime(self.bufferedActions) {
          self.bufferedActions.removeAll()
        }
        self.state = currentState
        self.isSending = false
        if !self.bufferedActions.isEmpty {
          if let task = self.send(
            self.bufferedActions.removeLast(),
            originatingFrom: originatingAction
          ) {
            tasks.withValue { $0.append(task) }
          }
        }
      }

      var index = self.bufferedActions.startIndex
      while index < self.bufferedActions.endIndex {
        defer { index += 1 }
        let action = self.bufferedActions[index] as! Action
        let effect = reducer.reduce(into: &currentState, action: action)

        switch effect.operation {
        case .none:
          break
        case let .publisher(publisher):
          var didComplete = false
          let boxedTask = Box<Task<Void, Never>?>(wrappedValue: nil)
          let uuid = UUID()
          let effectCancellable = withEscapedDependencies { continuation in
            publisher
              .receive(on: UIScheduler.shared)
              .handleEvents(receiveCancel: { [weak self] in self?.effectCancellables[uuid] = nil })
              .sink(
                receiveCompletion: { [weak self] _ in
                  boxedTask.wrappedValue?.cancel()
                  didComplete = true
                  self?.effectCancellables[uuid] = nil
                },
                receiveValue: { [weak self] effectAction in
                  guard let self else { return }
                  if let task = continuation.yield({
                    self.send(effectAction, originatingFrom: action)
                  }) {
                    tasks.withValue { $0.append(task) }
                  }
                }
              )
          }

          if !didComplete {
            let task = Task<Void, Never> { @MainActor in
              for await _ in AsyncStream<Void>.never {}
              effectCancellable.cancel()
            }
            boxedTask.wrappedValue = task
            tasks.withValue { $0.append(task) }
            self.effectCancellables[uuid] = effectCancellable
          }
        case let .run(priority, operation):
          withEscapedDependencies { continuation in
            let task = Task(priority: priority) { @MainActor in
              let isCompleted = LockIsolated(false)
              defer { isCompleted.setValue(true) }
              await operation(
                Send { effectAction in
                  if isCompleted.value {
                    reportIssue(
                      """
                      An action was sent from a completed effect:

                        Action:
                          \(debugCaseOutput(effectAction))

                        Effect returned from:
                          \(debugCaseOutput(action))

                      Avoid sending actions using the 'send' argument from 'Effect.run' after \
                      the effect has completed. This can happen if you escape the 'send' \
                      argument in an unstructured context.

                      To fix this, make sure that your 'run' closure does not return until \
                      you're done calling 'send'.
                      """
                    )
                  }
                  if let task = continuation.yield({
                    self.send(effectAction, originatingFrom: action)
                  }) {
                    tasks.withValue { $0.append(task) }
                  }
                }
              )
            }
            tasks.withValue { $0.append(task) }
          }
        }
      }

      guard !tasks.isEmpty else { return nil }
      return Task { @MainActor in
        await withTaskCancellationHandler {
          var index = tasks.startIndex
          while index < tasks.endIndex {
            defer { index += 1 }
            await tasks[index].value
          }
        } onCancel: {
          var index = tasks.startIndex
          while index < tasks.endIndex {
            defer { index += 1 }
            tasks[index].cancel()
          }
        }
      }
    }
    return _withoutPerceptionChecking {
      open(reducer: self.reducer)
    }
  }
}