Effect/Send protocol

Effects.swift

// MARK: - Extensions

extension Task where Failure == Never {
  public var cancellableValue: Success {
    get async {
      await withTaskCancellationHandler {
        await self.value
      } onCancel: {
        self.cancel()
      }
    }
  }
}

extension Task where Failure == Error {
  public var cancellableValue: Success {
    get async throws {
      try await withTaskCancellationHandler {
        try await self.value
      } onCancel: {
        self.cancel()
      }
    }
  }
}

// MARK: - Protocols

public protocol Send<Value>: Sendable {
  associatedtype Value

  func callAsFunction(_ value: Value) async
}

public protocol Effect<Value>: Sendable {
  associatedtype Value

  func run(_ send: some Send<Value>) async
}

public enum Sends {}

public enum Effects {}

// MARK: - None

extension Effects {
  public struct None<Value>: Effect {
    public init() {}

    public func run(_ send: some Send<Value>) {}
  }
}
extension Effect {
  public static func none<Value>() -> Self where Self == Effects.None<Value> {
    .init()
  }
}

public typealias None = Effects.None

// MARK: - Task

extension Effects {
  public struct TaskEffect<Value>: Effect {
    public let priority: TaskPriority?
    public let operation: @Sendable () async -> Value?

    public init(priority: TaskPriority? = nil, operation: @escaping @Sendable () async -> Value?) {
      self.priority = priority
      self.operation = operation
    }

    public init(
      priority: TaskPriority? = nil,
      operation: @escaping @Sendable () async throws -> Value?,
      catch handleError: @escaping @Sendable (Error) async -> Value?
    ) {
      self.init(priority: priority) {
        do {
          return try await operation()
        } catch is CancellationError {
          return nil
        } catch {
          return await handleError(error)
        }
      }
    }

    public func run(_ send: some Send<Value>) async {
      await Task(priority: priority) {
        guard let value = await operation() else { return }
        await send(value)
      }.cancellableValue
    }
  }
}
extension Effect {
  public static func task<Value>(
    priority: TaskPriority? = nil,
    operation: @escaping @Sendable () async -> Value?
  ) -> Self where Self == Effects.TaskEffect<Value> {
    Effects.TaskEffect(priority: priority, operation: operation)
  }

  public static func task<Value>(
    priority: TaskPriority? = nil,
    operation: @escaping @Sendable () async throws -> Value?,
    catch handleError: @escaping @Sendable (Error) async -> Value?
  ) -> Self where Self == Effects.TaskEffect<Value> {
    Effects.TaskEffect(priority: priority, operation: operation, catch: handleError)
  }
}

public typealias TaskEffect = Effects.TaskEffect

public struct Sync<Value>: Effect {
  public let value: @Sendable () -> Value

  public init(_ value: @escaping @Sendable () -> Value) {
    self.value = value
  }

  public init(_ value: @escaping @autoclosure @Sendable () -> Value) {
    self.value = value
  }

  public func run(_ send: some Send<Value>) async {
    await send(value())
  }
}

// MARK: - Run

extension Effects {
  public struct Run<Value>: Effect {
    public let priority: TaskPriority?
    public let operation: @Sendable (any Send<Value>) async -> Void

    public init(
      priority: TaskPriority? = nil,
      run: @escaping @Sendable (any Send<Value>) async -> Void
    ) {
      self.priority = priority
      self.operation = run
    }

    public init(
      priority: TaskPriority? = nil,
      run: @escaping @Sendable (any Send<Value>) async throws -> Void,
      catch handleError: @escaping @Sendable (Error, any Send<Value>) async -> Void
    ) {
      self.init(priority: priority) { send in
        do {
          try await run(send)
        } catch is CancellationError {
          return
        } catch {
          await handleError(error, send)
        }
      }
    }

