toy type inference

gistfile1.fs

module ToyTypeInfer

type Var = string

type Expr =
  | ENum of int
  | EBool of bool
  | EVar of Var
  | EPlus of Expr * Expr
  | EIf of Expr * Expr * Expr
  | ELet of Var * Expr * Expr
  | EFun of Var * Expr
  | EApp of Expr * Expr
  | ELetRec of Var * Expr * Expr

type Type =
  | TInt
  | TBool
  | TFun of Type * Type
  | TVar of int
  override this.ToString() =
    match this with
    | TInt -> "int"
    | TBool -> "bool"
    | TFun(t1,t2) -> t1.ToString() + " -> " + t2.ToString()
    | TVar(n) -> "a" + n.ToString()

///新しい一意な型変数を作る
let makeNewTVar =
  let count = ref 0
  fun () -> incr count; TVar(!count)

///型の方程式を摘出する
let rec extract tyEnv = function
  | ENum(_) -> [], TInt
  | EBool(_) -> [], TBool
  | EVar(var) -> [], Map.find var tyEnv
  | EPlus(e1, e2) ->
    let e1, t1 = extract tyEnv e1
    let e2, t2 = extract tyEnv e2
    let e3 = List.concat [e1; e2; [t1, TInt; t2, TInt]]
    e3, TInt
  | EIf(ce, te, fe) ->
    let e1, t1 = extract tyEnv ce
    let e2, t2 = extract tyEnv te
    let e3, t3 = extract tyEnv fe
    let e4 = List.concat [e1; e2; e3; [t1, TBool; t2, t3]]
    e4, t3
  | ELet(var, e1, e2) ->
    let e1, t1 = extract tyEnv e1
    let e2, t2 = extract (Map.add var t1 tyEnv) e2
    let e3 = List.append e1 e2
    e3, t2
  | EFun(var, e) ->
    let a = makeNewTVar()
    let e, t = extract (Map.add var a tyEnv) e
    e, TFun(a, t)
  | EApp(e1, e2) ->
    let e1, t1 = extract tyEnv e1
    let e2, t2 = extract tyEnv e2
    let a = makeNewTVar()
    let e3 = List.concat [e1; e2; [t1, TFun(t2,a)]]
    e3, a
  | ELetRec(var1, EFun(var2, e1), e2) ->
    let a1 = makeNewTVar()
    let a2 = makeNewTVar()
    let tyEnv = Map.add var1 a1 tyEnv
    let e1, t1 = extract (Map.add var2 a2 tyEnv) e1
    let e2, t2 = extract tyEnv e2
    let e3 = List.concat [e1; e2; [a1, TFun(a2,t1)]]
    e3, t2
  | _ -> failwith "invalid form"

///型の中にある自由変数を見つける
let rec FTV = function
  | TVar(n) -> [TVar(n)]
  | TBool -> []
  | TInt -> []
  | TFun(t1,t2) -> List.append (FTV t1) (FTV t2)
  
///ts内にaがあったらtに置き換える
let sub t a ts =
  ts
  |> List.map (fun (t1, t2) ->
    [t1; t2]
    |> List.map (fun t' ->
      if a = t' then t
      else t')
    |> function
    | [t1; t2] -> (t1, t2)
    | _ -> failwith "")

///型の方程式を解く関数を返す
let rec unify =
  let rec makeSub t1 a =
    function
    | TVar(_) as b when a = b -> t1
    | TFun(x, y) ->
      TFun(makeSub t1 a x, makeSub t1 a y)
    | t -> t
  function
  | [] -> id
  | (t1, t2)::env when t1=t2 -> unify env
  | ((TVar(_) as a), t2)::env when a<>t2 ->
    if List.exists ((=) a) <| FTV(t2) then
      failwith "unification failure"
    else
      let subs = unify <| sub t2 a env
      subs << makeSub t2 a
  | (t1, (TVar(_) as a))::env when a<>t1 ->
    unify <| (a, t1)::env
  | (TFun(t11, t12), TFun(t21, t22))::env ->
    unify <| (t11, t21)::(t12, t22)::env
  | _ -> failwith "unification failure"

let infer =
  extract Map.empty >>
  fun (x, ans) ->
    let s = unify x
    let mutable prev = ans
    let mutable curr = s prev
    while prev <> curr do
      prev <- curr
      curr <- s curr
    curr

///推論する式
/// let x =
///   let z = 123
///   (fun y -> y + z) z
/// fun y -> y x
let expr =
  ELet("x",
      ELet("z",
        ENum(123),
        EApp(EFun("y", EPlus(EVar("y"), EVar("z"))),
            EVar("z"))),
      EFun("y", EApp(EVar("y"), EVar("x"))))

[<EntryPoint>]
let main _ =
  expr
  |> infer
  |> string
  |> printfn "%A"
  0