(** This needs MetaOCaml -- see the exercise sheet for installation instructions: *)

(** The EXP signature.  This gives the type of the various interpreter
   modules.

   There is an abstract type 't' for terms, and signatures for
   functions that create values of type 't'.

   While 't' is abstract here, each module will define 't' in a
   different way to implement a different type of interpretation:
   evaluation, compilation, partial evaluation, cps transformation,
   etc. *)
module type EXP = sig
  type 'a t
  val int : int -> int t
  val add : int t -> int t -> int t 
  val bool : bool -> bool t
  val if_ : bool t -> 'a t -> 'a t -> 'a t
  val lam : ('a t -> 'b t) -> ('a -> 'b) t
  val app : ('a -> 'b) t -> 'a t -> 'b t
  val print : int t -> unit t
end

(** Now that we have the interface, we can write some example terms *)
module Example (E: EXP) =
struct
  open E
  (** The double application function λf.λx.f (f x) *)
  let t0 : ((int -> int) -> int -> int) E.t = 
      (lam (fun f ->
           lam (fun x ->
               app f (app f x))))

  (** An application of t0 to the identity *)
  let t1 = app t0 (lam (fun y -> y))

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

  (** A term that behaves differently under CBN and CBV
      (λb.λe.if b then e else print 1) false (print 2)

      CBV: print 2; print 1
      CBN: print 1
 *)
  let t3 =
    app
      (app
         (lam (fun b ->
              (lam (fun e ->
                   if_ b
                     e
                     (print (int 1))))))
         (bool false))
      (print (int 2))
end

(** A very simple interpretation of EXP: almost the identity,
    but we use thunks so that we can control the evaluation of
    'if' rather than always evaluating both branches *)
module Eval : EXP with type 'a t = unit -> 'a =
struct
  type 'a t = unit -> 'a
  let int x = fun () -> x
  let bool b = fun () -> b
  let add x y = fun () -> x () + y ()
  let lam f = fun () -> (fun x -> f (fun () -> x) ())
  let app f x () = f () (x ())
  let if_ b t e () = if b () then t () else e ()
  let print i () = print_int (i ())
end
(** Evaluate the examples using Eval *)
module Eval_examples = Example(Eval)
let t0_eval = Eval_examples.t0
let t1_eval = Eval_examples.t1
let t2_eval = Eval_examples.t2
let t3_eval = Eval_examples.t3

(** There follow some auxiliary definitions for a tracing
   interpretation.  Rather than executing the effects (i.e. the uses
   of 'print') in the interpreter, as Eval does, evalauting a term
   using the tracing interpretation will simply collect a trace of the
   effects that would be executed *)

(** A type of terms for the tracing interpreter.
    Each constructor pairs a value with a list of effects
    that the program executed during the construction of the value. *)
type effects = int list
type 'a trace =
  Fun : ('a trace -> 'b trace) * effects -> ('a -> 'b) trace
| Int : int * effects -> int trace
| Bool : bool * effects -> bool trace
| Unit : unit * effects -> unit trace

(** 'effects' retrieves the list of effects that were executed 
    during the construction of a value *)
let effects : type a. a trace -> effects = function
  | Fun (_, es) -> es
  | Int (_, es) -> es
  | Bool (_, es) -> es
  | Unit (_, es) -> es

(** 'cleartrace' removes the list of effects that were executed 
    during the construction of a value, leaving just the value *)
let cleartrace : type a. a trace -> a trace = function
  | Fun (v, _) -> Fun (v, [])
  | Int (v, _) -> Int (v, [])
  | Bool (v, _) -> Bool (v, [])
  | Unit (v, _) -> Unit (v, [])

(** '>>' composes two programs, catenating the lists of their effects.
    The result of 'a >> b' is the result of 'b', and the effects
    are the effects of 'a' followed by the effects of 'b'.
    That is, 'a >> b' behaves like 'a; b' in an imperative language. *)
let (>>) : type a b. a trace -> b trace -> b trace =
  fun x y ->
  match y with
  | Fun (v, es) -> Fun (v, effects x @ es)
  | Int (v, es) -> Int (v, effects x @ es)
  | Bool (v, es) -> Bool (v, effects x @ es)
  | Unit (v, es) -> Unit (v, effects x @ es)

