Algorithm W - or Hindley-Milner polymorphic type inference - in F#

AlgorithmW.fs

module AlgorithmW

// HINDLEY-MILNER TYPE INFERENCE
// Based on http://catamorph.de/documents/AlgorithmW.pdf

type Lit =
    | LInt of int
    | LBool of bool

type Exp =
    | EVariable of string
    | ELiteral of Lit
    | EApplication of Function : Exp * Argument : Exp
    | EFunction of ParameterName : string * Body : Exp
    | ELet of Name : string * Value : Exp * Body : Exp

type Type =
    | TInt
    | TBool
    | TFree of string
    | TFunction of ParameterType : Type * ReturnType : Type with
    member this.FreeTypeVariables =
        match this with
        | TInt
        | TBool -> Set.empty
        | TFree n -> Set.singleton n
        | TFunction (t1, t2) ->
            let v1 = t1.FreeTypeVariables
            let v2 = t2.FreeTypeVariables
            Set.union v1 v2
    member this.Apply (ts : TypeSubst) =
        match this with
        | TFree n ->
            match ts.TryFind n with
            | Some t -> t
            | None -> TFree n
        | TFunction (t1, t2) ->
            TFunction (t1.Apply ts, t2.Apply ts)
        | _ -> this

and TypeSubst = Map<string, Type>

/// A type like 'a -> 'b
type Polytype =
    {
        /// The type variables 'a, 'b, etc
        Variables : string list

        /// The type 'a -> 'a (?)
        Type : Type
    } with
    member this.FreeTypeVariables =
        this.Type.FreeTypeVariables - (Set.ofList this.Variables)
    member this.Apply (s : TypeSubst) =
        {
            Variables = this.Variables
            Type =
                let newSubst = List.fold (fun ns i -> Map.remove i ns) s this.Variables
                this.Type.Apply newSubst
        }

type TypeEnv =
    {
        Schemes : Map<string, Polytype>
    } with
    member this.Remove (var : string) =
        {
            Schemes = this.Schemes.Remove var
        }
    member this.FreeTypeVariables =
        Seq.fold (fun v (nts : System.Collections.Generic.KeyValuePair<string, Polytype>) ->
            Set.union v nts.Value.FreeTypeVariables) Set.empty this.Schemes
    member this.Apply (ts : TypeSubst) =
        { Schemes = this.Schemes |> Map.map (fun k v -> v.Apply ts) }


/// Missing in F# 3.1
/// Merges s1 over s2
let unionMap s1 s2 = Map.fold (fun acc key value -> Map.add key value acc) s2 s1

/// A Map with just one entry
let singletonMap k v = Map.ofSeq [k, v]

/// Apply `s1` to `s2` then merge the results
let composeSubst (s1 : TypeSubst) (s2 : TypeSubst) : TypeSubst =
    let ns2 = s2 |> Map.map (fun _ v -> v.Apply s1)
    unionMap ns2 s1

/// Abstracts a type over all type variables that are not free
let generalize (env : TypeEnv) (t : Type) : Polytype =
    {
        Variables = List.ofSeq (t.FreeTypeVariables - env.FreeTypeVariables)
        Type = t
    }

/// Generates a new type variable. STATEFUL.
let newTyVar =
    let nextIndex = ref 1
    fun n ->
        let nn = sprintf "%s%d" n !nextIndex
        nextIndex := !nextIndex + 1
        TFree nn

/// Replace all bound type variables with fresh variables
let instantiate (ts : Polytype) : Type =
    let nvars = List.map (fun _ -> newTyVar "a") ts.Variables
    let s = Map.ofSeq (Seq.zip ts.Variables nvars)
    let applied = ts.Type.Apply s
    applied

/// Bind a name to a Type and return that binding
let varBind (u : string) (t : Type) : TypeSubst =
    match t with
    | TFree u -> Map.empty
    | _ when t.FreeTypeVariables.Contains u ->
        failwith "Occur check fails: %s vs %A" u t
    | _ -> singletonMap u t

let rec unify (t1 : Type) (t2 : Type) : TypeSubst =
    match t1, t2 with
    | TFunction (l1, r1), TFunction (l2, r2) ->
        let s1 = unify l1 l2
        let s2 = unify (r1.Apply s1) (r2.Apply s1)
        composeSubst s1 s2
    | TFree u, t -> varBind u t
    | t, TFree u -> varBind u t
    | TInt, TInt -> Map.empty
    | TBool, TBool -> Map.empty
    | _ -> failwith "Types do not unify: %A vs %A" t1 t2

/// Type inference with pending substitutions
let rec ti (env : TypeEnv) (e : Exp) : TypeSubst * Type =
    match e with
    | EVariable n ->
        match env.Schemes.TryFind n with
        | None -> failwithf "Unbound variable: %s" n
        | Some sigma ->
            //let t = instantiate sigma
            Map.empty, sigma.Type
    | ELiteral l ->
        match l with
        | LInt _ -> Map.empty, TInt
        | LBool _ -> Map.empty, TBool
    | EFunction (n, e) ->
        let tv = newTyVar "a"
        let env2 : TypeEnv =
            let ts = { Type = tv; Variables = [] }
            { Schemes = unionMap (singletonMap n ts) env.Schemes }
        let s1, t1 = ti env2 e
        s1, TFunction (tv.Apply s1, t1)
    | EApplication (e1, e2) ->
        let tv = newTyVar "a"
        let s1, t1 = ti env e1
        let s2, t2 = ti (env.Apply s1) e2
        let s3 = unify (t1.Apply s2) (TFunction (t2, tv))
        List.fold composeSubst Map.empty [s3; s2; s1], tv.Apply s3
    | ELet (x, e1, e2) ->
        let s1, t1 = ti env e1 // get type of e1
        let tp = generalize (env.Apply s1) t1 // 
        let env2 = { Schemes = Map.add x tp env.Schemes }
        let s2, t2 = ti (env2.Apply s1) e2
        composeSubst s1 s2, t2

/// Type inference with all substitutions applied
let typeInference (env : Map<string, Polytype>) (e : Exp) =
    let s, t = ti { Schemes = env } e
    t.Apply s

/// Test this puppy
let test (e : Exp) =
    try
        let t = typeInference Map.empty e
        printfn "%A :: %A" e t
    with ex -> printfn "ERROR %O" ex


[<EntryPoint>]
let main argv = 
    [
        ELet ("id", EFunction ("x", EVariable "x"), EVariable "id")
        ELet ("id", ELiteral (LInt 42), EVariable "id")
    ]
    |> Seq.iter test
    0