(* CIS 341 *)
(* Example code about compilation of first class (higher-order) functions *)

module Fun = struct 
  type var = string
 
  (* Abstract syntax of "lambda terms" *)
  type exp = 
      Int of int              
    | Add of exp * exp
    | Var of var                (* local variables *)
    | Fun of var * exp          (* functions:  fun x -> e *)
    | App of exp * exp          (* function application *)

  type value =
    | IntV of int
    | FunV of var * exp

  (* values are just a subset of expressions *)
  let exp_of_value v = 
    match v with
      | IntV i -> Int i
      | FunV (arg, body) -> Fun(arg, body)

  (* Example term: *)
  (* ( ( (fun x -> fun y -> x + y) 3 ) 4 ) *)
  let ex = App (App (Fun ("x", Fun ("y", Add (Var "x", Var "y"))), 
                     Int 3), 
                Int 4)

end

(* A substitution-based evaluator *)
module Eval1 = struct
  open Fun

  (* Substitute the (closed) value v for x in e *)
  let rec subst v x e =
    match e with
      | Int i -> Int i
      | Add (e1, e2) -> Add(subst v x e1, subst v x e2)
      | Var y -> if y = x then exp_of_value v else Var y
      | Fun (arg, body) -> 
	  if (arg = x) then Fun (arg, body) 
	  else Fun(arg, subst v x body)
      | App (e1, e2) -> App(subst v x e1, subst v x e2)

  (* A substituion-based interpreter *)
  let rec eval e =
    match e with
      | Int i -> IntV i
      | Add (e1, e2) -> 
	  (match (eval e1, eval e2) with
	    | (IntV i1, IntV i2) -> IntV (i1 + i2)
	    | _ -> failwith "tried to add non-integers")
      | Var x -> failwith "reached a free variable"
      | Fun (arg, body) -> FunV (arg, body)
      | App (e1, e2) -> 
	  (match(eval e1, eval e2) with
	     | (FunV (x, body), v) -> eval(subst v x body)
	     | _ -> failwith "tried to apply non-function")

end

(* A *broken* attempt to use an environment-based evaluator *)
module Eval2 = struct
  open Fun

  type environment = (var * value) list
      
  let lookup x (env:environment) = List.assoc x env

  (* The following evaluator gives *dynamic* scoping, not
     static (or 'lexical') scoping! *)
  let rec eval e env =
    match e with 
      | Int i -> IntV i
      | Add (e1, e2) ->
	  (match (eval e1 env, eval e2 env) with
	     | (IntV i1, IntV i2) -> IntV (i1 + i2)
	     | _ -> failwith "tried to add non-integers")
      | Var x -> lookup x env
      | Fun (arg, body) -> FunV(arg, body)
      | App (e1, e2) ->
	  (match(eval e1 env, eval e2 env) with
	     | (FunV (x, body), v) -> eval body ((x,v)::env)
	     | _ -> failwith "tried to apply non-function")

end

(* A "closure"-based interpreter that fixes the problem with Eval2 *)
(* A closure is a pair consisting of a datastructure that represents 
   the evaluation context and a function *)
module Eval3 = struct
  open Fun

  (* Change the definition of value... *)
  type value =
    | IntV of int
    | Closure of environment * var * exp
  and environment = (var * value) list
      
  let lookup x (env:environment) = List.assoc x env

  (* The following evaluator fixes the problems with Eval2 *)
  let rec eval e env =
    match e with 
      | Int i -> IntV i
      | Add (e1, e2) ->
	  (match (eval e1 env, eval e2 env) with
	     | (IntV i1, IntV i2) -> IntV (i1 + i2)
	     | _ -> failwith "tried to add non-integers")
      | Var x -> lookup x env
      | Fun (arg, body) -> Closure(env, arg, body)
      | App (e1, e2) ->
	  (match(eval e1 env, eval e2 env) with
	     | (Closure(cenv, x, body), v) -> eval body ((x,v)::cenv)
	     | _ -> failwith "tried to apply non-function")

end

