algorithmw.swift

/*:
 
 This is an implementation of Algorithm W, as found in [Principal Types for functional programs](http://web.cs.wpi.edu/~cs4536/c12/milner-damas_principal_types.pdf).
 
 We'll start by defining literals and expressions:
 */

enum Literal {
    case string(String)
    case int(Int)
}

typealias ExpVariable = String

indirect enum Exp {
    case literal(Literal)
    case variable(ExpVariable)
    case app(Exp, Exp)
    case lam(ExpVariable, Exp)
    case let_(ExpVariable,Exp,Exp)
}

//: A primitive type is either a `string` or an `int`:

indirect enum PrimitiveType {
    case string, int
}

//: We'll use `String`s as type variables:
typealias TypeVariable = String

//: A type is either a variable, a primitive or a function type.
indirect enum Type {
    case variable(TypeVariable)
    case primitive(PrimitiveType)
    case function(Type,Type)
}

//: A type scheme is either a type, or a quantified type scheme. Alternatively, we could have defined this as `struct TypeScheme { var quantified: [String], var baseType: Type }`:
indirect enum TypeScheme {
    case type(Type)
    case forall(TypeVariable, TypeScheme)
}

//: `PrimitiveType`, `Type` and `TypeScheme` are made `Equatable` and `CustomStringConvertible`. The definitions are uninteresting but convenient.

extension PrimitiveType: Equatable, CustomStringConvertible {
    static func ==(l: PrimitiveType, r: PrimitiveType) -> Bool {
        switch (l,r) {
        case (.string,.string): return true
        case (.int,.int): return true
        default: return false
        }
    }
    
    var description: String {
        switch self {
        case .int:
            return "int"
        case .string:
            return "string"
        }
    }
}

extension Type: CustomStringConvertible, Equatable {
    var description: String {
        switch self {
        case .variable(let v): return v
        case .function(let l, let r): return "(\(l)) -> (\(r))"
        case .primitive(let p): return p.description
        }
    }

    static func ==(lhs: Type, rhs: Type) -> Bool {
        switch (lhs, rhs) {
        case let (.variable(l), .variable(r)): return l == r
        case let (.primitive(l), .primitive(r)): return l == r
        case let (.function(l1, l2), .function(r1, r2)): return l1 == r1 && l2 == r2
        default: return false
        }
    }
}

extension TypeScheme: Equatable, CustomStringConvertible {
    
    static func ==(lhs: TypeScheme, rhs: TypeScheme) -> Bool {
        switch (lhs, rhs) {
        case let (.type(tl), .type(tr)): return tl == tr
        case let (.forall(vl, tl), .forall(vr, tr)): return vl == vr && tl == tr
        default: return false
        }
    }
    
    private var rep: ([TypeVariable], Type) {
        switch self {
        case let .forall(x, rest):
            var (vars,type) = rest.rep
            vars.append(x)
            return (vars,type)
        case let .type(t):
            return ([], t)
        }
    }
    
    var description: String {
        let (vars,type) = rep
        guard !vars.isEmpty else { return type.description }
        let all = vars.joined(separator: " ")
        return "Ɐ \(all). \(type)"
    }
}

//: Some convenience extensions to write programs more functionally (closer to the paper)

extension Dictionary {
    func removing(_ key: Key) -> Dictionary {
        var copy = self
        copy[key] = nil
        return copy
    }
    
    func inserting(_ key: Key, _ value: Value) -> Dictionary {
        var copy = self
        copy[key] = value
        return copy
    }
}

extension Set {
    func removing(_ element: Element) -> Set {
        var result = self
        result.remove(element)
        return result
    }
}


//: A substition maps type variables to their types. Alternatively, we could represent this as a function `(TypeVariable) -> Type`, but making it a `Dictionary` allows for easy inspection.

typealias Substition = [TypeVariable: Type]

//: Assumptions are a mapping of variables to their types. For example, when we start inferencing types, this could contain the built-in functions and schemes.
typealias Assumptions = [ExpVariable: TypeScheme]

extension Type {
//: To apply a substitution, we recursively traverse a type. In case of variables, we check if the variable is in the map:
    mutating func apply(_ s: Substition) {
        switch self {
        case .primitive(_): ()
        case .variable(let v): self = s[v] ?? self
        case .function(var t1, var t2):
            t1.apply(s)
            t2.apply(s)
            self = .function(t1, t2)
        }
    }
    
