(* A meta-circular CBV interpreter for PNF in continuation-passing style *)
(* Models user function values with explicit closures. *)

module Pnfinterpc :
sig 
  type value = 
      Vint of int 
    | Vrecord of (Pnf.label * value) list
    | Vvariant of Pnf.label * value
    | Vfn of (value -> value) (* primitive functions *)
    | Vclos of Pnf.id * Pnf.exp * value Env.env ref (* user-defined functions *)

  val evaluate : Pnf.exp -> value
end =
struct
  open Pnf

(* Values manipulated by the interpreter. *)
  type value = 
      Vint of int 
    | Vrecord of (label * value) list
    | Vvariant of label * value
    | Vfn of (value -> value) (* primitive functions *)
    | Vclos of id * exp * value Env.env ref (* user-defined functions *)

(* Various exceptions raised if invalid program is interpreted. *)
  exception Bad_exp
  exception Bad_app
  exception Bad_primarg
  exception Bad_switch
  exception Bad_select
  exception Bad_fix


(* The main interpreter loop. *)
  let rec eval (env:value Env.env) (e:exp) (k:value -> value) : value =
    let rec trivial (e:trexp) =
      match e with
	Var (id,_) -> Env.lookup env id
      | Abs((id,_),e1) -> Vclos(id,e1,ref env)
      | Int i -> Vint i
    and eval0 (e: exp) (k:value->value) : value = 
      match e with
	Triv trexp -> k(trivial trexp)
      |	Let((x,_),a,b) ->
	  let k' = fun a' -> eval (Env.extend env x a') b k in
	  eval0 a k'
      |	App(tre1,tre2) -> apply (trivial tre1) (trivial tre2) k
      |	Fix (fs,e2) ->
	  let knotenv = ref env in
	  let ext env ((fid,_),e0) = 
	    match e0 with
	      Abs((id,_),e1) -> Env.extend env fid (Vclos(id,e1,knotenv)) 
	    | _ -> raise Bad_fix in
	  let newenv = List.fold_left ext (!knotenv) fs in
	  knotenv := newenv;
	  eval newenv e2 k
      | Record el ->  k (Vrecord(List.map (fun (l,tre) -> (l,trivial tre)) el))
      | Select(l, tre) ->  
	  begin
	    match trivial tre with
	      Vrecord(vl) ->
		k(try List.assoc l vl
		  with Not_found -> raise Bad_select)
	    | _ -> raise Bad_select
	  end
      | Variant(l,tre) -> k(Vvariant(l,trivial tre))
      | Switch(tre, el) ->
	  begin
	    match trivial tre with
	      Vvariant(l,v) ->
		begin
		  try 
		    let tre = List.assoc l el in
		    apply (trivial tre) v k
		  with Not_found -> raise Bad_switch
		end
	    | _ ->  raise Bad_switch
	  end in
    eval0 e k

  and apply v1 v2 k =
    match v1 with
      Vclos (id,exp,env) -> eval (Env.extend (!env) id v2) exp k
    | Vfn f -> k(f v2)
    | _ -> raise Bad_app
	  

  let idk v = v

(* Initial environment defining primitive operations *)
  let initial_env =
    let vFalse = eval Env.empty (Tcore2pnf.translate Tcore.eFalse) idk in
    let vTrue = eval Env.empty (Tcore2pnf.translate Tcore.eTrue) idk in
    let vFst = eval Env.empty (Tcore2pnf.translate Tcore.eFst) idk in
    let vSnd = eval Env.empty (Tcore2pnf.translate Tcore.eSnd) idk in
    let iadd_prim = 
      Vfn (fun v -> 
        match (apply vFst v idk,apply vSnd v idk) with
          (Vint i1,Vint i2) -> Vint(i1+i2)
        | _ -> raise Bad_primarg) in
    let imul_prim = 
      Vfn (fun v -> 
        match (apply vFst v idk ,apply vSnd v idk) with
          (Vint i1,Vint i2) -> Vint(i1*i2)
        | _ -> raise Bad_primarg) in
    let ieql_prim =
      Vfn (fun v ->
        match (apply vFst v idk,apply vSnd v idk) with
	  (Vint i1,Vint i2) -> if i1 = i2 then vTrue else vFalse
	| _ -> raise Bad_primarg) in
    List.fold_left2 Env.extend Env.empty Tcore.primops [iadd_prim;imul_prim;ieql_prim] 

(* The top level evaluator *)
  let evaluate (e : exp) : value = eval initial_env e idk

end