
(* Title:   Staging a Featherweight Java interpreter
   Author:  Walid Taha and Valery Trifonov
   Date:    Sat Oct 13 17:00:01 EDT 2001 *)

Trx.init_times ()

exception Yiikes

type ('a,'b) map = ('a * 'b) list

let empty = []

let (++) x l = x::l 

let list2map l = l (* i.e. foldr (op ++) empty *)

let rec lookup' k' x =
  match x with 
    ((k,v)::l) -> if k=k' then Some v else lookup' k' l
  | _          -> None

let lookup k env = .<lookup' k .~env>.

let rec zip (a,b) =
  match (a,b) with
    (x::xs, y::ys) -> (x,y) :: zip (xs,ys)
  | _              -> []

(* Syntax *)

type var       = string (* including "this" *)
type className = string
type methName  = string
type fieldName = string

type exp
  = Var of var 
  | Inv of exp * methName * exp list 
  | Sel of exp * fieldName
  | New of className * exp list

type methDef    = var list * exp
type classDecl  = className * fieldName list * (methName, methDef) map
type classTable = (className, classDecl) map
type program    = classTable * exp

(* Values *)

type value = Obj of className * value list

let rec lookupMeth' t m c =
    match lookup' c t with 
      Some (super,_,meths) -> 
    match lookup' m meths with
          Some x -> x
	| None   -> lookupMeth' t m super

let lookupMeth t m c = .<lookupMeth' .~t m .~c>.

let rec lookupField' t f (c,vs) =
    match lookup' c t with 
      Some (super,fields,_) ->
    let rec search (a,b) = 
        match (a,b) with
          (n::ns, v::vs) -> if f=n then v else search (ns,vs)
        | ([],vs)        -> lookupField' t f (super,vs)
    in search (fields,vs)

let lookupField t f v = .<lookupField' .~t f .~v>.

let fst (x,y) = x
let snd (x,y) = y

let rec map'' f l =
  match l with 
    []      -> []
  | (x::xs) -> (f x) :: map'' f xs