    public func run(_ send: some Send<Value>) async {
      let task = Task(priority: priority) {
        await operation(send)
      }
      await task.cancellableValue
    }
  }
}
extension Effect {
  public static func run<Value>(
    priority: TaskPriority? = nil,
    run: @escaping @Sendable (any Send<Value>) async -> Void
  ) -> Self where Self == Effects.Run<Value> {
    Effects.Run(priority: priority, run: run)
  }

  public static func run<Value>(
    priority: TaskPriority? = nil,
    run: @escaping @Sendable (any Send<Value>) async throws -> Void,
    catch handleError: @escaping @Sendable (Error, any Send<Value>) async -> Void
  ) -> Self where Self == Effects.Run<Value> {
    Effects.Run(priority: priority) { send in
      do {
        try await run(send)
      } catch is CancellationError {
        return
      } catch {
        await handleError(error, send)
      }
    }
  }
}

public typealias Run = Effects.Run

// MARK: - Map

extension Sends {
  public struct Pullback<Base: Send, Value>: Send {
    public let base: Base
    public let transform: @Sendable (Value) async -> Base.Value

    @inlinable
    public init(base: Base, transform: @escaping @Sendable (Value) async -> Base.Value) {
      self.base = base
      self.transform = transform
    }

    @inlinable
    public func callAsFunction(_ value: Value) async {
      await base(transform(value))
    }
  }
}

extension Effects {
  public struct Map<Base: Effect, Value>: Effect {
    public let base: Base
    public let transform: @Sendable (Base.Value) async -> Value

    @inlinable
    public func run(_ send: some Send<Value>) async {
      await base.run(Sends.Pullback(base: send, transform: transform))
    }
  }
}
extension Effect {
  public func map<NewValue>(
    _ transform: @escaping @Sendable (Value) async -> NewValue
  ) -> Effects.Map<Self, NewValue> {
    Effects.Map(base: self, transform: transform)
  }
}

// MARK: - Merge

extension Effects {
  public struct Merge2<First: Effect, Second: Effect>: Effect where First.Value == Second.Value {
    public typealias Value = First.Value

    public let first: First
    public let second: Second

    public init(first: First, second: Second) {
      self.first = first
      self.second = second
    }

    public func run(_ send: some Send<Value>) async {
      async let a: Void = first.run(send)
      async let b: Void = second.run(send)
      _ = await (a, b)
    }
  }
}
extension Effect {
  public func merged<Other: Effect<Value>>(with other: Other) -> Effects.Merge2<Self, Other> {
    .init(first: self, second: other)
  }
}

// MARK: - Concatenate

extension Effects {
  public struct Chain2<First: Effect, Second: Effect>: Effect
  where First.Value == Second.Value {
    public typealias Value = First.Value

    public let first: First
    public let second: Second

    public init(first: First, second: Second) {
      self.first = first
      self.second = second
    }

    public func run(_ send: some Send<Value>) async {
      await first.run(send)
      await second.run(send)
    }
  }
}
extension Effect {
  public func chaining<Other: Effect<Value>>(_ other: Other) -> Effects.Chain2<Self, Other> {
    .init(first: self, second: other)
  }
}

// MARK: - Cancellation

public actor Cancellables {
  private var tasks: [ObjectIdentifier: Task<Void, Never>] = [:]

  public init() {}

  public init(rebasingFrom other: Cancellables) async {
    self.tasks = await other.tasks
  }

  @TaskLocal public static var current: Cancellables = .init()

  public func enqueueTask<ID>(_ task: Task<Void, Never>, forID: ID.Type, cancelInFlight: Bool) {
    let key = ObjectIdentifier(ID.self)
    if cancelInFlight, let existingTask = tasks.removeValue(forKey: key) {
      existingTask.cancel()
    }
    tasks[key] = task
  }

  public func cancelTask<ID>(forID: ID.Type) {
    guard let task = tasks.removeValue(forKey: ObjectIdentifier(ID.self)) else { return }
    task.cancel()
  }
}

extension Effects {
  public struct Cancellable<Base: Effect, ID>: Effect {
    public typealias Value = Base.Value

