An implementation of the catamorphism recursion scheme in Swift

Catamorphism.swift

// NOTE: This code was written when I first figured out how to encode HKT in Swift.
//       There is a lot that can be improved and I would write it somewhat differently today.

// This example shows how higher-kinded types can be emulated in Swift today.
// It acheives correct typing at the cost of some boilerplate, manual lifting and an existential representation.
// The technique below was directly inspired by the paper Lightweight Higher-Kinded Polymorphism
// by Jeremy Yallop and Leo White found at http://ocamllabs.io/higher/lightweight-higher-kinded-polymorphism.pdf

/// `ConstructorTag` represents a type constructor.
/// `Argument` represents an argument to the type constructor.
struct Apply<ConstructorTag, Argument> {
    /// An existential containing a value of `Constructor<Argument>`
    /// Where `Constructor` is the type constructor represented by `ConstructorTag`
    let tag: ConstructorTag
}

protocol TypeConstructor2 {
    associatedtype Tag
    associatedtype Argument1    
    associatedtype Argument2
    typealias Applied = Self
}
protocol Apply2Protocol {
    associatedtype Tag
    associatedtype Argument1    
    associatedtype Argument2
    //associatedtype Applied
}
protocol Apply2Tag {
    associatedtype Tag
    // This has arbitrary type arguments. It is only used as an argument to Replace2
    associatedtype Applied: TypeConstructor2 where Applied.Tag == Self
}
struct Apply2<ConstructorTag: Apply2Tag, Arg1, Arg2>: Apply2Protocol {
    /// An existential containing a value of `Constructor<Argument1, Argument2>`
    /// Where `Constructor` is the type constructor represented by `ConstructorTag`
    let tag: ConstructorTag
    
    typealias Tag = ConstructorTag
    typealias Argument1 = Arg1
    typealias Argument2 = Arg2
    typealias Applied = ConstructorTag.Applied
}
protocol Replace2Protocol {
    associatedtype Constructor: TypeConstructor2
    associatedtype Argument1    
    associatedtype Argument2    
    associatedtype Replaced: TypeConstructor2 where Replaced.Tag == Constructor.Tag, Replaced.Argument1 == Argument1, Replaced.Argument2 == Argument2
}
struct ArrayReplacer<A: TypeConstructor2, Argument> where A.Tag == ArrayTag {
    typealias Replaced = [Argument]
}

/// A protocol all type constructors must conform to.
protocol TypeConstructor {
    /// The existential type that erases `Argument`.
    /// This should only be initializable with values of types created by the current constructor.
    associatedtype Tag
    /// The argument that is currently applied to the type constructor in `Self`.
    associatedtype Argument
    
    /// `self` stored in the Tag existential
    var apply: Apply<Tag, Argument> { get } 
    /// Must unwrap the `app.tag` existential.
    static func unapply(_ apply: Apply<Tag, Argument>) -> Self
}

struct ArrayTag {
    fileprivate let array: Any
    // Private access to the initializer is what makes this a safe technique.
    // Creating an `Apply` (where the type information is stored)
    // requires creating a `Tag` first.
    // Using access control we can restrict that to the same file that defines
    // the `Array: TypeConstructor` conformance below to ensure that
    // `Apply<ArrayTag, T>` instances are only created with the correct type of
    // array values.
    init<T>(_ array: [T]) {
        self.array = array
    }
}
extension Array: TypeConstructor {
    typealias Tag = ArrayTag
    var apply: Apply<Tag, Element> {
        return Apply<Tag, Element>(tag: ArrayTag(self))
    }
    static func unapply(_ apply: Apply<Tag, Element>) -> Array {
        return apply.tag.array as! Array
    }
}

struct OptionalTag {
    fileprivate let optional: Any
    init<T>(_ optional: T?) {
        self.optional = optional as Any
    }
}
extension Optional: TypeConstructor {
    typealias Tag = OptionalTag
    var apply: Apply<Tag, Wrapped> {
        return Apply<Tag, Wrapped>(tag: OptionalTag(self))
    }
    static func unapply(_ apply: Apply<Tag, Wrapped>) -> Optional {
        return apply.tag.optional as? Wrapped
    }
}