(* A "closure-based" intermediate language based on Eval3 *)
module ClosureConvert = struct
  type var = string

  (* A closure is pair consisting of an environment and a code pointer  *)
  (* The environment of a closure is a tuple of values representing the *)
  (* values of the context at the time the closure is created.          *)

  (* Since we need pairs and tuples,   *)
  (* the target language has to be extended *)
  type exp =
    | Val of value
    | Var of var               (* local variables *)
    | Global of var            (* global variables *)
    | Add of exp * exp
    | App of exp * exp
    | Let of var * exp * exp   (* introduce local variables *)
    | Tuple of exp list
    | Nth of exp * int

  and value =
    | IntV of int
    | Code of var * var * exp   (* Environment name, arg name, body *)
    | TupleV of value list


  (* printing *)
  let ps = print_string
  let rec pp_e e =
    match e with
      | Val v -> pp_v v
      | Var v -> ps v
      | Global v -> ps v
      | Add (e1, e2) -> ps "Add("; pp_e e1; ps ", "; pp_e e2; ps ")"
      | App (e1, e2) -> ps "App("; pp_e e1; ps ", "; pp_e e2; ps ")"
      | Let (v, e1, e2) -> ps "Let("; ps v; ps ", "; pp_e e1;
          ps ", "; pp_e e2; ps ")"
      | Tuple elist -> ps "Tuple("; List.iter (fun e -> pp_e e; ps ", ") elist; ps ")"
      | Nth (e, i) -> ps "Nth("; pp_e e; print_string (", " ^ string_of_int i)
  and pp_v v =
    match v with
      | IntV i -> print_int i
      | Code(env, x, body) -> ps "ClosureV("; ps env; ps ", "; ps x; ps ", ";
	  pp_e body; ps ")"
      | TupleV vlist -> ps "TupleV("; List.iter (fun v -> pp_v v; ps ", ") vlist; ps ")"
	  
	  
  type environment = var list

  let mk_tmp =
    let ctr = ref 0 in fun () ->
      let x = !ctr in ctr := x + 1; ("TMP" ^ string_of_int(x))

  (*************************************)
  (* The main closure conversion code. *)
  (*************************************)

  (* A function prologue that sets up the expected local variables. *)
  let build_local_env env tmp body =
    let (_, code) = 
      List.fold_left (fun (i, code) -> fun x ->
                        (i+1, Let(x, Nth(Var tmp, i), code))) 
           (0, body) env
    in
      code

  (* In practice, this would store only the free variables in the function. *)
  (* If so, then build_local_env would be modified appropriately. *)
  let build_closure_env env =
    Tuple (List.map (fun x -> Var x) env)

  let rec convert (e:Fun.exp) env =
    match e with
      | Fun.Int i -> Val (IntV i)
      | Fun.Add (e1, e2) -> Add (convert e1 env, convert e2 env)
      | Fun.Var x -> Var x
      | Fun.Fun (arg, body) -> 
	  let env_name = mk_tmp () in
	  let body' = build_local_env env env_name (convert body (arg::env)) in
	  let env' = build_closure_env env in
	    Tuple [env'; Val(Code(env_name, arg, body'))]
      | Fun.App (e1, e2) -> App (convert e1 env, convert e2 env)

  (* Assumes that all closures are closed and *) 
  (* moves all closures to the top level.     *)
  (* In practice, hoisting and closure conversion *)
  (* could be combined *)
  let hoist e =
    let rec help_e (e:exp):((var * value) list * exp) =
      match e with
	| Val(Code(env, x, body)) -> 
	    let (c1, r1) = help_e body in
	    let tmp = mk_tmp () in
	      ((tmp, Code(env, x, r1))::c1, Global tmp)
	| Val(v) -> 
	    let (c1, r1) = help_v v in
	      (c1, Val r1)
	| Var x -> ([], Var x)
	| Global x -> ([], Global x)
	| Add(e1, e2) ->
	    let (c1, r1) = help_e e1 in
	    let (c2, r2) = help_e e2 in
	      (c1@c2, Add(r1, r2))
	| App(e1, e2) ->
	    let (c1, r1) = help_e e1 in
	    let (c2, r2) = help_e e2 in
	      (c1@c2, App(r1, r2))
	| Let(x, e1, e2) ->
	    let (c1, r1) = help_e e1 in
	    let (c2, r2) = help_e e2 in
	      (c1@c2, Let(x, r1, r2))
	| Tuple(elist) ->
	    let (cs, rs) = List.split(List.map help_e elist) in
	      (List.concat cs, Tuple rs)
	| Nth(e1, i) ->
	    let (c1, r1) = help_e e1 in
	      (c1, Nth(r1, i))

    and help_v v =
      match v with
	| IntV i -> ([], IntV i)
	| TupleV(vlist) ->
	    let (cs, rs) = List.split(List.map help_v vlist) in
	      (List.concat cs, TupleV rs)
	| _ -> failwith "impossible"
    in	    
      help_e e

  (* Assumes that Code is closed *)
  let rec subst v x e =
    match e with 
      | Val u -> Val u     (* Note that we don't substitute through values *)
      | Var y -> if y = x then Val v else e
      | Global x -> Global x
      | Add(e1, e2) -> Add(subst v x e1, subst v x e2)
      | App(e1, e2) -> App(subst v x e1, subst v x e2)
      | Let(y, e1, e2) ->
	  let e1' = subst v x e1 in
	    if x = y then Let(y, e1', e2) else Let(y, e1', subst v x e2)
      | Tuple(elist) -> Tuple(List.map (subst v x) elist)
      | Nth(e, i) -> Nth(subst v x e, i)

  let rec eval globals e =
    let rec eval_e e = 
      pp_e e; ps "\n";
      match  e with
	| Var x -> failwith "tried to evaluate free variable"
	| Global x -> List.assoc x globals
	| Val u -> u
	| Add(e1, e2) -> 
	    (match (eval_e e1, eval_e e2) with
	       | (IntV i1, IntV i2) -> IntV (i1 + i2)
	       | _ -> failwith "tried to add two non-integers")

	(* Since closures are pairs, application expects to find a tuple. *)
        (* Since all function code takes two parameters, this interpreter *)
        (* "hardwires" the function call.*)
	| App(e1, e2) -> 
	    (match (eval_e e1, eval_e e2) with
	       | (TupleV [env; Code(env_name, x, b)], v) -> 
		   eval_e (subst v x (subst env env_name b))
	       | _ -> failwith "tried to apply non-closure") 
	| Let(x, e1, e2) -> eval_e (subst (eval_e e1) x e2)
	| Tuple(elist) -> TupleV(List.map eval_e elist)
	| Nth(e, i) -> 
	    (match  eval_e e with
	       | TupleV vlist -> List.nth vlist i
	       | _ -> failwith "tried to project from non tuple")
    in eval_e e
	 
end