Skip to content

Instantly share code, notes, and snippets.

@chriseidhof
Last active February 3, 2018 23:11
Show Gist options
  • Save chriseidhof/148b32685a34096b7c57b3d5eaf641e5 to your computer and use it in GitHub Desktop.
Save chriseidhof/148b32685a34096b7c57b3d5eaf641e5 to your computer and use it in GitHub Desktop.
/*:
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment