flatMapBuild.swift

import Combine

@_functionBuilder
struct FlatMapBuilder {
    static func buildBlock<C: Publisher>(_ component: C) -> C
    {
        component
    }

    static func buildEither<F: Publisher, S: Publisher>(first component: F) -> EitherPublisher<F, S>
    where
        F.Output == S.Output,
        F.Failure == S.Failure
    {
        .left(component)
    }

    static func buildEither<F: Publisher, S: Publisher>(first component: F) -> EitherPublisher<Publishers.SetFailureType<F, S.Failure>, S>
    where
        F.Output == S.Output,
        F.Failure == Never
    {
        .left(component.setFailureType(to: S.Failure.self))
    }

    static func buildEither<F: Publisher, S: Publisher>(second component: S) -> EitherPublisher<F, S>
    where
        F.Output == S.Output,
        F.Failure == S.Failure
    {
        .right(component)
    }

    static func buildEither<F: Publisher, S: Publisher>(first component: S) -> EitherPublisher<F, Publishers.SetFailureType<S, F.Failure>>
    where
        F.Output == S.Output,
        S.Failure == Never
    {
        .right(component.setFailureType(to: F.Failure.self))
    }
}

enum EitherPublisher<L: Publisher, R: Publisher>: Publisher
where
    L.Output == R.Output, L.Failure == R.Failure
{
    typealias Output = L.Output
    typealias Failure = L.Failure
    case left(L)
    case right(R)

    func receive<S>(subscriber: S) where S : Subscriber, Self.Failure == S.Failure, Self.Output == S.Input {
        switch self {
        case .left(let value):
            value.receive(subscriber: subscriber)
        case .right(let value):
            value.receive(subscriber: subscriber)
        }
    }
}

extension Publisher {
    func flatMapBuild<O, P>(
        @FlatMapBuilder _ builder: @escaping (Self.Output) -> P
    ) -> Publishers.FlatMap<P, Self>
    where O == P.Output, Failure == P.Failure, P: Publisher
    {
        flatMap(builder)
    }

    func flatMapBuild<O, P>(
        @FlatMapBuilder _ builder: @escaping (Self.Output) -> P
    ) -> Publishers.FlatMap<Publishers.SetFailureType<P, Self.Failure>, Self>
    where O == P.Output, P: Publisher, P.Failure == Never
    {
        flatMap(builder)
    }
}

extension Publisher where Failure == Never {
    func flatMapBuild<O, P>(
        @FlatMapBuilder _ builder: @escaping (Self.Output) -> P
    ) -> Publishers.FlatMap<P, Publishers.SetFailureType<Self, P.Failure>>
    where O == P.Output, P: Publisher
    {
        flatMap(builder)
    }

    func flatMapBuild<O, P>(
        @FlatMapBuilder _ builder: @escaping (Self.Output) -> P
    ) -> Publishers.FlatMap<P, Self>
    where O == P.Output, P: Publisher, P.Failure == Never
    {
        flatMap(builder)
    }
}