solver.fs

exception TypeError of string

// Some useful extensions to F#'s map module
module Map =
    let keys m = m |> Map.toSeq |> Seq.map fst |> Set.ofSeq

    let fromKeys f keys =
        Seq.fold (fun accum k -> Map.add k (f k) accum) Map.empty keys

    let singleton key value =
        Map.add key value Map.empty

type Subst = Map<TyVar, TyExpr>
let idSubst : Subst = Map.empty

let varsubst theta (v : TyVar) : TyExpr =
    match Map.tryFind v theta with
    | Some tau -> tau
    | None -> TVarExpr v

let compose (theta2 : Subst) (theta1 : Subst) =
    let domain = Set.union (Map.keys theta2) (Map.keys theta1)
    let replaceTheta = (varsubst theta1) >> (tysubst theta2)
    Map.fromKeys replaceTheta domain

let rec consubst theta con =
    match con with
    | Equal (tau1, tau2, err) ->
        let ts = tysubst theta
        Equal (ts tau1, ts tau2, err)
    | And (c1, c2) ->
        And (consubst theta c1, consubst theta c2)
    | Trivial ->
        Trivial

let (|--->) (a : TyVar) (tau : TyExpr) =
    match tau with
    | TVarExpr b ->
        if a = b then
            Some idSubst
        else
            let (TyVar (_, kindA)) = a
            let (TyVar (_, kindB)) = b
            if kindA = kindB then
                Some (Map.singleton a tau)
            else
                None
    | _ ->
        if Set.contains a (freevars tau) then
            None
        else
            Some (Map.singleton a tau)

let rec solve (con : Constraint) : Subst =
    match con with
    | Trivial -> idSubst
    | And (left, right) ->
        let theta1 = solve left
        let theta2 = solve (consubst theta1 right)
        compose theta2 theta1
    | Equal (tau1, tau2, err) ->
        let failMsg = lazy (sprintf "Type error: The types %s and %s are not equal.\n\n%s" (tyExprString tau1) (tyExprString tau2) (err.Force()))
        match (tau1, tau2) with
        | ((TVarExpr a, tau') | (tau', TVarExpr a)) ->
            match (a |---> tau') with
            | Some answer -> answer
            | None -> raise <| TypeError (failMsg.Force())
        | (TConExpr mu, TConExpr mu') ->
            if mu = mu' then
                idSubst
            else
                raise <| TypeError (failMsg.Force())
        | (TApExpr (l, r), TApExpr (l', r')) ->
            solve (((l =~= l') err) &&& ((r =~= r') err))
        | _ ->
            raise <| TypeError (failMsg.Force())