(** A CBV tracing implementation of EXP. *)
module Trace : EXP with type 'a t = 'a trace =
struct
  type 'a t = 'a trace
  let int x = Int (x, [])
  let bool x = Bool (x, [])
  let add (Int (x,es)) (Int (y,es')) = Int (x + y, es @ es')
  (** (a): complete the remaining functions, replacing
      'assert false' with your implementation.

      Note: use CBV application order
 *)
  let lam f                       = assert false (* (a) TODO *)
  let app f                       = assert false (* (a) TODO *)
  let if_ (Bool _) t e            = assert false (* (a) TODO *)
  let print (Int (i, es))         = assert false (* (a) TODO *)
end

module TraceCBN : EXP with type 'a t = 'a trace =
struct
  type 'a t = 'a trace
  let int x = Int (x, [])
  let bool x = Bool (x, [])
  let add (Int (x,es)) (Int (y,es')) = Int (x + y, es @ es')
  (** (a): complete the remaining functions, replacing
      'assert false' with your implementation

      Note: use CBN application order *)
  let lam f                       = assert false (* (b) TODO *)
  let app f                       = assert false (* (b) TODO *)
  let if_ (Bool _) t e            = assert false (* (b) TODO *)
  let print (Int (i, es))         = assert false (* (b) TODO *)
end

(** Naive compilation using MetaOCaml's brackets and escapes  *)
module Compile : EXP with type 'a t = 'a code =
struct
  type 'a t = 'a code
  let int (x:int) = .<x>.
  let bool (x:bool) = .<x>.
  let add x y = .< .~x + .~y >.
  let lam f = .< fun x -> .~(f .<x>.) >.
  let app f x = .< .~f .~x >.
  let if_ b t e = .< if .~b then .~t else .~e >.
  let print e = .< print_int .~e >.
end
(** Compile the examples using Compile *)
module Compile_examples = Example(Compile)
let t0_compile = Compile_examples.t0
let t1_compile = Compile_examples.t1
let t2_compile = Compile_examples.t2
let t3_compile = Compile_examples.t3

(** A simple partial evaluator, similar to the one shown in lectures *)
type _ static1 =
  | Int : int -> int static1
  | Bool : bool -> bool static1
  | Unknown : _ static1
type 'a sd1 = {
    sta: 'a static1;
    dyn: 'a code;
  }
let rec resid1 : type a. a sd1 -> a code = function
  | { sta = Unknown; dyn } -> dyn
  | { sta = Int x } -> .< x >.
  | { sta = Bool x } -> .< x >.

module PE_basic : EXP with type 'a t = 'a sd1 =
struct
  (** (c) complete the implementation of PE *)
  type 'a t = 'a sd1
  let int (i : int) = { sta = Int i; dyn = .<i>. }
  let bool (b : bool) = { sta = Bool b; dyn = .<b>. }
  let add x y = match x.sta, y.sta with
    | Int 0, Int _ -> y
    | Int _, Int 0 -> x
    | Int l, Int r -> let s = l + r in
                      { sta = Int s; dyn = .< s >. }
    | _ -> { sta = Unknown; dyn = .< .~(x.dyn) + .~(y.dyn) >. }
  let lam f   = let dyn = .< fun x -> .~(resid1 (f {sta = Unknown; dyn = .<x>.})) >.
                in { sta = Unknown; dyn = dyn }
  let app f x = { sta = Unknown; dyn = .< .~(resid1 f) .~(resid1 x) >. }
  let print e = { sta = Unknown; dyn = .< print_int .~(resid1 e) >. }
  let if_ b t e =
    match b.sta with
      Bool true -> t
    | Bool false -> e
    | Unknown -> let dyn = .< if .~(b.dyn) then .~(resid1 t) else .~(resid1 e) >. in
                 { sta = Unknown; dyn = dyn }
end
				 

(** Next we'll build a higher-order partial evaluator.  The type 'sd'
   is a representation of partially-static data, where each value is a
   pair of an optional static (known) value and a dynamic (unknown)
   value. *)
type _ static =
  | Int : int -> int static
  | Bool : bool -> bool static
  | Fun (* (c) TODO: give Fun an appropriate type. *)
  | Unknown : _ static
and 'a sd = {
    sta: 'a static;
    dyn: 'a code;
  }
(** The 'resid' function converts partially-static values into
   fully dynamic values — i.e. into code.

   We favour using the static portions over the dynamic portions
   wherever possible, for the best (i.e. most-evaluated) possible
   results *)
let rec resid : type a. a sd -> a code = function
  | { sta = Unknown; dyn } -> dyn
  | { sta = Int x } -> .< x >.
  | { sta = Bool x } -> .< x >.
  | { sta = Fun; dyn } -> assert false (* (c) TODO *)

(** Here's a starting point for the higher-order partial evaluator —
   an interpretation of EXP where the term type is instantiated to sd
   *)
module PE : EXP with type 'a t = 'a sd =
struct
  (** (c) complete the implementation of PE *)
  type 'a t = 'a sd
  let int (i : int) = { sta = Int i; dyn = .<i>. }
  let bool (b : bool) = { sta = Bool b; dyn = .<b>. }
  let add x y = match x.sta, y.sta with
    | Int 0, Int _ -> y
    | Int _, Int 0 -> x
    | Int l, Int r -> let s = l + r in
                      { sta = Int s; dyn = .< s >. }
    | _ -> { sta = Unknown; dyn = .< .~(x.dyn) + .~(y.dyn) >. }
  let lam f   = assert false (* (c) TODO *)
  let app f   = assert false (* (c) TODO *)
  let print e = assert false (* (c) TODO *)
  let if_ b t e =
    match b.sta with
      Bool true -> t
    | Bool false -> e
    | Unknown -> assert false (* (c) TODO *)
end
(** Partially-evaluate the example terms.  *)
(* module PE_examples = Example(PE)
 * let t0_pe =  resid PE_examples.t0
 * let t1_pe =  resid PE_examples.t1
 * let t2_pe =  resid PE_examples.t2
 * let t3_pe =  resid PE_examples.t3 *)

(** The 'equalp' function checks whether two 'static' values
    are equal.  If it is not possible to determine
    equality then 'equalp' returns 'Unknown'. *)
type equal = Yes | No | Unknown
let rec equalp : type a. a static -> a static -> equal =
  fun l r -> assert false (* (d) TODO *)

module PE2 : EXP with type 'a t = 'a sd =
struct
  (* (d) complete the implementation of PE2.
     Your implementation should make appropriate use of 'equalp',
     but it should be otherwise equivalent to PE *)
  type 'a t = 'a sd
  let int (i : int) = { sta = Int i; dyn = .<i>. }
  let bool (b : bool) = { sta = Bool b; dyn = .<b>. }
  let add x y = match x.sta, y.sta with
    | Int 0, Int _ -> y
    | Int _, Int 0 -> x
    | Int l, Int r -> let s = l + r in
                      { sta = Int s; dyn = .< s >. }
    | _ -> { sta = Unknown; dyn = .< .~(x.dyn) + .~(y.dyn) >. }
  let lam f   = assert false (* (d) TODO *)
  let app f   = assert false (* (d) TODO *)
  let print e = assert false (* (d) TODO *)
  let if_ b t e =
    match b.sta with
      Bool true -> t
    | Bool false -> e
    | Unknown -> assert false (* (d) TODO *)
end
(* module PE2_examples = Example(PE2)
 * let t0_pe2 =  resid PE2_examples.t0
 * let t1_pe2 =  resid PE2_examples.t1
 * let t2_pe2 =  resid PE2_examples.t2
 * let t3_pe2 =  resid PE2_examples.t3 *)


(** A CPS interpretation of EXP, where each term
    is interpreted as a continuation-accepting function.

    Making 'cps' a record type makes it possible to use higher-order
    polymorphism --- i.e. to pass around polymorphic functions
    as arguments and return values.

    The type of cps is equivalent to   ∀b. ('a → b) → b    *)
type 'a cps = {k: 'b. ('a -> 'b) -> 'b}
module CPS : EXP with type 'a t = 'a cps =
struct
  (** (e) complete the implementation of CPS. *)
  type 'a t = 'a cps
  let int i =
    {k=fun k -> k i}
  let bool b =
    {k=fun k -> k b}
  let add x y =
    assert false (* (e) TODO *)
  let lam f =
    assert false (* (e) TODO *)
  let app f x =
    assert false (* (e) TODO *)
  let if_ c t e =
    assert false (* (e) TODO *)
  let print e =
    assert false (* (e) TODO *)
end
(** Interpret the example terms using CPS *)
(* module CPS_examples = Example(CPS) *)
(** Pass in the identity continuation to recover the values *)
(* let t0_cps =  CPS_examples.t0
 * let t1_cps =  CPS_examples.t1
 * let t2_cps =  CPS_examples.t2
 * let t3_cps =  CPS_examples.t3 *)

(** Another version of CPS, parameterized by a second interpretation
   of EXP.  Rather than ultimately evaluating in OCaml and passing the
   results to the continuation, we'll evaluate using the second
   interpretation, E. *)
module CPS2 (E: EXP) :
sig
  type 'a cps_ = {k: 'b. ('a E.t -> 'b E.t) -> 'b E.t}
  include EXP with type 'a t = 'a cps_
end =
struct
  type 'a cps_ = {k: 'b. ('a E.t -> 'b E.t) -> 'b E.t}
  type 'a t = 'a cps_
  let int i =
    {k=fun k -> k (E.int i)}
  let bool b =
    {k=fun k -> k (E.bool b)}
  let add x y =
    assert false (* (f) TODO *)
  let lam f =
    assert false (* (f) TODO *)
  let app f x =
    assert false (* (f) TODO *)
  let if_ c t e =
    assert false (* (f) TODO *)
  let print e =
    assert false (* (f) TODO *)
end

(** With the parameterized CPS interpretation we can combine
    CPS and PE2.  Are the results any better than PE2 alone? *)
(* module CPS_PE = CPS2(PE) *)
(** Let's interpret the terms using CPS_PE, and look at the results. *)
(* module CPS_PE_examples = Example(CPS_PE) *)
(* let t0_cps_pe =  CPS_PE_examples.t0
 * let t1_cps_pe =  CPS_PE_examples.t1
 * let t2_cps_pe =  CPS_PE_examples.t2
 * let t3_cps_pe =  CPS_PE_examples.t3 *)

(** Finally, a normalizing implementation with some constraints
    on the form of generated terms (for which see the exercise sheet) *)
type 'a normalized =
  | Return : 'a value -> 'a normalized
  | Let : 'b effect * ('b value -> 'a normalized) -> 'a normalized
  | If : bool value * 'a normalized * 'a normalized -> 'a normalized
and 'a value =
  | Var : 'a code -> 'a value
  | Int : int -> int value
  | Bool : bool -> bool value
  | Lam : ('a value -> 'b normalized) -> ('a -> 'b) value
and 'a effect (* = *)
   (* (g) TODO *)

let rec residn : type a. a normalized -> a code = function
    Return v -> residnv v
  | Let (e, k) -> .< let x = .~(residne e) in .~(residn (k (Var .<x>.))) >.
  | If (c, t, e) -> .< if .~(residnv c) then .~(residn t) else .~(residn e) >.
and residnv : type a. a value -> a code = function
  | Var x -> x
  | Int i -> .<i>.
  | Bool b -> .<b>.
  | Lam f -> .< fun x -> .~(residn (f (Var .<x>.))) >.
and residne : type a. a effect -> a code = function
    _ -> assert false (* TODO *)

module Normal : EXP with type 'a t = 'a normalized =
struct
  type 'a t = 'a normalized

  (* Hint: write a function '>>=' with the given type
     and use it to combine values of type 't' in your
     implementations of the functions below. *)
  let rec (>>=) : type a b. a t -> (a value -> b t) -> b t =
    fun m k -> assert false (* (g) TODO, if you like *)

  let int x = Return (Int x)
  let bool x = Return (Bool x)
  let app f =
    assert false (* (g) TODO *)
  let lam f =
    assert false (* (g) TODO *)
  let if_ c t e =
    assert false (* (g) TODO *)
  let add x y =
    assert false (* (g) TODO *)
  let print e =
    assert false (* (g) TODO *)
end