Emulating HKT in Swift

HKT.swift

// 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
}

/// 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 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 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)) // []