    func applying(_ s: Substition) -> Type {
        var copy = self
        copy.apply(s)
        return copy
    }

//: A set of all the variables in a type:
    var variables: Set<TypeVariable> {
        switch self {
        case .primitive(_): return []
        case let .function(l, r): return l.variables.union(r.variables)
        case .variable(let v): return [v]
        }
    }

//: Given some assumptions, we can quantify a type scheme. This adds a `forall` for each type variable in the type, except the types that are in `a`.
    func quantify(_ a: Assumptions) -> TypeScheme {
        let freeVariables = variables.subtracting(a.freeVariables)
        return freeVariables.reduce(.type(self), { .forall($1, $0) })
    }
}

//: We can also apply a substition to a `TypeScheme`:
extension TypeScheme {
    func applying(_ s: Substition) -> TypeScheme {
        switch self {
        case .type(let t): return .type(t.applying(s))
        case .forall(let v, let t):
            return .forall(v, t.applying(s.removing(v)))
        }
    }
    
    mutating func apply(_ s: Substition) {
        self = applying(s)
    }
    
    var freeVariables: Set<TypeVariable> {
        switch self {
        case .type(let t): return t.variables
        case let .forall(x, t): return t.freeVariables.removing(x)
        }
    }
}

//: Finally, we can apply substitions to `Assumptions`:
extension Dictionary where Value == TypeScheme {
    mutating func apply(_ s: Substition) {
        for i in keys {
            self[i]!.apply(s) // todo mapValues or something?
        }
    }
    
    func applying(_ s: Substition) -> Dictionary {
        var copy = self
        copy.apply(s)
        return copy
    }
    
    var freeVariables: Set<TypeVariable> {
        return values.reduce([], { $0.union($1.freeVariables) })
    }
}

//: We can combine two substitions:
extension Dictionary where Key == TypeVariable, Value == Type {
    func unioned(_ other: Dictionary) -> Dictionary {
        return merging(other.mapValues { $0.applying(self) }, uniquingKeysWith: { x, _ in x }) // todo why can we ignore $1?
    }
}


extension String: Error { }

extension Literal {
    var primitiveType: PrimitiveType {
        switch self {
        case .int: return .int
        case .string: return .string
        }
    }
    var type: Type {
        return .primitive(primitiveType)
    }
}


//: The main inferencer state
struct State {
    
//: We need a fresh variable supply, and we'll use integers to make them
    private var freshVariableSupply = (0...).makeIterator()
    
//: Creating a type with a fresh variable
    mutating func freshType() -> Type {
        let num = freshVariableSupply.next()!
        return .variable("τ_\(num)")
    }

//: Unification tries to check if the types are "the same". For example, in case of two variables, it will return a substition that substitutes one variable for the other. Functions are unified recursively, and primitives should just be the same:
    func unify(_ l: Type, _ r: Type) throws -> Substition {
        switch (l, r) {
        case let (.function(l1, l2), .function(r1,r2)):
            let sub1 = try! unify(l1,r1)
            let sub2 = try! unify(l2.applying(sub1),r2.applying(sub1))
            return sub1.unioned(sub2)
        case let (.variable(name), _):
            return [name: r]
        case let (_, .variable(name)):
            return [name: l]
        case let (.primitive(p1), .primitive(p2)) where p1 == p2:
            return [:]
        default:
            throw "Cannot unify \(l) and \(r)"
        }
    }

//: Instantiating a type scheme (e.g. `forall a b . a -> b -> b`) to a type (with fresh variables, e.g. `a_0 -> b_0 -> b_0`)
    mutating func instantiate(_ t: TypeScheme) -> Type {
        switch t {
        case let .forall(v, t):
            let b = freshType()
            return instantiate(t.applying([v:b]))
        case .type(let t):
            return t
        }
    }

//: Public "interface": calculating a type scheme from an expression and some assumptions:
    mutating func scheme(e: Exp, a: Assumptions) throws -> TypeScheme {
        let (s,t) = try infer(e: e, a: a)
        return t.applying(s).quantify(a.applying(s))
    }

//: The actual type-inference method.
    mutating func infer(e: Exp, a: Assumptions) throws -> (Substition, Type) {
        switch e {
//: Inferring a variable is just looking it up in the context:
        case .variable(let s):
            let t = try a[s] !? "Cannot find variable \(s)"
            return ([:], instantiate(t))
//: Inferring a literal is assign its type:
        case .literal(let l):
            return ([:], l.type)
        case .app(let l, let r):
//: When inferring an application, we first infer the right-hand side (the parameter) and the left-hand side (the function).
            let (s1, t2) = try infer(e: r, a: a)
            let (s2, t1) = try infer(e: l, a: a.applying(s1))
//: Now we create a fresh type, and try to unify t1 (the left-hand side's type) with a new function type (taking `t2` and returning `beta`).
            let beta = freshType()
            let v = try unify(t1.applying(s2), .function(t2, beta))
            return (v.unioned(s2.unioned(s1)), beta.applying(v))
//: For function abstraction, we create fresh type, infer the body, and return a function type
        case .lam(let n, let s):
            let b = freshType()
            let (s,t1) = try infer(e: s, a: a.inserting(n,.type(b)))
            return (s, .function(b.applying(s), t1))
//: The most tricky one, a let binding.
        case let .let_(v, e1, e2):
//: We start by inferring the type of `e1`:
            let (s1, t1) = try infer(e: e1, a: a)
            let t1_1 = t1.applying(s1)
//: Now we create a new context with `v:t_1` in there. Note that the type is quantified, that is: this is where we introduce the `forall`s.
            let a1 = a.applying(s1)
            let aNew = a1.inserting(v, t1_1.quantify(a1))
//: We can then infer the type of e2 in our new context:
            let (s2, t2) = try infer(e: e2, a: aNew)
            return (s2.unioned(s1), t2)
        }
        
    }
    
