Simple bidirectional type checker in Swift

bidirectional_typechecker.swift

// NOTA BENE:
//
// The supports a _very_ tiny AST, with only a few possible
// types and just a few more expressions.
//
// Modelled on the excellent blog post here:
//
//    https://jimmyhmiller.com/easiest-way-to-build-type-checker

indirect enum Type: Equatable, CustomStringConvertible {
    case number
    case string
    case function(argTypes: [Type], returnType: Type)

    var description: String {
        switch self {
        case .number: return "Number"
        case .string: return "String"
        case .function(let argTypes, let returnType):
            let allArgTypes = argTypes.map(\.description).joined(separator: ", ")
            return "Function (\(allArgTypes)) -> \(returnType)"
        }
    }
}

indirect enum Expression {
    case number(value: Double)
    case string(value: String)
    case variable(name: String)
    case function(args: [(String, Type)], body: Expression)
    case call(function: Expression, argVals: [Expression])
    case `let`(name: String, value: Expression, type: Type?)
    case block(statements: [Expression], returnVal: Expression)
    case add(left: Expression, right: Expression)
}

typealias Context = Dictionary<String, Type>

enum TypeError: Error {
    case cannotCallNonfunction
    case cannotInferType
    case cannotInferFunctionType
    case typeMismatch
    case unboundVariable
    case incorrectArity
}

func infer(context: inout Context,
           expression: Expression) throws -> Type {
    switch expression {
    case .number:
        return .number
    case .string:
        return .string
    case .variable(name: let name):
        guard let type = context[name] else {
            throw TypeError.unboundVariable
        }

        return type
    case .add(left: let left, right: let right):
        if try (check(context: &context, expression: left, expectedType: .number) &&
            check(context: &context, expression: right, expectedType: .number)) {
            return .number
        }

        throw TypeError.typeMismatch
    case .call(function: let funcExpr, argVals: let argVals):
        let funcType = try infer(context: &context, expression: funcExpr)
        guard case .function(argTypes: let argTypes,
                             returnType: let returnType) = funcType else {
            throw TypeError.cannotCallNonfunction
        }

        for (idx, argVal) in argVals.enumerated() {
            if !(try check(context: &context,
                           expression: argVal,
                           expectedType: argTypes[idx])) {
                throw TypeError.typeMismatch
            }
        }

        return returnType
    case .function(args: let args, body: let bodyExpr):
        var newCtx: Context = [:]
        var argTypes: [Type] = []
        for (argName, argType) in args {
            newCtx[argName] = argType
            argTypes.append(argType)
        }

        let returnType = try infer(context: &newCtx, expression: bodyExpr)

        return .function(argTypes: argTypes, returnType: returnType)
    case .let(name: let name, value: let valExpr, type: let maybeType):
        let valType = try infer(context: &context, expression: valExpr)
        if let type = maybeType {
            if valType != type {
                throw TypeError.typeMismatch
            }
        }

        context[name] = valType
        return valType
    case .block(statements: let exprs, returnVal: let returnVal):
        var newCtx: Context = [:]
        for expr in exprs {
            try infer(context: &newCtx, expression: expr)
        }

        return try infer(context: &context, expression: returnVal)
    }
}

func check(context: inout Context,
           expression: Expression,
           expectedType: Type) throws -> Bool {
    switch expression {
    case .function(args: let args, body: let bodyExpr):
        guard case .function(argTypes: let argTypes,
                             returnType: let returnType) = expectedType else {
            return false
        }

        if argTypes.count != args.count {
            throw TypeError.incorrectArity
        }

        var newCtx: Context = [:]
        for (idx, (argName, argType)) in args.enumerated() {
            if argType != argTypes[idx] {
                throw TypeError.typeMismatch
            }
            newCtx[argName] = argTypes[idx]
        }

        return try check(context: &newCtx, expression: bodyExpr, expectedType: returnType)
    case .block(statements: let exprs, returnVal: let returnVal):
        var newCtx: Context = [:]
        for expr in exprs {
            try infer(context: &newCtx, expression: expr)
        }

        return try check(context: &newCtx, expression: returnVal, expectedType: expectedType)
    default:
        let actualType = try infer(context: &context, expression: expression)
        return actualType == expectedType
    }
}

var context: Context = [:]

let testExpr: Expression =
    .function(args: [("a", .number), ("b", .number)],
              body: .block(statements: [],
                           returnVal: .add(left: .variable(name: "a"),
                                           right: .variable(name: "b"))))

let testType = try infer(context: &context, expression: testExpr2)
print(testType)