protocol Functor: TypeConstructor {
    func map<T>(_ transform: (Argument) -> T) -> Apply<Tag, T>
}
protocol Monad: TypeConstructor {
    static func wrap<T>(_ value: T) -> Apply<Tag, T>
    func flatMap<T>(_ continuation: (Argument) -> Apply<Tag, T>) -> Apply<Tag, T>
}
extension Array: Monad {
    static func wrap<T>(_ value: T) -> Apply<Tag, T> {
        return [value].apply
    }
    func flatMap<T>(_ continuation: (Element) -> Apply<Tag, T>) -> Apply<Tag, T> {
        return flatMap { [T].unapply(continuation($0)) }.apply
    }
}
extension Optional: Monad {
    static func wrap<T>(_ value: T) -> Apply<Tag, T> {
        return (value as T?).apply
    }
    func flatMap<T>(_ continuation: (Wrapped) -> Apply<Tag, T>) -> Apply<Tag, T> {
        return flatMap { T?.unapply(continuation($0)) }.apply
    }
}

// Here we use flatMap on values of types [Int] and Int?.
// The result is automatically lifted into the corresponding emulated HKT.
// [1, 2, 3, 2, 4, 6, 3, 6, 9, 4, 8, 12]
Array.unapply([1, 2, 3, 4].flatMap { [$0, $0 * 2, $0 * 3].apply })
Optional.unapply((42 as Int?).flatMap { (($0 * 2) as Int?).apply }) // 84
Optional.unapply((nil as Int?).flatMap { _ in (nil as Int?).apply }) // nil

protocol FunctorTag {
    static func map<T, U>(_ value: Apply<Self, T>, _ transform: (T) -> U) -> Apply<Self, U>
}
protocol MonadTag {
    static func wrap<T>(_ value: T) -> Apply<Self, T>
    static func flatMap<T, U>(_ value: Apply<Self, T>, _ continuation: (T) -> Apply<Self, U>) -> Apply<Self, U>
}
extension ArrayTag: MonadTag {
    static func wrap<T>(_ value: T) -> Apply<ArrayTag, T> {
        return [value].apply
    }
    static func flatMap<T, U>(_ value: Apply<ArrayTag, T>, _ continuation: (T) -> Apply<ArrayTag, U>) -> Apply<ArrayTag, U> {
        return Array.unapply(value).flatMap(continuation)
    }
}
extension OptionalTag: MonadTag {
    static func wrap<T>(_ value: T) -> Apply<OptionalTag, T> {
        return (value as T?).apply
    }
    static func flatMap<T, U>(_ value: Apply<OptionalTag, T>, _ continuation: (T) -> Apply<OptionalTag, U>) -> Apply<OptionalTag, U> {
        return Optional.unapply(value).flatMap(continuation)
    }
}

/// We will soon be able to declare conformances for the emulated HKT existentials themselves!
extension Apply/*: Monad */ where ConstructorTag: MonadTag {
    static func wrap<T>(_ value: T) -> Apply<ConstructorTag, T> {
        return ConstructorTag.wrap(value)
    }
    func flatMap<T>(_ continuation: (Argument) -> Apply<ConstructorTag, T>) -> Apply<ConstructorTag, T> {
        return ConstructorTag.flatMap(self, continuation)
    }
}

// Here we use flatMap directly on the emulated HKT values of types Apply<ArrayTag, Int> 
// and Array<OptionalTag, Int> and observe the same results as flatMap applied to the base types.
// [1, 2, 3, 2, 4, 6, 3, 6, 9, 4, 8, 12]
Array.unapply([1, 2, 3, 4].apply.flatMap { [$0, $0 * 2, $0 * 3].apply })
Optional.unapply((42 as Int?).apply.flatMap { (($0 * 2) as Int?).apply }) // 84
Optional.unapply((nil as Int?).apply.flatMap { _ in (nil as Int?).apply }) // nil

protocol NaturalTransformation {
    associatedtype FromTag
    associatedtype ToTag
    
    static func apply<T>(to value: Apply<FromTag, T>) -> Apply<ToTag, T>
}

// A natural transformation from T? to [t]
enum OptionalToArray: NaturalTransformation {
    static func apply<T>(to optional: Apply<OptionalTag, T>) -> Apply<ArrayTag, T> {
        return [Optional.unapply(optional)].flatMap { $0 }.apply
    }
}

extension Apply {
    func transform<Transformation: NaturalTransformation>(using transformation: Transformation.Type) -> Apply<Transformation.ToTag, Argument> where Transformation.FromTag == ConstructorTag {
        return Transformation.apply(to: self)
    }
}