    public let base: Base
    public let cancelInFlight: Bool

    public init(base: Base, cancelInFlight: Bool) {
      self.base = base
      self.cancelInFlight = cancelInFlight
    }

    public func run(_ send: some Send<Base.Value>) async {
      let task = Task {
        await base.run(send)
      }
      await Cancellables.current.enqueueTask(
        task,
        forID: ID.self,
        cancelInFlight: cancelInFlight
      )
      await task.cancellableValue
    }
  }

  public struct Cancel<Value, ID>: Effect {
    public init(id: ID.Type) {}

    public func run(_ send: some Send<Value>) async {
      await Cancellables.current.cancelTask(forID: ID.self)
    }
  }
}
extension Effect {
  public func cancellable<ID>(
    id: ID.Type,
    cancelInFlight: Bool = false
  ) -> Effects.Cancellable<Self, ID> {
    .init(base: self, cancelInFlight: cancelInFlight)
  }

  public static func cancel<Value, ID>(id: ID.Type) -> Self
  where Self == Effects.Cancel<Value, ID> {
    .init(id: id)
  }

  public func cancel<ID>(id: ID.Type) -> Effects.Chain2<Self, Effects.Cancel<Value, ID>> {
    self.chaining(.cancel(id: id))
  }
}

public typealias Cancel = Effects.Cancel

// MARK: - AsyncSequence

extension Effects {
  public struct AsyncSequenceEffect<S: AsyncSequence & Sendable>: Effect {
    public typealias Value = S.Element

    public let sequence: S
    public let handleError: (@Sendable (Error, any Send<S.Element>) async -> Void)?

    public func run(_ send: some Send<S.Element>) async {
      do {
        for try await value in sequence {
          await send(value)
        }
      } catch is CancellationError {
        return
      } catch {
        guard let handleError else {
          return assertionFailure()
        }
        await handleError(error, send)
      }
    }
  }
}
extension AsyncSequence {
  public func effect(
    catch handleError: (@Sendable (Error, any Send<Element>) async -> Void)?
  ) -> Effects.AsyncSequenceEffect<Self> {
    .init(sequence: self, handleError: handleError)
  }

  public func effect(
    catch handleError: (@Sendable (Error) async -> Element?)? = nil
  ) -> Effects.AsyncSequenceEffect<Self> {
    if let handleError {
      return .init(sequence: self, handleError: { error, send in
        if let value = await handleError(error) {
          await send(value)
        }
      })
    } else {
      return .init(sequence: self, handleError: nil)
    }
  }
}

// MARK: - Combine

import Combine

extension Effects {
  public struct PublisherEffect<P: Publisher>: Effect where P.Failure == Never {
    public typealias Value = P.Output

    public let publisher: @Sendable () -> P

    public init(_ publisher: @escaping @Sendable () -> P) {
      self.publisher = publisher
    }

    public func run(_ send: some Send<Value>) async {
      for await value in publisher().values {
        guard !Task.isCancelled else { return }
        await send(value)
        guard !Task.isCancelled else { return }
      }
    }
  }
}

extension Publisher where Failure == Never {
  public var effect: Effects.PublisherEffect<Self> {
    .init({ self })
  }
}

public typealias PublisherEffect = Effects.PublisherEffect

// MARK: - Examples

enum Thingy {
  case first, second, third

  static func random() -> Self {
    switch Int.random(in: 1...3) {
    case 1: return .first
    case 2: return .second
    case 3: return .third
    default: fatalError()
    }
  }
}

func blahteonuh() -> any Effect<Int> {
  switch Thingy.random() {
  case .first:
    let effect = Run<Int> { send in
      await send(3 + 8)
      await Task.yield()
      await send(4 * 9)
    }
    .cancellable(id: Thingy.self, cancelInFlight: true)

    if Bool.random() {
      return effect.chaining(Sync(9))
    } else {
      return effect
    }
  case .second:
    return None()
  case .third:
    return Cancel(id: Thingy.self)
  }
}