Typed Expressions in Swift

TypedExpr.swift

// Variables just contain an integer. We can have a maximum of `Int.max` variables in our program.  ¯\_(ツ)_/¯

private struct Var {
    static var freshVarIx = 0
    let ix: Int
    init() {
        Var.freshVarIx+=1
        ix = Var.freshVarIx
    }
}

// Expressions are untyped. We only expose `TExpr` construction functions that allow us to build well typed terms.

private indirect enum Expr {
    case Literal(x: Any)
    case Variable(x: Var)
    case Abs(x: Var, f: Expr)
    case App(x: Expr, y: Expr)
}

// Evaluation. This is what a closure looks like: it contains a context (the variables it closes over), the variable we need to substitute, and the body of the closure.

private struct Closure {
    let context: [Int:Any]
    let variable: Var
    let body: Expr
}

// Because the public interface only allows us to construct well-typed terms, the ! and fatalError are safe here.
extension Expr {
    private func eval(context: [Int:Any]) -> Any {
        switch self {
        case .Literal(let x): return x
        case .Variable(let ix): return context[ix.ix]!
        case let .Abs(variable, body): return Closure(context: context, variable: variable, body: body)
        case let .App(lhs, rhs):
            let r = rhs.eval(context)
            guard let closure = lhs.eval(context) as? Closure else { fatalError("Expected a lambda") }
            var newContext = closure.context
            newContext[closure.variable.ix] = r
            return closure.body.eval(newContext)
        }
    }
}


// The public facing type. Note that it can only be constructed using the functions below.

public struct TExpr<T> {
    private let expr: Expr
}

// A simple literal

public func literal<A>(value: A) -> TExpr<A> {
    return TExpr(expr: Expr.Literal(x: value))
}

// Creating lambda's
public func lambda<A,B>(f: TExpr<A> -> TExpr<B>) -> TExpr<A -> B> {
    let var_ = Var()
    let variable = TExpr<A>(expr: Expr.Variable(x: var_))
    return TExpr(expr: .Abs(x: var_, f: f(variable).expr))
}

// An apply function. Maybe we could use subscript syntax instead?
public func apply<A,B>(f: TExpr<A -> B>, _ arg: TExpr<A>) -> TExpr<B> {
    return TExpr(expr: .App(x: f.expr, y: arg.expr))
}

// Evaluating the expr. Again, because it's well-typed, we know we can safely cast. The exception: if T is a function type, we will crash.
extension TExpr {
    public func eval() -> T {
        return expr.eval([:]) as! T
    }
}

// The first lambda forgets the first parameter, the second lambda forgets the second parameter. The type annotations are optional: Swift will infer it for us.

let hello: TExpr<String> = apply(apply((lambda { x in lambda { y in y } }), literal(1)), literal("Hello"))
let one:   TExpr<Int>    = apply(apply((lambda { x in lambda { y in x } }), literal(1)), literal("Hello"))

print(hello.eval())
print(one.eval())