amelia
/
ahc-old


module M = import "data/map.ml"

module G = import "./lib/graph.ml"

open import "./lang.ml"

open import "amulet/exception.ml"

open import "prelude.ml"


type tc_tyvar 'a = Tv of {

 name : string, level : int, var : ref (option 'a)

}


instance eq (tc_tyvar 'a) begin

 let Tv x == Tv y = x.name == y.name

end


instance ord (tc_tyvar 'a) begin

 let Tv x `compare` Tv y = x.name `compare` y.name

end


type tc_kappa =

 | K_arr of tc_kappa * tc_kappa

 | K_star

 | K_var of tc_tyvar tc_kappa


type tc_rho =

 | T_uvar of tc_tyvar tc_rho

 | T_var of string

 | T_con of string

 | T_app of tc_rho * tc_rho

 | T_arr of tc_rho * tc_rho


instance show tc_rho begin

 let show =

 let rec show_arg = function

 | T_app _ as x -> "(" ^ go x ^ ")"

 | x -> show_domain x

 and show_domain = function

 | T_arr _ as x -> "(" ^ go x ^ ")"

 | x -> go x

 and go = function

 | T_uvar (Tv n) ->

 match !n.var with

 | Some t -> go t

 | None -> n.name

 | T_var v -> v

 | T_con v -> v

 | T_app (f, x) -> go f ^ " " ^ show_arg x

 | T_arr (a, b) -> show_domain a ^ " -> " ^ go b

 go

end


instance show tc_kappa begin

 let show x =

 let rec go = function

 | K_star -> "*"

 | K_var (Tv v) -> "?" ^ v.name

 | K_arr (a, b) -> show_domain a ^ " -> " ^ go b

 and show_domain = function

 | K_arr _ as x -> "(" ^ show x ^ ")"

 | x -> go x

 go x

end


type tc_sigma =

 Forall of {

 vars : list string,

 body : tc_rho

 }


let rec free_unif_vars = function

 | T_uvar v -> S.singleton v

 | T_var _ -> S.empty

 | T_con _ -> S.empty

 | T_app (f, x) -> S.union (free_unif_vars f) (free_unif_vars x)

 | T_arr (a, b) -> S.union (free_unif_vars a) (free_unif_vars b)


let new_name =

 let c = ref 0

 fun () ->

 c := !c + 1

 "alpha" ^ show !c


let new_tcvar level =

 let name = new_name ()

 Tv { name, level, var = ref None }


let rec zonk = function

 | T_uvar (Tv r) as rho ->

 match !r.var with

 | Some rho -> zonk rho

 | None -> rho

 | T_var v -> T_var v

 | T_con v -> T_con v

 | T_app (f, x) -> T_app (zonk f, zonk x)

 | T_arr (f, x) -> T_arr (zonk f, zonk x)


let empty (Tv r) =

 match !r.var with

 | None -> true

 | Some (T_uvar (Tv r')) -> r.name == r'.name

 | _ -> false


let generalise level rho =

 let rho = zonk rho

 let vars =

 free_unif_vars rho

 |> S.filter (fun (Tv r) -> r.level > level && empty (Tv r))

 |> S.members

 flip iter vars @@ fun (Tv r) ->

 r.var := Some (T_var r.name)

 Forall { vars = map (fun (Tv r) -> r.name) vars, body = zonk rho }


let rec unify a b =

 let solve r s =

 match !r.var with

 | Some t -> unify t s

 | None -> r.var := Some s

 match a, b with

 | T_uvar (Tv r), b -> solve r b

 | a, T_uvar (Tv r) -> solve r a

 | T_var a, T_var b when a == b -> ()

 | T_con a, T_con b when a == b -> ()

 | T_app (f, x), T_app (f', x') ->

 unify f f'

 unify x x'

 | T_arr (a, b), T_arr (a', b') ->

 unify a a'

 unify b b'

 | a, b -> error @@ "Types " ^ show a ^ " and " ^ show b ^ " are not equal"


let rec unify_kappa a b =

 let solve r s =

 match !r.var with

 | Some t -> unify_kappa t s

 | None -> r.var := Some s

 match a, b with

 | K_var (Tv r), b -> solve r b

 | a, K_var (Tv r) -> solve r a

 | K_star, K_star -> ()

 | K_arr (a, b), K_arr (a', b') ->

 unify_kappa a a'

 unify_kappa b b'

 | a, b -> error @@ "Kinds " ^ show a ^ " and " ^ show b ^ " are not equal"


type scheme 'a = Poly of tc_sigma | Mono of 'a


instance show 'a => show (scheme 'a) begin

 let show = function

 | Poly (Forall {vars,body}) ->

 foldl (fun s i -> s ^ " " ^ i) "forall" vars ^ ". " ^ show body

 | Mono x -> show x

end


let mono m = function

 | Mono x -> x

 | Poly _ -> error @@ "Unexpected polytype " ^ m


let get_scope map var =

 match M.lookup var map with

 | Some v -> v

 | None -> error @@ "Name not in scope: " ^ var


let is_function_kind level tau =

 match tau with

 | K_arr (a, b) -> (a, b)

 | _ ->

 let a = new_tcvar level |> K_var

 let b = new_tcvar level |> K_var

 unify_kappa tau (K_arr (a, b))

 (a, b)


let rec infer_kind scope = function

 | Tyvar v ->

 let kappa = get_scope scope v |> mono "(kinds aren't ever polymorphic)"

 (T_var v, kappa)

 | Tycon v ->

 let kappa = get_scope scope v |> mono "(kinds aren't ever polymorphic)"

 (T_con v, kappa)

 | Tyapp (f, x) ->

 let (f, k_f) = infer_kind scope f

 let (x, k_x) = infer_kind scope x

 let (domain, result) = is_function_kind 0 k_f

 unify_kappa domain k_x

 (T_app (f, x), result)

 | Tyarr (a, b) ->

 let a = check_is_type scope a

 let b = check_is_type scope b

 (T_arr (a, b), K_star)

 | Tytup [] -> (T_con "Unit#", K_star)

 | _ -> error "Tuple types not supported"

and check_is_type scope t =

 let (t, k) = infer_kind scope t

 unify_kappa k K_star

 t


let rec default_to_star = function

 | K_var (Tv r) ->

 match !r.var with

 | Some k -> default_to_star k

 | None -> K_star

 | K_star -> K_star

 | K_arr (a, b) -> K_arr (default_to_star a, default_to_star b)


type dt_info <-

 { name : string, d_args : list string, c_args : list tc_rho, c_ret : tc_rho }


let mk_con_info (d_name : string) (d_args : list string) : list (string * list tc_rho) -> list dt_info =

 let go (name, args) =

 { name, c_args = args, d_args, c_ret = foldl (fun f x -> T_app (f, T_var x)) (T_con d_name) d_args }

 map go


let infer_data_group_kind scope (group : list _) =

 let init_kind (group, names) (name, args, constr) =

 let args =

 args |> map (fun v -> (v, new_tcvar 0 |> K_var |> Mono))

 let kind = foldl (fun t (_, r) -> K_arr (t, mono "" r)) K_star args

 let scope = M.from_list args

 ((name, kind, constr, scope, args) :: group, M.insert name (Mono kind) names)


 let (group, scope') = foldl init_kind ([], M.empty) group


 let scope = M.union scope scope'


 let group : list (string * tc_kappa * list string * list (string * list tc_rho)) =

 flip map group @@ fun (name, kind, constrs, args, args') ->

 let scope = M.union scope args

 constrs

 |> map (fun (Constr (name, args)) -> (name, map (check_is_type scope) args))

 |> (name,kind,[x|with (x,_)<-args'],)


 flip map group @@ fun (name, kind, args, constrs) ->

 (name, default_to_star kind, constrs, mk_con_info name args constrs)


let rec subst_vars f = function

 | T_var v as t ->

 match f v with

 | None -> t

 | Some t -> t

 | T_uvar (Tv v) as t ->

 match !v.var with

 | Some t -> subst_vars f t

 | None -> t

 | T_con v -> T_con v

 | T_app (a, b) -> T_app (subst_vars f a, subst_vars f b)

 | T_arr (a, b) -> T_arr (subst_vars f a, subst_vars f b)


let instantiate level (Forall { vars, body }) =

 let vars =

 vars

 |> map (fun v -> (v, new_tcvar level |> T_uvar))

 |> M.from_list

 subst_vars (flip M.lookup vars) body


let lookup_ty level scope v =

 get_scope scope v |> function

 | Mono t -> t

 | Poly s -> instantiate level s


let is_function_type level tau =

 match tau with

 | T_arr (a, b) -> (a, b)

 | _ ->

 let a = new_tcvar level |> T_uvar

 let b = new_tcvar level |> T_uvar

 unify tau (T_arr (a, b))

 (a, b)


(* TODO: Rank-N types *)

let is_subtype = unify


let rec infer dt_info level scope = function

 | Ref v -> lookup_ty level scope v |> (Ref v,)

 | App (f, x) ->

 let (f, arg, res) =

 infer dt_info level scope f

 |> second (is_function_type level)

 let x = check dt_info level scope arg x

 (App (f, x), res)

 | Lit i -> (Lit i, T_con "Int")

 | Let (bindings, body) ->

 let (bindings, scope') =

 infer_binding_group dt_info level scope bindings

 let (body, body_t) = infer dt_info level (scope `M.union` map force scope') body

 (Let (bindings, body), body_t)

 | x ->

 let t = new_tcvar level |> T_uvar

 let x = check dt_info level scope t x

 (x, t)


and check dt_info level scope wanted = function

 | Lam (arg, body) ->

 let (arg_t, body_t) = is_function_type level wanted

 let body =

 (* TODO: Rank-N types *)

 check dt_info level (M.insert arg (Mono arg_t) scope) body_t body

 Lam (arg, body)

 | Case (_, []) -> error "Empty case"

 | Case (scrutinee, Cons ((con, _, _), _) as patterns) ->

 let data =

 match M.lookup con dt_info with

 | Some data -> data

 | None -> error @@ "Constructor " ^ con ^ " doesn't belong to a type"


 let (scrutinee, scrut_t) = infer dt_info level scope scrutinee


 let go_arm {name, d_args, c_args, c_ret} (con, args, expr) =

 if name <> con then

 error @@ "Constructors out of order: expected this pattern to match " ^ name

 else ()


 if length c_args <> length args then

 error @@ "Constructor "

 ^ con ^ " has "

 ^ show (length c_args)

 ^ " but is being matched against with " ^ show (length args)

 ^ " variables"

 else ()


 let d_args =

 d_args

 |> map (fun v -> (v, new_tcvar level |> T_uvar))

 |> M.from_list

 let c_args = map (Mono # subst_vars (flip M.lookup d_args)) c_args

 let c_ret = subst_vars (flip M.lookup d_args) c_ret


 unify c_ret scrut_t


 let scope' = M.from_list (zip args c_args) `M.union` scope

 (con, args, check dt_info level scope' wanted expr)


 Case (scrutinee, zip_with go_arm data patterns)

 | If (cond, e_then, e_else) ->

 let cond = check dt_info level scope (T_con "Bool") cond

 let e_t = check dt_info level scope wanted e_then

 let e_e = check dt_info level scope wanted e_else

 Case (cond, [ ("True", [], e_t), ("False", [], e_e) ])

 | x ->

 let (x, t) = infer dt_info level scope x

 is_subtype t wanted

 x


and infer_binding_group dt_info level (scope : M.t string _) bindings : _ * M.t string _ =

 let inner = level + 1

 let initial_types =

 bindings

 |> map (fun (name, _) -> (name, new_tcvar inner |> T_uvar |> Mono))

 |> M.from_list


 let initial_types = initial_types |> M.union scope


 let go_binding (bindings : list _, scope' : M.t _ _) (name : string, body : expr) =

 let (body, body_ty) =

 (fun () -> infer dt_info inner initial_types body)

 `catch` fun (e : some exception) ->

 error (describe_exception e ^ "\nwhen type checking " ^ name)

 M.lookup name scope

 |> function

 | Some (Mono t) -> unify t body_ty

 | _ -> ()

 (

 (name, body) :: bindings,

 M.insert name (lazy (generalise level body_ty |> Poly)) scope'

 )

 foldl go_binding ([], M.empty) bindings


let dependency_graph defs =

 let rec free_vars_of_cons t m (Constr (name, args)) =

 let cons =

 foldl (fun s t -> S.union s (free_cons t)) (S.singleton t)

 args

 M.insert name cons m

 let define n x m =

 M.alter (function

 | Some _ -> error @@ "Redefinition of value " ^ n

 | None -> Some x)

 n m

 let go (graph, defs) = function

 | Foreign (Fimport { var }) as x ->

 (M.insert var S.empty graph, define var x defs)

 | Decl (name, args, expr) as x ->

 let fvs =

 free_vars expr

 |> flip S.difference (S.from_list args)

 |> S.delete name

 (M.insert name fvs graph, define name x defs)

 | Data (name, _, cons) as x ->

 M.union graph (foldl (free_vars_of_cons name) M.empty cons)

 |> M.insert name S.empty

 |> (, define name x defs)

 let (graph, defs) = foldl go (M.empty, M.empty) defs

 (G.groups_of_sccs graph, defs)


let mk_lam args body = foldr (curry Lam) body args

let rec unlambda = function

 | Lam (v, x) ->

 let (args, x) = unlambda x

 (v :: args, x)

 | e -> ([], e)


let rec replicate n x =

 if n <= 0 then

 []

 else

 x :: replicate (n - 1) x


let rec add_missing_vars scope = function

 | Tyvar v ->

 match M.lookup v scope with

 | Some _ -> scope

 | None ->

 let k = new_tcvar 0 |> K_var

 M.insert v (Mono k) scope

 | Tycon _ -> scope

 | Tyapp (a, b) -> add_missing_vars (add_missing_vars scope b) a

 | Tyarr (a, b) -> add_missing_vars (add_missing_vars scope b) a

 | Tytup xs -> foldl add_missing_vars scope xs


let tc_program value_exports type_exports (prog : list decl) =

 let (plan, defs) = dependency_graph prog


 let tc_one (dt_info, val_scope, ty_scope, out) group =

 let defs = [ x | with name <- S.members group, with Some x <- [M.lookup name defs] ]

 match defs with

 | [] -> (dt_info, val_scope, ty_scope, out)

 | [Foreign (Fimport {var, ftype}) as def] ->

 let ty_scope' = add_missing_vars M.empty ftype

 let t = check_is_type (M.union ty_scope' ty_scope) ftype

 (

 dt_info,

 M.insert var (Forall { vars = M.keys ty_scope', body = t } |> Poly) val_scope,

 ty_scope, def :: out

 )

 | Cons (Foreign (Fimport {var}), _) ->

 error @@ "Foreign definition " ^ var ^ " is part of a group. How?"

 | Cons (Decl (name, args, body), ds) ->

 let bindings =

 (name, mk_lam args body)

 :: [ (name, mk_lam args body) | with Decl (name, args, body) <- ds ]

 let (bindings, scope') = infer_binding_group dt_info -1 val_scope bindings

 let decs =

 [ Decl (name, unlambda expr) | with (name, expr) <- bindings ]

 (dt_info, M.union (map force scope') val_scope, ty_scope, foldr (::) decs out)

 | Cons (Data d, ds) ->

 let datas = d :: [ d | with Data d <- ds ]

 let r = infer_data_group_kind ty_scope datas

 let fix_constr (name, rhos : list tc_rho) =

 Constr (name, replicate (length rhos) (Tycon "#"))

 let rec go dt ty (vl : M.t string (scheme tc_rho)) ds = function

 | [] -> (dt, vl, ty, ds)

 | Cons ((name, kind, constrs, info : list dt_info), rest) ->

 go

 (foldl (fun i {name} -> M.insert name info i) dt info)

 (M.insert name (Mono kind) ty)

 (foldl

 (fun s {name,d_args,c_args,c_ret} ->

 M.insert name (Forall { vars = d_args, body = foldr (curry T_arr) c_ret c_args} |> Poly) s)

 vl info)

 (Data (name, [], fix_constr <$> constrs) :: ds)

 rest

 go dt_info ty_scope val_scope out r

 let (_, vals, types, p) = foldl tc_one (M.empty, M.empty, M.empty, []) plan

 (

 [x | with k <- value_exports, with Some x <- [M.lookup k vals]],

 [x | with k <- type_exports, with Some x <- [M.lookup k types]],

 p

 )