Abstract Binding Trees ~(http://www.cs.cmu.edu/~rjsimmon/15312-s14/hws/hw1update2-handout.pdf) in Swift

ABT.swift

/// An operator is given by a list of arities, where each element indicates the number of
/// variables bound by the operator at that position and the length of the list determines the
/// number of variables the operator accepts.  The full generalization of arity is called
/// the operator's "valence".
///
/// For example, if I have a little calculator language with let-bindings, its operators
/// would look like this:
///
/// 
///		enum CalcOps : Operator {
///			case Plus
///			case Minus
///			case Mult
///			case Div
///			case Let
///			
///			var arity : [UInt] {
///				switch self {
///				case .Plus:
///					return [0, 0] // x + y binds no variables
///				case .Minus:
///					return [0, 0] // x - y binds no variables
///				case .Mult:
///					return [0, 0] // x * y binds no variables
///				case .Div:
///					return [0, 0] // x / y binds no variables
///				case .Let:
///					return [0, 1] // let x = e binds 1 variable (x).
///				}
///			}
///			
///			var description : String {
///				switch self {
///				case .Plus:
///					return "+"
///				case .Minus:
///					return "-"
///				case .Mult:
///					return "*"
///				case .Div:
///					return "/"
///				case .Let:
///					return "let"
///				}
///			}
///		}
protocol Operator : Equatable, CustomStringConvertible {
	var arity : [UInt] { get }
}

/// A variable is a fixed term.  ABTs require variables have a "freshness guarantee".  That is,
/// newly created variables are wholly distinct from previously created variables.
protocol Variable : Equatable, Comparable, Hashable, CustomStringConvertible {
	var fresh : Self { get }
	var copy : Self { get }
}

/// An ABT capturing the design spec in the homework assignment from 
/// CMU (http://www.cs.cmu.edu/~rjsimmon/15312-s14/hws/hw1update2-handout.pdf)
/// about ABTs.
///
/// The original concept comes from Per Martin-Lof's "Method of Arities"
/// (http://www.cse.chalmers.se/research/group/logic/book/book.pdf).
protocol ABT {
	associatedtype V : Variable
	associatedtype O : Operator
	
	/// Fold a View back into an ABT.
	static func into(v : View<V, O>) -> Self
	/// Unfold an ABT into a View.
	func outOf() -> View<V, O>
}

/// A "single-step unfolding" of the ABT.  The View of the ABT acts as a mediator between the ABT and
/// "the outside world", which in this case is Swift.  On the one hand, forcing the user to go through
/// the View means we get to perform validation of invariants.  On the other, it allows a much, much better
/// way to pattern match on the useful parts of an ABT.
enum View<V : Variable, O : Operator> {
	/// A variable. Free, never bound.
	case Var(V)
	/// An abstraction, given by the bound variable and the rest of the term.
	case Abstract(V, NamelessABT<V, O>)
	/// An application, given by an operator of proper arity applied to a list of terms.
	case Apply(O, [NamelessABT<V, O>])
}

/// A Locally Nameless implementation of ABTs.
///
/// If an AST is a tree of variables applied to operators whose arguments are classified by sorts, 
/// then an ABT is an AST that is aware of binding.  We augment the traditional `Apply`-only AST with
/// variables (here, because it is *locally* nameless we use whatever the user wants for free variables
/// but use de-bruijn indexing for bound terms) and a way to introduce bound variables in a lambda
/// abstraction yielding an unfolding of the ABT in which the given variable is bound.  As before,
/// variables are given significance only through the process of substitution.  Only this time, we can
/// view substitution through the lens of a set of (possibly disjoint) binding scopes.
indirect enum NamelessABT<V : Variable, O : Operator> : ABT, CustomStringConvertible {
	/// Ideally, these cases are all private.
	case Free(V)
	case Bound(UInt)
	case Abstract(V, NamelessABT<V, O>)
	case Apply(O, [NamelessABT<V, O>])
	
	/// Convert a `View` into a term.
	static func into(v : View<V, O>) -> NamelessABT<V, O> {
		switch v {
		/// Transport free variables to free variables in the ABT.
		case .Var(let v):
			return .Free(v)
		/// Convert an abstraction binding a variable into the ABT's flavor of abstraction making sure
		/// the previously free variable is properly bound in the body of the ABT.
		case .Abstract(let v, let abt):
			return .Abstract(v.copy, abt.byBindingVariable(v))
		/// Converts an application 
		case .Apply(let o, let es):
			let ar = o.arity
			// Arity check: Does the operator's arity match the number of arguments we actually have on hand?
			assert(ar.count == es.count)
			// Moreover, for each argument of the operator, does its arity match with what we would expect?
			assert(zip(es, ar).reduce(true) { b, t in b && t.0.matchArity(t.1) })
			return .Apply(o, es)
		}
	}
	
	/// from https://github.com/freebroccolo/rust-abt
	private func matchArity(n : UInt) -> Bool {
		switch self {
		case .Abstract(_, _):
			return n != 0
		default:
			return n == 0
		}
	}
	
	/// Convert a term into a `View`.
	func outOf() -> View<V, O> {
		switch self {
		/// Transport free variables to free variables in the View.
		case .Free(let v):
			return .Var(v)
		/// Bound variables have no place in the View.  Mostly because the only way that
		/// we can get here is if something has gone terribly wrong because you can't
		/// bind variables arbitrarily, only with the `Abstract` case.
		case .Bound(_):
			fatalError()
		/// Converts an abstraction involving a bound variable to one involving a free
		/// variable and the rest of the ABT where the new free variable has been added.
		case .Abstract(let x, let e):
			let v = x.copy
			return .Abstract(v.copy, e.byFreeingVariable(v))
		/// Don't have to do any work for applications.
		case .Apply(let o, let es):
			return .Apply(o, es)
		}
	}
	
	/// Attempts to find the variable `v` and bind it.
	///
	/// Imprisonment.
	func byBindingVariable(v : V, lvl : UInt = 0) -> NamelessABT<V, O> {
		switch self {
		// For free variables:
		case .Free(let fv):
			if v == fv {
				// if we find what we're looking for, then we have to bind it at the current
				// index level.
				return .Bound(lvl)
			} else {
				// else we leave it alone.
				return .Free(fv)
			}
		// If it's already bound there's nothign to do.
		case .Bound(let bv):
			return .Bound(bv)
		// Recurse into the body of the abstraction trying to bind the variable in the next level.
		case .Abstract(let x, let e):
			return .Abstract(x, e.byBindingVariable(v, lvl: lvl + 1))
		// Recurse into the bodies of the operator's arguments trying to bind the variable at the current level.
		case .Apply(let o, let es):
			let res = es.map { $0.byBindingVariable(v.copy, lvl: lvl) }
			return .Apply(o, res)
		}
	}
	
	/// Inserts a new free variable starting at position `n`, renaming bound variables as needed.
	///
	/// Jailbreak!
	func byFreeingVariable(v : V, lvl : UInt = 0) -> NamelessABT<V, O> {
		switch self {
		// Leave free variables alone.
		case .Free(let fv):
			return .Free(fv)
		// For bound variables:
		case .Bound(let m):
			if m == lvl {
				// if we find what we're looking for, then we free the variable.
				return .Free(v)
			} else {
				// else we leave it alone.
				return .Bound(m)
			}
		// Recurse into the body of the abstraction trying to free the variable in the next level.
		case .Abstract(let x, let e):
			return .Abstract(x, e.byFreeingVariable(v, lvl: lvl + 1))
			// Recurse into the bodies of the operator's arguments trying to free the variable at the current level.
		case .Apply(let o, let es):
			let res = es.map({ e in e.byFreeingVariable(v.copy, lvl: lvl) })
			return .Apply(o, res)
		}
	}

	/// Returns the result of substituting the in the current ABT for the given variable and applying the
	/// expression to the result.
	func bySubstituting(tm: NamelessABT<V, O>, for vari : V) -> NamelessABT<V, O> {
		switch self.outOf() {
		/// HACK: I have to rework variable equality in the named case.
		case let .Var(v) where vari.description == v.description:
			return tm // Perform the sub
		/// Search the body of the binding for things to sub for.
		case let .Abstract(v, expb):
			return abstract(v, expb.bySubstituting(tm, for: vari))
		/// Search the operands for things to sub for.
		case let .Apply(o, ts):
			return apply(o, ts.map { $0.bySubstituting(tm, for: vari) })
		default:
			return self
		}
	}
	
	/// Returns a lift containing the free variables in the ABT with each occuring at most once.
	var freevars : Set<V> {
		func rec(exp : NamelessABT<V, O>) -> [V] {
			switch exp {
			/// Trivial
			case let .Free(v):
				return [v]
			case .Bound(_):
				return []
			/// The free variables in an abstraction are all the free variables in the body
			/// minus the variable we captured.
			case let .Abstract(v, expi):
				return rec(expi).filter({ $0 != v })
			/// The free variables in an operator application is the union of free variables in all
			/// the operands.
			case let .Apply(_, exps):
				return exps.flatMap(rec)
			}
		}
		return Set(rec(self))
	}
	
	var description : String {
		switch self.outOf() {
		case .Var(let v):
			return v.description
		case .Abstract(let v, let e):
			return v.description + "." + e.description
		case .Apply(let o, let es) where es.isEmpty:
			return o.description
		case .Apply(let o, let es):
			let ss = es.map { $0.description + ", " }.reduce("", combine: +)
			return o.description + "(" + ss[ss.startIndex..<ss.endIndex.predecessor().predecessor()] + ")"
		}
	}
}


/// Helpers:
///
/// The example below is an implementation of the Lambda calculus, it only makes sense to have lambda terms
/// be easy to construct.  

func abstract<V : Variable, O : Operator>(x : V, _ e : NamelessABT<V, O>) -> NamelessABT<V, O> {
	return NamelessABT<V, O>.into(.Abstract(x, e))
}

func apply<V : Variable, O : Operator>(o : O, _ es : [NamelessABT<V, O>]) -> NamelessABT<V, O> {
	return NamelessABT<V, O>.into(.Apply(o, es))
}

func variable<V : Variable, O : Operator>(x : V) -> NamelessABT<V, O> {
	return NamelessABT<V, O>.into(.Var(x))
}

extension NamelessABT : Equatable { }

/// Alpha-equivalence
///
/// Respects NamelessABT<V, O>.into(abt.outOf()) == abt
func == <V : Variable, O : Operator>(l : NamelessABT<V, O>, r : NamelessABT<V, O>) -> Bool {
	switch (l, r) {
	case let (.Free(i), .Free(j)):
		return i == j
	case let (.Bound(x), .Bound(y)):
		return x == y
	case let (.Abstract(_, t1), .Abstract(_, t2)):
		return t1 == t2
	case let (.Apply(f1, _), .Apply(f2, _)):
		return f1 == f2
	default:
		return false
	}
}

Helpers.swift

infix operator <^> {
  associativity left
  precedence 140
}

infix operator <*> {
  associativity left
  precedence 140
}

public func <^> <A, B>(f: A -> B, a: A?) -> B? {
	return a.map(f)
}

public func <^> <A, B>(f: A -> B, a: [A]) -> [B] {
	return a.map(f)
}

public func <*> <A, B>(f: [(A -> B)], a: [A]) -> [B] {
	var re = [B]()
	for g in f {
		for h in a {
			re.append(g(h))
		}
	}
	return re
}

public func <*> <A, B>(f: (A -> B)?, a: A?) -> B? {
	if f != nil && a != nil {
		return (f!(a!))
	} else {
		return .None
	}
}

LambdaCalculus.swift

// The first-order dependently typed lambda calculus.
enum LambdaCalculus : Operator {
	case Lam
	case App
	case Pi
	case Unit
	case Axiom
	
	var arity : [UInt] {
		switch self {
		case .Lam: // λx.e binds x
			return [1]
		case .App: // (λx.e)y binds nothing
			return [0, 0]
		case .Pi: // Πx.e(x) binds x only in the body e(x)
			return [0, 1]
		case .Unit: // constant
			return []
		case .Axiom: // constant
			return []
		}
	}
	
	var description : String {
		switch self {
		case .Lam:
			return "λ"
		case .App:
			return "ap"
		case .Pi:
			return "pi"
		case .Unit:
			return "unit"
		case .Axiom:
			return "<>"
		}
	}
}

typealias LambdaTerm = NamelessABT<Named, LambdaCalculus>

let x = "x"
let y = "y"
let identity : LambdaTerm = apply(.Lam, [abstract(Named(x), variable(Named(x)))])
let app : LambdaTerm = apply(.App, [identity, identity])
identity.bySubstituting(identity, for: Named(x))

LambdaTerm.into(identity.outOf()) == identity

Naming.swift

/// Nameless variables ["de Bruijn indices"]
struct Nameless : Variable {
	private let unVariable : UInt
	
	init() {
		self.unVariable = Nameless.freshVariable()
	}
	
	private init(copying : UInt) {
		self.unVariable = copying
	}
	
	var fresh : Nameless {
		return Nameless()
	}
	
	var copy : Nameless {
		return Nameless(copying: self.unVariable)
	}
	
	var hashValue : Int {
		return Int(self.unVariable)
	}
	
	var description : String {
		return String(self.unVariable)
	}
	
	static var _varPool : UInt = 0
	
	static func freshVariable() -> UInt {
		_varPool += 1;
		return _varPool
	}
}

func <=(lhs : Nameless, rhs : Nameless) -> Bool {
	return lhs.unVariable <= rhs.unVariable
}

func >=(lhs : Nameless, rhs : Nameless) -> Bool {
	return lhs.unVariable >= rhs.unVariable
}

func >(lhs : Nameless, rhs : Nameless) -> Bool {
	return lhs.unVariable > rhs.unVariable
}

func <(lhs : Nameless, rhs : Nameless) -> Bool {
	return lhs.unVariable < rhs.unVariable
}

func ==(lhs : Nameless, rhs : Nameless) -> Bool {
	return lhs.unVariable == rhs.unVariable
}

/// Named variables; Reference equality is needed because strings have value semantics.
final class Named : Variable, StringLiteralConvertible {
	private let unVariable : String
	typealias StringLiteralType = String
	
	init(unicodeScalarLiteral value: UnicodeScalar) {
		self.unVariable = String(value)
	}
	
	init(extendedGraphemeClusterLiteral value: String) {
		self.unVariable = value
	}
	
	init(stringLiteral value: String) {
		self.unVariable = value
	}
	
	init(_ copying : String) {
		self.unVariable = copying
	}
	
	var fresh : Named {
		return Named(self.unVariable + "'")
	}
	
	var copy : Named {
		return Named(self.unVariable)
	}
	
	var hashValue : Int {
		return self.unVariable.hashValue
	}
	
	var description : String {
		return self.unVariable
	}
}

func <=(lhs : Named, rhs : Named) -> Bool {
	return lhs.unVariable <= rhs.unVariable
}

func >=(lhs : Named, rhs : Named) -> Bool {
	return lhs.unVariable >= rhs.unVariable
}

func >(lhs : Named, rhs : Named) -> Bool {
	return lhs.unVariable > rhs.unVariable
}

func <(lhs : Named, rhs : Named) -> Bool {
	return lhs.unVariable < rhs.unVariable
}

func ==(lhs : Named, rhs : Named) -> Bool {
	return lhs === rhs
}

TypeCheck.swift

typealias Context = Dictionary<Named, LambdaTerm>

func eval(tm : LambdaTerm) -> LambdaTerm {
	func step(tm : LambdaTerm) -> Optional<LambdaTerm> {
		switch tm.outOf() {
		case let .Apply(.App, es):
			let m = es.first!
			switch m.outOf() {
			case let .Apply(.Lam, xe):
				return xe.first!
			default:
				return (applyPi <^> step(m) <*> es[1]) ?? (applyPi <^> .Some(m) <*> step(es[1]))
			}
		case let .Apply(.Pi, tys):
			let a = tys[0]
			let b = tys[1]
			return pi <^> step(a) <*> step(b)
		default:
			return nil
		}
	}
	return step(tm) ?? tm
}

func infer(inout gamma : Context, tm : LambdaTerm) -> LambdaTerm {
	switch tm.outOf() {
	case .Var(let v):
		if let ty = gamma.filter({ $0.0.unVariable == v.unVariable }).first {
			return eval(ty.1)
		}
		fatalError("Ill-scoped variable")
	case let .Apply(.App, es):
		let m = es.first!
		let n = es[1]
		switch infer(&gamma, tm: m).outOf() {
		case let .Apply(.Pi, tys):
			let a = tys[0]
			let b = tys[1]
			check(&gamma, tm: n, ty: a)
			return eval(b)
		default:
			fatalError("Expected Pi type for lambda abstraction")
		}
	case let .Abstract(m, n): // YOLO
		gamma[m] = unit()
		check(&gamma, tm: variable(m), ty: unit())
		return .Abstract(m, infer(&gamma, tm: eval(n)))
	default:
		fatalError("Only first-order neutral terms may be inferred")
	}
}

func check(inout gamma : Context, tm : LambdaTerm, ty : LambdaTerm) {
	let ntm = eval(tm)
	let nty = eval(ty)
	switch (ntm.outOf(), nty.outOf()) {
	case let (.Apply(.Lam, es), .Apply(.Pi, tys)):
		switch es.first! {
		case .Abstract(let x, _):
			let z = x.fresh
			_ = es.map({ (ez) -> () in
				gamma[z] = tys.first!
				check(&gamma, tm: ez, ty: tys[1])
				return ()
				
			})
			return ()
		default:
			fatalError("Ill-scoped variable")
		}
	case (.Apply(.Axiom, _), .Apply(.Axiom, _)):
		return ()
	default:
		let ty2 = infer(&gamma, tm: tm)
		if ty2 == nty {
			return ()
		} else {
			fatalError("Type error")
		}
	}
}

func lam(e : LambdaTerm) -> LambdaTerm {
	return .Apply(.Lam, [e])
}

func axiom() -> LambdaTerm {
	return .Apply(.Axiom, [])
}

func unit() -> LambdaTerm {
	return .Apply(.Unit, [])
}

func applyPi(a : LambdaTerm) -> LambdaTerm -> LambdaTerm {
	return { b in .Apply(.App, [a, b]) }
}

func pi(a : LambdaTerm) -> LambdaTerm -> LambdaTerm {
	return { b in .Apply(.Pi, [a, b]) }
}

var context = Context()
let vx = Named(x)
check(&context, tm: lam(abstract(vx, variable(Named(x)))), ty: pi(unit())(abstract(vx, unit())))