(* A meta-circular CBV interpreter for "extra-normalized" PNF *)

module Pnfinterpn :
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
  exception Bad_let

(* The main interpreter loop. *)
  let rec eval (env:value Env.env) (e : exp) : value =
    let trivial (e:trexp) =
      match e with
	Var (id,_) -> Env.lookup env id
      | Abs((id,_),e1) -> Vclos(id,e1,ref env)
      | Int i -> Vint i in
    match e with
      Triv trexp -> trivial trexp
    | Let((x,_),App(tre1,tre2),b) ->
	let v1 = trivial tre1 in
	let v2 = trivial tre2 in
	begin
	  match v1 with
	    Vclos(id,exp,envr) -> (* user-defined funciton *)
	      let a' = eval (Env.extend (!envr) id v2) exp in
	      eval (Env.extend env x a') b
	  | Vfn f -> (* primitive function *)
	      eval (Env.extend env x (f v2)) b
	  | _ -> raise Bad_app
	end
    | Let((x,_),Triv trexp,b) ->
	let a' = trivial trexp in
	eval (Env.extend env x a') b
    | Let((x,_),Record el,b) ->
	let a' = Vrecord(List.map (fun (l,tre) -> (l,trivial tre)) el) in
	eval (Env.extend env x a') b
    | Let((x,_),Select (l,tre),b) ->
	let a' = 
	  begin
	    match trivial tre with
	      Vrecord(vl) ->
		begin
		  try List.assoc l vl  
		  with Not_found -> raise Bad_select
		end
	    | _ -> raise Bad_select
	  end in
	eval (Env.extend env x a') b
    | Let((x,_),Variant (l,tre),b) ->
	let a' = Vvariant(l,trivial tre) in
	eval (Env.extend env x a') b
    | Let _ -> raise Bad_let
    | 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
    | Switch(tre, el) ->
	begin
	  match trivial tre with
	    Vvariant(l,v) ->
	      let arm = 
		try List.assoc l el 
		with Not_found -> raise Bad_switch in
	      begin
		match arm with
		  Abs((x,_),e1) -> eval (Env.extend env x v) e1
		| _ ->  raise Bad_switch
	      end
	  | _ -> raise Bad_switch
	end 
    | _ -> raise Bad_exp


(* Initial environment defining primitive operations.
   This version is dirtier than before, because we don't want to use the standard
   evaluation machinery inside primops. *)
  let initial_env =
    let vFalse = eval Env.empty (Pnfnorm.norm (Tcore2pnf.translate Tcore.eFalse)) in
    let vTrue = eval Env.empty (Pnfnorm.norm (Tcore2pnf.translate Tcore.eTrue)) in
    let do_select label v =
      match v with
	Vrecord vl -> 
	  begin
	    try List.assoc label vl
	    with Not_found -> raise Bad_select
	  end
      |	_ -> raise Bad_select in
    let vFst = eval Env.empty (Pnfnorm.norm (Tcore2pnf.translate Tcore.eFst)) in
    let vSnd = eval Env.empty (Pnfnorm.norm (Tcore2pnf.translate Tcore.eSnd)) in
    let iadd_prim = 
      Vfn (fun v -> 
        match (do_select fst_label v,do_select snd_label v) with
          (Vint i1,Vint i2) -> Vint(i1+i2)
        | _ -> raise Bad_primarg) in
    let imul_prim = 
      Vfn (fun v -> 
        match (do_select fst_label v,do_select snd_label v) with
          (Vint i1,Vint i2) -> Vint(i1*i2)
        | _ -> raise Bad_primarg) in
    let ieql_prim =
      Vfn (fun v ->
        match (do_select fst_label v,do_select snd_label v) 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 

end