(* Convert Tcore to Pnf, recording types on all variables as we go. *)

module Tcore2pnf :
sig 
  val translate : Tcore.exp -> Pnf.exp
end = 
struct
  open Pnf

  exception Non_trivial
  exception Bad_fix

  let rec trans (env:Tcore.typ Env.env) (e:Tcore.exp) : Pnf.exp =
    match e with
      Tcore.App(Tcore.Abs(v,t,b),a) ->
	Let((v,t),trans env a,trans (Env.extend env v t) b)
    | Tcore.App(e1,e2) -> 
	trivialize env e1 (fun tre1 -> 
	  trivialize env e2 (fun tre2 -> 
	    App(tre1,tre2)))
    | Tcore.Fix(fs,b) ->  
	let ext env (f,e0,rt) = 
	  match e0 with
	    Tcore.Abs(_,vt,_) -> Env.extend env f (ArrowT(vt,rt)) 
	  | _ -> raise Bad_fix in
	let env1 = List.fold_left ext env fs in
	Fix(List.map (fun (f,e,_) -> ((f,Env.lookup env1 f),translate_triv env1 e)) fs,
	    trans env1 b)
    | Tcore.Record r -> 
	trivialize_row env r (fun trr -> Record trr)
    | Tcore.Select(l,e) ->
	trivialize env e (fun tre ->
	  Select(l,tre))
    | Tcore.Variant(l,e) ->
	trivialize env e (fun tre ->
	  Variant(l,tre))
    | Tcore.Switch(e,r) -> 
	trivialize env e (fun tre ->
	  trivialize_row env r (fun trr ->
	    Switch(tre,trr)))
    | _ -> Triv(translate_triv env e)
  
  and translate_triv env (e:Tcore.exp) : Pnf.trexp =
    match e with
      Tcore.Var v -> Var (v,Env.lookup env v)
    | Tcore.Int i -> Int i
    | Tcore.Abs(v,t,e) -> Abs((v,t),trans (Env.extend env v t) e)
    | _ -> raise Non_trivial

  and trivialize env
                 (e:Tcore.exp)  
                 (k:Pnf.trexp -> Pnf.exp) : Pnf.exp =
    try 
      k(translate_triv env e)
    with Non_trivial ->
      let xt = gensym(Check.check_in_env env e) in 
      Let(xt,trans env e,k(Var xt))

  and trivialize_row env lel (k:Pnf.row -> Pnf.exp) : Pnf.exp =
    let rec f lel xl =
      match lel with
	(l,e)::ler ->
	  trivialize env e 
	    (fun tre -> 
	      f ler ((l,tre)::xl))
      |	[] -> k (List.rev xl) in
    f lel []

  let translate = trans Check.initial_env

end