let rec map' f l =
  match l with 
    []      -> .<[]>.
  | (x::xs) -> .< .~(f x) :: .~(map' f xs) >.

let lift x = .<x>.

let rec intE' tv = 
    let interpT' env =
	let rec interpTE' =
	   fun a -> match a with 
             Var x -> 
               let Some v = (lookup' x env) in v
	    | Inv (e,m,es) -> 
		  let v = (interpTE' e)
               in let c = match v with (Obj (c,_)) -> c
	       in let vs = (map'' interpTE' es)
		  and z = (lookupMeth' (!tv) m c)
	          in 
                     let (xs,lv') = z
                      in lv' 
                         (("this",v) ++ (zip (xs,vs)))
                              
	    | Sel (e,f)    -> 
		  let Obj (a,b) = (interpTE' e)
	          in (lookupField' (!tv) f (a,b))
	    | New (c,es)   -> Obj (c, (map'' interpTE' es))
        in interpTE'
    in interpT'
and intT' tv0 t =  
         let cp (cn, (sn,fnl,mt)) = ((cn),
                                     ((sn),
                                      (fnl),(
                let mtc' (mn,(vl,e)) = 
                      ((mn),
                       ((vl), 
                       fun env -> (intE' tv0 env e)))
                in map'' mtc' mt)))
         in map'' cp t

let rec intE tv = 
    let interpT env =
	let rec interpTE =
	   fun a -> match a with 
             Var x -> 
               .<match .~(lookup x env) with Some v -> v | _ -> assert false>.
	    | Inv (e,m,es) -> .<
		  let v = .~(interpTE e)
               in let c = match v with (Obj (c,_)) -> c
	       in let vs = .~(map' interpTE es)
		  and z = .~(lookupMeth (.< ! .~tv>.) m .<c>.)
	          in 
                     let (xs,lv') = z
                      in lv' 
                         (("this",v) ++ (zip (xs,vs)))
                              >.
	    | Sel (e,f)    -> .<
		  let Obj (a,b) = .~(interpTE e)
	          in .~(lookupField (.< ! .~tv>.) f .<(a,b)>.)>.
	    | New (c,es)   -> .<Obj (c, .~(map' interpTE es))>.
        in interpTE
    in interpT
and intT tv0 topenv =  
         let cp (cn, (sn,fnl,mt)) = .<(cn,
                                     (sn,
                                      fnl,.~(
                let mtc (mn,(vl,e)) = 
                      .<(mn,
                       (vl, 
                       fun env -> .~(intE tv0 .<env>. e)))>.
                in map' mtc mt)))>.
         in map' cp topenv

let interp' topenv env e = let temp = ref (Obj.magic 1)
                           in
                           let rec tv0 () = (intT' temp topenv)
                           in let _ = (temp := (tv0 ()))
                              in (intE' temp env e)

let interp topenv env e = .<let temp = ref (Obj.magic 1)
                            in
                            let rec tv0 () = .~(intT .<temp>. topenv)
                            in let _ = (temp := (tv0 ()))
                               in .~(intE .<temp>. env e)>.

(* Examples *)

let myObject = ("",[],empty)

(* class thunk ext myObject { myObject val; myObject get () { this.let } } *)
let thunk = ("myObject",
	     ["val"],
	     list2map [("get", ([], Sel (Var "this", "val")))])

(* class bool ext myObject { ; myObject if (thunk x, thunk y) { x.get() } } *)
let bool = ("myObject",
	    [],
	    list2map [("if", (["x";"y"], Inv (Var "x", "get", [])))])

(* class true ext bool {} *)
let myTrue = ("bool", [], empty)

(* class myFalse ext bool { ; myObject if (thunk x, thunk y) { y.get() } } *)
let myFalse = ("bool",
	     [],
	     list2map [("if", (["x";"y"], Inv (Var "y", "get", [])))])

(* interface numCont ext myObject { num app (num arg) } *)

(* -- modulo some casts --
   class num ext myObject { ;
     myObject ifz (thunk z, numCont s) { z.get() }
     myObject iter (thunk z, numCont s) { z.get() }
     num add (num x) { this.iter (new thunk (x), new addC ()) }
     num mul (num x) { this.iter (new thunk (new zero ()), new mulC (x)) }
     num exp (num x) { x.iter (new thunk (new succ (new zero ())),
			       new expC (this)) }
   }
   class addC ext myObject imp numCont { ; num app (num x) { new succ (x) } }
   class mulC ext myObject imp numCont {
     num cand;
     num app (num x) { this.cand.add (x) }
   }
   class expC ext myObject imp numCont {
     num base;
     num app (num x) { this.base.mul (x) }
   }
   class zero ext num {}
   class succ ext num {
     num pred;
     myObject ifz (thunk z, numCont s) { s.app (this.pred) }
     myObject iter (thunk z, numCont s) { s.app (this.pred.iter (z,s)) }
   }
   *)

let num = ("myObject",
	   [],
	   list2map [
	     ("ifz", (["z";"s"], Inv (Var "z", "get", [])));
	     ("iter", (["z";"s"], Inv (Var "z", "get", [])));
	     ("add", (["x"], Inv (Var "this",
				  "iter",
				  [New ("thunk", [Var "x"]);
				   New ("addC", [])])));
	     ("mul", (["x"], Inv (Var "this",
				  "iter",
				  [New ("thunk", [New ("zero", [])]);
				   New ("mulC", [Var "x"])])));
	     ("exp", (["x"], Inv (Var "x",
				  "iter",
				  [New ("thunk",
					[New ("succ", [New ("zero", [])])]);
				   New ("expC", [Var "this"])])))])

let addC = ("myObject",
	    [],
	    list2map [("app", (["x"], New ("succ", [Var "x"])))])

let mulC = ("myObject",
	    ["cand"],
	    list2map [("app", (["x"], Inv (Sel (Var "this", "cand"),
					   "add",
					   [Var "x"])))])

let expC = ("myObject",
	    ["base"],
	    list2map [("app", (["x"], Inv (Sel (Var "this", "base"),
					    "mul",
					    [Var "x"])))])

let zero = ("num", [], empty)

let succ = ("num",
	    ["pred"],
	    list2map [
              ("ifz", (["z";"s"],
		       Inv (Var "s", "app", [Sel (Var "this", "pred")])));
	      ("iter", (["z";"s"],
			Inv (Var "s",
			     "app",
			     [Inv (Sel (Var "this", "pred"),
				   "iter",
				   [Var "z"; Var "s"])])))])

let topenv = list2map [
  ("myObject",myObject); ("thunk",thunk); ("bool",bool); ("myTruetrue",myTrue);
  ("myFalse",myFalse); ("num",num); ("zero",zero); ("succ",succ);
  ("addC",addC); ("mulC",mulC); ("expC",expC)]


let eval' e = interp' topenv empty e

let eval e = interp topenv .<empty>. e

let newObj c l = New (c,l)
let ( <-* ) x m = fun l -> Inv (x,m,l)
let ( <<* ) x f = Sel (x,f)

let rec count o =
  match o with
    (Obj ("zero",[]))   -> 0
  | (Obj ("succ", [v])) -> count v + 1

(* ; *)

let zero = newObj "zero" []
let one = newObj "succ" [zero]
let two = (one <-* "add") [one]
let three = (two <-* "add") [one]

(* Examples *)

(* #use "mex/7.ml";; *)

let unstagedRunning = 
     Trx.time 10
        "unstaged running"
        (fun () -> count (eval' ((one <-* "exp") [((((two <-* "exp") [two]) <-* "exp") [three])])))

let stage1 = 
     Trx.time 10
        "stage 1"
        (fun () -> eval ((one <-* "exp") [((((two <-* "exp") [two]) <-* "exp") [three])]))

let compilation = 
     Trx.time 1
        "compilation"
        (fun () -> .! .<fun () -> .~stage1>.)

let stage2 = 
     Trx.time 10
        "stage 2 (compiling and running)"
        (fun () -> count (compilation ()))

(* A second experiment *)

let unstagedRunning = 
     Trx.time 1
        "unstaged running"
        (fun () -> count (eval' ((one <-* "exp") [((((three <-* "exp") [two]) <-* "exp") [three])])))

let stage1 = 
     Trx.time 1
        "stage 1"
        (fun () -> eval ((one <-* "exp") [((((three <-* "exp") [two]) <-* "exp") [three])]))

let compilation = 
     Trx.time 1
        "compilation"
        (fun () -> .! .<fun () -> .~stage1>.)

let stage2 = 
     Trx.time 1
        "stage 2 (compiling and running)"
        (fun () -> count (compilation ()))

(* A third experiment *)

let unstagedRunning = 
     Trx.time 1
        "unstaged running"
        (fun () -> count (eval' ((((three <-* "exp") [three]) <-* "exp") [three])))

let stage1 = 
     Trx.time 1
        "stage 1"
        (fun () -> eval ((((three <-* "exp") [three]) <-* "exp") [three]))

let compilation = 
     Trx.time 1
        "compilation"
        (fun () -> .! .<fun () -> .~stage1>.)

let stage2 = 
     Trx.time 1
        "stage 2 (compiling and running)"
        (fun () -> count (compilation ()))

let _ = Trx.print_times ()

(*

Types before and after staging:

val interp' :
  (className *
   (className * fieldName list * (methName * (var list * exp)) list))
  list -> (var * value) list -> exp -> value = <fun>
val interp :
  (className *
   (className * fieldName list * (methName * (var list * exp)) list))
  list -> (var * value) list code -> exp -> value code = <fun>

*)