    mutating func algorithmM(e: Exp, a: Assumptions, rho: Type) throws -> Substition {
        switch e {
        case .literal(let l):
            return try unify(rho, l.type)
        case .variable(let s):
            let t = try a[s] !? "Cannot find variable \(s)"
            return try unify(rho, instantiate(t))
        case .lam(let x, let e):
            let b1 = freshType()
            let b2 = freshType()
            let s1 = try unify(rho, .function(b1, b2))
            var newAssumptions = a.applying(s1)
            newAssumptions[x] = .type(b1.applying(s1))
            let s2 = try algorithmM(e: e, a: newAssumptions, rho: b2.applying(s1))
            return s2.unioned(s1)
        case .app(let e1, let e2):
            let b = freshType()
            let s1 = try algorithmM(e: e1, a: a, rho: .function(b, rho))
            let s2 = try algorithmM(e: e2, a: a.applying(s1), rho: b.applying(s1))
            return s2.unioned(s1)
        case let .let_(x, e1, e2):
            let beta = freshType()
            let s1 = try algorithmM(e: e1, a: a, rho: beta)
            var newAssumptions = a.applying(s1)
            newAssumptions[x] = beta.applying(s1).quantify(newAssumptions)
            let s2 = try algorithmM(e: e2, a: newAssumptions, rho: rho.applying(s1))
            return s2.unioned(s1)
        }
        
    }
    
    mutating func runM(e: Exp, a: Assumptions) throws -> TypeScheme {
        let rho = freshType()
        let s = try algorithmM(e: e, a: a, rho: rho)
        return rho.applying(s).quantify(a.applying(s))
    }
}

infix operator !?
func !?<A>(lhs: A?, rhs: @autoclosure () -> Error) throws -> A {
    guard let x = lhs else { throw rhs() }
    return x
}


extension Exp: ExpressibleByIntegerLiteral {
    init(integerLiteral value: Int) {
        self = .literal(.int(value))
    }
    
}

let id: Exp = .variable("id")
let idLam: Exp = .lam("x", .variable("x"))
let const: Exp = .lam("x", .lam("y", .variable("x")))
let apply: Exp = .lam("f", .lam("x", .app(.variable("f"), .variable("x"))))

let flip: Exp = .lam("f", .lam("x", .lam("y", .app(.app(.variable("f"), .variable("y")), .variable("x")))))

var x = State()
// let t = try! x.infer(e: .let_("main", flip, .variable("main")), a: ["id": .forall("a", .type(.function(.variable("a"), .primitive(.string))))])
let t = try! x.runM(e: .let_("main", flip, .variable("main")), a: ["id": .forall("a", .type(.function(.variable("a"), .primitive(.string))))])
print(t)