(***************** DEEP embedding *******************

This file illustrates the process of going from an 
interpreter for a deep embedding to a stack-based compiler. 
The process starts with a simple recursive interpreter
folled by a CPS-ed version of that interpreter.  This is 
then followed by a sequence of "defunctionalisations" 
and refactorings. 

---Tim Griffin (for Metaprogramming L305)
   November 2018 

***)

(* lots and lots of "tags" .... *)   
type simple_ast =
  | Var of string 
  | Int of int
  | Bool of bool
  | Add  of simple_ast * simple_ast
  | App  of simple_ast * simple_ast
  | Lam  of string * simple_ast 
  | If   of simple_ast * simple_ast * simple_ast
			
			    
(* Signature for DEEP embedding. 
   This is meant to look similar to the file taglesscps.ml
*) 		    
module type DEXP = sig
  type ast 
  type result
  type trace 
  val var : string -> ast 	 
  val int : int -> ast
  val bool : bool -> ast 		     
  val add : ast -> ast -> ast 
  val if_ : ast -> ast -> ast -> ast 
  val lam : string -> ast -> ast 
  val app : ast -> ast -> ast 
  val eval : ast -> result
  val trace_eval : ast -> trace 
end

(* some examples *) 		     
module DeepExample (E: DEXP) =
struct
  open E
  (** The double application function λf.λx.f (f x) *)
  let t0 = lam "f" (lam "x" (app (var "f") (app (var "f") (var "x"))))

  (** An application of t0 to the identity *)
  let t1 = app t0 (lam "y" (var "y"))

  (** A function that doesn't specialize (partially evaluate) very well:
        λb.(2 + if b then 3 else 4)
   *)
  let t2 = lam "b" ((add (int 2) (if_ (var "v") (int 3) (int 4))))

  (* (λf. f(1) + f(2))(λx. x + 17) 
     should eval to 37
   *) 	       
  let t3 = app (lam "f" (add (app (var "f") (int 1)) (app (var "f") (int 2)))) (lam "x" (add (var "x") (int 17)))

  (* (λf. (λx. f x) 17)((λx. λy. x + y) 1) 
     should eval to 18 (this is a test that environments are implemented correctly) 
   *) 	       
  let t4 = app (lam "f" (app (lam "x" (app (var "f") (var "x"))) (int 17))) (app (lam "x" (lam "y" (add (var "x") (var "y")))) (int 1))

  (* should eval to 15 *) 	       
  let t5 = add (add (add (int 0) (int 1) ) (int 2)) (add (add (int 3) (int 4)) (int 5))   	       
end

(* ... more "tags" .... *)   
type value = 
  | RuntimeError 
  | IntVal of int
  | BoolVal of bool
  | FunVal of (value -> value)
		     
