Traits flatMap with FB

fbsample.swift

struct Single<T> {}
struct Maybe<T> {}

protocol MaybeCov {
    associatedtype T
    func asMaybe() -> Maybe<T>
}

extension Single: MaybeCov {
    func asMaybe() -> Maybe<T> { fatalError() }
}

extension Maybe: MaybeCov {
    func asMaybe() -> Maybe<T> { fatalError() }
}

@_functionBuilder
struct Builder {
    enum Either<First, Second> {
        case first(First)
        case second(Second)
    }

    public static func buildBlock() -> Void {

    }

    public static func buildBlock<C>(_ c: C) -> C {
        c
    }

    public static func buildEither<T, F>(first: T) -> Either<T, F> {
        .first(first)
    }

    public static func buildEither<T, F>(second: F) -> Either<T, F> {
        .second(second)
    }
}

extension Builder.Either: MaybeCov where First: MaybeCov, Second: MaybeCov, First.T == Second.T {
    func asMaybe() -> Maybe<First.T> {
        switch self {
        case .first(let first): return first.asMaybe()
        case .second(let second): return second.asMaybe()
        }
    }
}

extension Single {
    func flatMap<M: MaybeCov>(@Builder _ f: (T) -> M) -> Maybe<M.T> {
        f(fatalError() as! T).asMaybe()
    }
}


let res = Single<Int>().flatMap { (arg) in
    if arg % 2 == 0 {
        Single<String>()
    } else {
        Maybe<String>()
    }
}
// res is Maybe<String>