// Apply the natural transformation to values of the emulated HKT type Apply<OptionalTag, Int>
// to receive values of emulated HKT type Apply<ArrayTag, Int> and then unwrap the result.
Array.unapply((42 as Int?).apply.transform(using: OptionalToArray.self)) // [42]
Array.unapply((nil as Int?).apply.transform(using: OptionalToArray.self)) // []

struct Fix<ConstructorTag> {
    var unfix: Apply<ConstructorTag, Fix<ConstructorTag>>
    init(_ value: Apply<ConstructorTag, Fix<ConstructorTag>>) {
        unfix = value
    }
}

struct ExprTag {
    fileprivate let expr: Any
    init<T>(_ expr: Expr<T>) {
        self.expr = expr
    }
}

enum Expr<T> {
    case int(Int)
    indirect case add(T, T)
    indirect case mul(T, T)
}
extension Expr: TypeConstructor {
    typealias Tag = ExprTag
    
    var apply: Apply<ExprTag, T> {
        return Apply(tag: ExprTag(self))
    }
    static func unapply(_ apply: Apply<ExprTag, T>) -> Expr {
        return apply.tag.expr as! Expr
    }
}
extension Apply where ConstructorTag == ExprTag {
    static func int(_ i: Int) -> Apply {
        return Expr<Argument>.int(i).apply
    }
    static func add(_ lhs: Argument, _ rhs: Argument) -> Apply {
        return Expr<Argument>.add(lhs, rhs).apply
    }
    static func mul(_ lhs: Argument, _ rhs: Argument) -> Apply {
        return Expr<Argument>.mul(lhs, rhs).apply
    }
}
extension Expr: Functor {
    func map<U>(_ transform: (T) -> U) -> Apply<ExprTag, U> {
        switch self {
        case .int(let i): return Expr<U>.int(i).apply
        case .add(let lhs, let rhs): 
            return Expr<U>.add(transform(lhs), transform(rhs)).apply
        case .mul(let lhs, let rhs):
            return Expr<U>.mul(transform(lhs), transform(rhs)).apply
        }
    }
}
extension ExprTag: FunctorTag {
    static func map<T, U>(_ value: Apply<ExprTag, T>, _ transform: (T) -> U) -> Apply<ExprTag, U> {
        return Expr<T>.unapply(value).map(transform)
    }
}
extension Expr where T == Int {
    var eval: T {
        switch self {
            case .int(let val): return val
            case .add(let lhs, let rhs): return lhs + rhs
            case .mul(let lhs, let rhs): return lhs * rhs
        }
    }
}
extension Apply where ConstructorTag == ExprTag, Argument == Int {
    var eval: Argument {
        return Expr<Int>.unapply(self).eval
    }
}
func cata<Tag: FunctorTag, R, A>(
    _ fAlgebra: @escaping (Apply<Tag, A>) -> A, 
    _ out: @escaping (R) -> Apply<Tag, R>, 
    _ value: R
    ) -> A {
    return fAlgebra(Tag.map(out(value)) { cata(fAlgebra, out, $0 ) })
}
func evalAlgebra(_ expr: Apply<ExprTag, Int>) -> Int {
    return Expr<Int>.unapply(expr).eval
}
func out(_ fix: Fix<ExprTag>) -> Apply<ExprTag, Fix<ExprTag>> {
    return fix.unfix
}
func eval(_ fix: Fix<ExprTag>) -> Int {
    return cata(evalAlgebra, out, fix)
}
func cataFix<Tag: FunctorTag, A>(
    _ value: Fix<Tag>,
    _ fAlgebra: @escaping (Apply<Tag, A>) -> A
    ) -> A {
    return fAlgebra(Tag.map(value.unfix) { cataFix($0, fAlgebra) })
}
let evalFix = { cataFix($0) { Expr<Int>.unapply($0).eval }} // $0.unapply.eval

let expr = Fix<ExprTag>(.mul(
    Fix<ExprTag>(.int(42)), 
    Fix<ExprTag>(.add(
        Fix<ExprTag>(.int(42)), 
        Fix<ExprTag>(.int(42))
        ))))
let evaled = eval(expr)
let evaledFix = evalFix(expr)
let val = 42 * (42+42)