module Deep : DEXP with type ast = simple_ast and type result = value = 
  struct

  type ast = simple_ast     
  type result = value
  type trace = unit 
	 
  let var s = Var s 
  let int x = Int x
  let bool b = Bool b 
  let add x y = Add (x, y) 
  let lam x e = Lam(x, e)
  let app e1 e2 = App (e1, e2) 
  let if_ e1 e2 e3 = If(e1, e2, e3)

  (* a simple interpreter *) 		       
  let rec eval' (e, env) =
  match e with
  | Var x -> env x 
  | Int n -> IntVal n  
  | Bool b -> BoolVal b 
  | Add(e1, e2) ->
     (match eval'(e1, env) with
      | IntVal v1 -> (match eval' (e2, env) with
		      | IntVal v2 -> IntVal(v1 + v2)
		      | _ -> RuntimeError)
      | _ -> RuntimeError)
  | App(e1, e2) ->
     (match eval'(e1, env) with
      | FunVal f -> (match eval' (e2, env) with
		     | RuntimeError -> RuntimeError		       
		     | v -> f v
		    )
      | _ -> RuntimeError)     
  | Lam(x, e) -> FunVal(fun z -> eval'(e, fun s -> if x = s then z else env s))
  | If(e1, e2, e3) ->
     (match eval'(e1, env) with
      | BoolVal b -> if b then eval' (e2, env) else eval'(e3, env)
      | _ -> RuntimeError)			  

  let init_env s = RuntimeError
					   
  let eval e = eval' (e, init_env)

  let trace_eval e = () 
end

module Deep_examples = DeepExample(Deep)
let t0_deep = Deep.eval Deep_examples.t0
let t1_deep = Deep.eval Deep_examples.t1
let t2_deep = Deep.eval Deep_examples.t2
let t3_deep = Deep.eval Deep_examples.t3
let t4_deep = Deep.eval Deep_examples.t4
let t5_deep = Deep.eval Deep_examples.t5									

(* Next, we apply a CPS translation on eval' 

This means we need to change types .... 
*) 			
type cps_value = 
  | CPS_RuntimeError 
  | CPS_IntVal of int
  | CPS_BoolVal of bool
  | CPS_FunVal of ((cps_value * cnt) -> cps_value)
  and cnt = cps_value -> cps_value 
			
module Deep_CPS : DEXP with type ast = simple_ast and type result = cps_value = 
  struct

  type ast = simple_ast     
  type result = cps_value
  type trace = unit
		 
  let var s = Var s 
  let int x = Int x
  let bool b = Bool b 
  let add x y = Add (x, y) 
  let lam x e = Lam(x, e)
  let app e1 e2 = App (e1, e2) 
  let if_ e1 e2 e3 = If(e1, e2, e3)

  (* note that eval' is tail recursive *) 		       
  let rec eval' (e, env, cnt) =
  match e with
  | Var x -> cnt (env x)
  | Int n -> cnt (CPS_IntVal n)
  | Bool b -> cnt (CPS_BoolVal b)
  | Add(e1, e2) ->
     eval'(e1, env, fun r1 ->
      match r1 with 
      | CPS_IntVal v1 -> eval' (e2, env, fun r2 ->
		      match r2 with  
		      | CPS_IntVal v2 -> cnt (CPS_IntVal(v1 + v2))
		      | _ -> cnt CPS_RuntimeError)
      | _ -> cnt CPS_RuntimeError)
  | App(e1, e2) ->
     eval'(e1, env, fun r1 ->
      match r1 with 
      | CPS_FunVal f -> eval' (e2, env, fun r2 ->
		     match r2 with 
		     | CPS_RuntimeError -> cnt CPS_RuntimeError		       
		     | v -> f (v, cnt))
      | _ -> cnt CPS_RuntimeError)
  | Lam(x, e) -> cnt (CPS_FunVal(fun (z, c) -> eval'(e, (fun s -> if x = s then z else env s), c)))
  | If(e1, e2, e3) ->
     eval'(e1, env, fun r ->
      match r with 
      | CPS_BoolVal b -> if b then eval' (e2, env, cnt) else eval'(e3, env, cnt)
      | _ -> cnt CPS_RuntimeError)			  

  let init_env s = CPS_RuntimeError
  let init_cnt v = v 
					   
  let eval e = eval' (e, init_env, init_cnt)

  let trace_eval e = () 		     
end
			
module Deep_CPS_examples = DeepExample(Deep_CPS)
let t0_deep_cps = Deep_CPS.eval Deep_CPS_examples.t0
let t1_deep_cps = Deep_CPS.eval Deep_CPS_examples.t1
let t2_deep_cps = Deep_CPS.eval Deep_CPS_examples.t2
let t3_deep_cps = Deep_CPS.eval Deep_CPS_examples.t3
let t4_deep_cps = Deep_CPS.eval Deep_CPS_examples.t4
let t5_deep_cps = Deep_CPS.eval Deep_CPS_examples.t5											


(* Defunctionalisation 

Look at the line above : 

  | Add(e1, e2) ->
     eval'(e1, env, fun r1 ->
      match r1 with 
      | CPS_IntVal v1 -> eval' (e2, env, fun r2 ->
		      match r2 with  
		      | CPS_IntVal v2 -> cnt (CPS_IntVal(v1 + v2))
		      | _ -> cnt CPS_RuntimeError)
      | _ -> cnt CPS_RuntimeError)

Notice that this introduces two continuations: 

 CNT1(v1, cnt) = fun r2 -> match r2 with | CPS_IntVal v2 -> cnt (CPS_IntVal(v1 + v2)) | _ -> cnt CPS_RuntimeError

 CNT2(e2, cnt) = fun r1 -> match r1 with | CPS_IntVal v1 -> eval' (e2, env, CNT1(v1, cnt)) | _ -> cnt CPS_RuntimeError

So we can replace these functions with instances of a new datatype : 

 type cnt = CNT1 of cps_value * cnt | CNT2 of simple_ast * cnt | ... 

and and apply function (mutually recursive with eval') such as 

   apply(CNT1(v1, cnt), CPS_IntVal v2) -> apply(cnt, CPS_IntVal(v1 + v2))
   apply(CNT2(e2, cnt), CPS_IntVal v1) -> eval' (e2, env, CNT1(v1, cnt))		      
   ... 

Note also that all of the cnt constructors will contain a continuation, so the constructors
look a lot like many flavours of list cons. So let's treat the continuations as lists! 

This requires another modification of types... 				
 *) 
