Monomorphic type system in F#

Monomorphism.fs

module TypeSystem

type HashTable<'T, 'V> = System.Collections.Generic.Dictionary<'T, 'V>

let varId = ref 0
let nextId() =
    let r = !varId
    varId := r + 1
    r

type Ty
    = Unknown
    | TUnit
    | TBool
    | TInt
    | TFloat
    | TVar  of int
    | TAbs  of Ty * Ty

type Exp
    = EUnit
    | EBool     of bool
    | EInt      of int
    | EFloat    of float
    | ELet      of string * Exp * Exp
    | EAbs      of string * Exp
    | EVar      of string
    | EIf       of Exp * Exp * Exp
    | EApp      of Exp * Exp

type TypeEnv
    = { hm              : HashTable<int, Ty>
        mutable varId   : int }
with
    static member empty() = { hm = HashTable(); varId = 0 }
    member x.newType() =
        let vid = x.varId
        x.varId <- vid + 1
        x.hm.Add(vid, Unknown)
        TVar vid
    
    member x.Item with get i =
                    let rec getTy i =
                        match x.hm.[i] with
                        | TVar i'   -> getTy i'   // recursively find the item
                        | y         -> y
                    getTy i
                  and  set i v = x.hm.[i] <- v

    member x.contains k = x.hm.ContainsKey k        
    
let printVars (tyEnv: TypeEnv) (varEnv: HashTable<string, Ty>) =

    let rec printTy (ty: Ty) : string =
        match ty with
        | Unknown       -> "Unknown"
        | TUnit         -> "Unit"
        | TBool         -> "Bool"
        | TInt          -> "Int"
        | TFloat        -> "Float"
        | TVar  i       -> printTy tyEnv.[i]
        | TAbs(p, r)    -> sprintf "%s -> %s" (p |> printTy) (r |> printTy)

    varEnv
    |> Seq.iter(fun kv -> printfn "%s : %s" kv.Key (kv.Value |> printTy))

let rec varAlphaConvert (varMap: Map<string, string>) (exp: Exp) : Exp =
    match exp with
    | EUnit     
    | EBool _   
    | EInt  _   
    | EFloat _  -> exp
    | EVar x    -> varMap.[x] |> EVar

    | ELet (x, b, r) ->
        let nb = varAlphaConvert varMap b
        let vName = sprintf "%s.%d" x (nextId())
        let env = varMap.Add(x, vName)
        let nr = varAlphaConvert env r
        ELet (vName, nb, nr)

    | EAbs (x, b) ->
        let vName = sprintf "%s.%d" x (nextId())
        let env = varMap.Add(x, vName)
        let nb = varAlphaConvert env b
        EAbs (vName, nb)

    | EIf (b, t, e) ->
        let nb = varAlphaConvert varMap b
        let nt = varAlphaConvert varMap t
        let ne = varAlphaConvert varMap e
        EIf (nb, nt, ne)

    | EApp (f, p) ->
        let nf = varAlphaConvert varMap f
        let np = varAlphaConvert varMap p
        EApp (nf, np)

let rec unify (tyEnv: TypeEnv) (a: Ty) (b: Ty) =
    match a, b with
    | a, b when a = b -> ()
    | TAbs (a0, a1), TAbs (b0, b1) ->
        unify tyEnv a0 b0
        unify tyEnv a1 b1

    | TVar ta, TVar tb ->
        unify tyEnv tyEnv.[ta] tyEnv.[tb]
        tyEnv.[ta] <- b

    | TVar t, x
    | x, TVar t ->
        match tyEnv.[t] with
        | s when s = x  -> ()
        | Unknown       -> tyEnv.[t] <- x
        | _             -> failwith (sprintf "could not unify %d with %A" t x)

    | _ -> failwith (sprintf "unable to unify %A with %A" a b)

let rec infer (tyEnv: TypeEnv) (varEnv: HashTable<string, Ty>) (exp: Exp) : Ty =
    match exp with
    | EUnit     -> TUnit
    | EBool _   -> TBool
    | EInt  _   -> TInt
    | EFloat _  -> TFloat
    | EVar x    -> varEnv.[x]
    | ELet (x, b, r) ->
        let bTy = infer tyEnv varEnv b
        varEnv.[x] <- bTy
        infer tyEnv varEnv r

    | EAbs (x, b) ->
        let ty = tyEnv.newType()
        varEnv.Add(x, ty)
        let r = infer tyEnv varEnv b
        TAbs(ty, r)

    | EIf (b, t, e) ->
        let bTy = infer tyEnv varEnv b
        unify tyEnv bTy TBool
        let tTy = infer tyEnv varEnv t
        let eTy = infer tyEnv varEnv e
        unify tyEnv tTy eTy
        tTy

    | EApp (f, p) ->
        let fTy = infer tyEnv varEnv f
        let pTy = infer tyEnv varEnv p
        let tv = tyEnv.newType()
        unify tyEnv fTy (TAbs(tv, pTy))
        pTy

let test() =
    let doit(exp: Exp) =
        let tyEnv = TypeEnv.empty()
        let varEnv = HashTable()

        let nExp
            = exp
            |> varAlphaConvert Map.empty
        
        infer tyEnv varEnv nExp
        |> printfn "%A"
        varEnv
        |> printVars tyEnv

    let test1() =
        ELet("x", EInt 1, EUnit)
        |> doit

    let test2() =
        ELet("f", EAbs("x", EVar "x"), EApp(EVar "f", EInt 1))
        |> doit

    let test3() =
        ELet("f", EAbs("x", EVar "x"), 
            ELet("g", EAbs("x",
                        EApp(EVar "f", EVar "x")),
                     EApp(EVar "g", EInt 1)))
        |> doit

    test1()
    test2()
    test3()