type cps_defun_value = 
  | CPS_Defun_RuntimeError 
  | CPS_Defun_IntVal of int
  | CPS_Defun_BoolVal of bool
  | CPS_Defun_FunVal of ((cps_defun_value * cnt) -> cps_defun_value)
and cnt_element =
  | ADD  of simple_ast * env 
  | ADD2 of int 
  | IF   of simple_ast * simple_ast * env 
  | ARG  of simple_ast * env 
  | APP  of ((cps_defun_value * cnt) -> cps_defun_value) 
and cnt = cnt_element list 								 
and env = string -> cps_defun_value		       									  

module Deep_CPS_Defun : DEXP with type ast = simple_ast and type result = cps_defun_value = 
  struct

  type ast = simple_ast     
  type result = cps_defun_value
  type trace = unit
		 
  let var s = Var s 
  let int x = Int x
  let bool b = Bool b 
  let add x y = Add (x, y) 
  let lam x e = Lam(x, e)
  let app e1 e2 = App (e1, e2) 
  let if_ e1 e2 e3 = If(e1, e2, e3)

  (* Note that eval' and apply are mutually tail recursive *) 
  let rec eval' (e, env, c) =
  match e with
  | Var x          -> apply(c, env x)
  | Int n          -> apply(c, CPS_Defun_IntVal n)
  | Bool b         -> apply(c, CPS_Defun_BoolVal b)
  | Add(e1, e2)    -> eval'(e1, env, ADD(e2, env)::c)
  | App(e1, e2)    -> eval'(e1, env, ARG(e2, env)::c)
  | Lam(x, e)      -> apply(c, CPS_Defun_FunVal(fun (z, k) -> eval'(e, (fun s -> if x = s then z else env s), k)))
  | If(e1, e2, e3) -> eval'(e1, env, IF(e2, e3, env)::c) 
		       
  and apply = function 
    | ([],                                   v) -> v
    | (_,               CPS_Defun_RuntimeError) -> CPS_Defun_RuntimeError
    | (ADD2(v1)::c,        CPS_Defun_IntVal v2) -> apply(c, CPS_Defun_IntVal(v1 + v2))
    | (ADD(e2, env)::c,    CPS_Defun_IntVal v1) -> eval' (e2, env, (ADD2 v1) ::c)
    | (IF(e2, e3, env)::c, CPS_Defun_BoolVal b) -> if b then eval' (e2, env, c) else eval'(e3, env, c)
    | (APP(f)::c,                            v) -> f (v, c)
    | ((ARG(e2, env)::c),   CPS_Defun_FunVal f) -> eval' (e2, env, (APP f) ::c)
    | (c,                                    _) -> apply(c, CPS_Defun_RuntimeError)

						 
  let init_env s = CPS_Defun_RuntimeError
					   
  let eval e = eval' (e, init_env, [])

  let trace_eval e = () 		     		     
end
			
module Deep_CPS_Defun_examples = DeepExample(Deep_CPS_Defun)
let t0_deep_cps_defun = Deep_CPS_Defun.eval Deep_CPS_Defun_examples.t0
let t1_deep_cps_defun = Deep_CPS_Defun.eval Deep_CPS_Defun_examples.t1
let t2_deep_cps_defun = Deep_CPS_Defun.eval Deep_CPS_Defun_examples.t2
let t3_deep_cps_defun = Deep_CPS_Defun.eval Deep_CPS_Defun_examples.t3
let t4_deep_cps_defun = Deep_CPS_Defun.eval Deep_CPS_Defun_examples.t4								    
let t5_deep_cps_defun = Deep_CPS_Defun.eval Deep_CPS_Defun_examples.t5

(* Now, let's defunctionalise apply and eval'!  

   This is done by introducing a new type 

  state =
  | Apply of cnt2 * cps_defun2_value
  | Eval of simple_ast * env2 * cnt2

*) 				  
type cps_defun2_value = 
  | CPS_Defun2_RuntimeError 
  | CPS_Defun2_IntVal of int
  | CPS_Defun2_BoolVal of bool
  | CPS_Defun2_FunVal of ((cps_defun2_value * cnt2) -> state)
and cnt2_element =
  | ADD'  of simple_ast * env2 
  | ADD2' of int 
  | IF'   of simple_ast * simple_ast * env2 
  | ARG'  of simple_ast * env2 
  | APP'  of ((cps_defun2_value * cnt2) -> state) 
and cnt2 = cnt2_element list 								 
and env2 = string -> cps_defun2_value		       									  
and state =
  | Apply of cnt2 * cps_defun2_value
  | Eval of simple_ast * env2 * cnt2

module Deep_CPS_Defun2 : DEXP with type ast = simple_ast and type result = cps_defun2_value and type trace = state list = 
  struct

  type ast = simple_ast     
  type result = cps_defun2_value
  type trace = state list 		  

  let var s = Var s 
  let int x = Int x
  let bool b = Bool b 
  let add x y = Add (x, y) 
  let lam x e = Lam(x, e)
  let app e1 e2 = App (e1, e2) 
  let if_ e1 e2 e3 = If(e1, e2, e3)

  (* this is starting to look like a "machine" with these steps: *) 		       
  let step = function
    | Apply(ADD2'(v1)::c,           CPS_Defun2_IntVal v2) -> Apply(c, CPS_Defun2_IntVal(v1 + v2))
    | Apply(ADD'(e2, env)::c,       CPS_Defun2_IntVal v1) -> Eval(e2, env, (ADD2' v1) ::c)
    | Apply(IF'(e2, _, env)::c,  CPS_Defun2_BoolVal true) -> Eval(e2, env, c)
    | Apply(IF'(_, e3, env)::c, CPS_Defun2_BoolVal false) -> Eval(e3, env, c) 								 
    | Apply(APP'(f)::c,                                v) -> f (v, c)
    | Apply((ARG'(e2, env)::c),      CPS_Defun2_FunVal f) -> Eval(e2, env, (APP' f) ::c)
    | Apply(c,                                         _) -> Apply(c, CPS_Defun2_RuntimeError)
      
    | Eval(Var x, env, c)          -> Apply(c, env x)
    | Eval(Int n, env, c)          -> Apply(c, CPS_Defun2_IntVal n)
    | Eval(Bool b, env, c)         -> Apply(c, CPS_Defun2_BoolVal b)
    | Eval(Add(e1, e2), env, c)    -> Eval(e1, env, ADD'(e2, env)::c)
    | Eval(App(e1, e2), env, c)    -> Eval(e1, env, ARG'(e2, env)::c)
    | Eval(Lam(x, e), env, c)      -> Apply(c, CPS_Defun2_FunVal(fun (z, k) -> Eval(e, (fun s -> if x = s then z else env s), k)))
    | Eval(If(e1, e2, e3), env, c) -> Eval(e1, env, IF'(e2, e3, env)::c)

  (* And here is the "driver" *) 					  
  let rec eval' = function
    | Apply([], v) -> v
    | Apply(_, CPS_Defun2_RuntimeError) -> CPS_Defun2_RuntimeError
    | s -> eval' (step s)
						 
  let init_env s = CPS_Defun2_RuntimeError
					   
  let eval e = eval' (Eval(e, init_env, []))

  let rec trace_eval'(s, l) =
    match s with 
    | Apply([], v) -> s :: l 
    | Apply(_, CPS_Defun2_RuntimeError) -> s :: l 
    | s -> trace_eval' (step s, s :: l)

  (* finally, something to trace! *) 		       
  let trace_eval e = List.rev (trace_eval' (Eval(e, init_env, []), []))
end
    
module Deep_CPS_Defun2_examples = DeepExample(Deep_CPS_Defun2)
let t0_deep_cps_defun2 = Deep_CPS_Defun2.eval Deep_CPS_Defun2_examples.t0
let t1_deep_cps_defun2 = Deep_CPS_Defun2.eval Deep_CPS_Defun2_examples.t1
let t2_deep_cps_defun2 = Deep_CPS_Defun2.eval Deep_CPS_Defun2_examples.t2
let t3_deep_cps_defun2 = Deep_CPS_Defun2.eval Deep_CPS_Defun2_examples.t3
let t4_deep_cps_defun2 = Deep_CPS_Defun2.eval Deep_CPS_Defun2_examples.t4
let t5_deep_cps_defun2 = Deep_CPS_Defun2.eval Deep_CPS_Defun2_examples.t5;; (* need the ;; here because of the directives below! *)

(* For tracing I recommend using these repl directives *)

#print_depth 10000000;;
#print_length 10000000;;
					      
let t4_deep_cps_defun2_trace = Deep_CPS_Defun2.trace_eval Deep_CPS_Defun2_examples.t4
let t5_deep_cps_defun2_trace = Deep_CPS_Defun2.trace_eval Deep_CPS_Defun2_examples.t5								    			
(* We still have a way to go. We want a function 

    compile : ast -> code 

   (for some type code) so that eval can be defined as 

      let eval e = eval' ( ..., compile e, ....) 

   How can we refactor the Deep_CPS_Defun2 implementation to arrive as such an implementation? 

 *)