(*

 Looking up variables in the environment by name is slow. We can
 pre-compute the location of a variable binding in the "stack" of
 values that is the environment.

 We break our interpreter into two parts: a compiler, and a "bytecode"
 interpreter. Bytecodes are still large, so they look a lot like source
 expressions, except that variable names are replaced with stack
 indices.

*)

type xvar = string

type xval = Num of int
          | Fun of (xval -> xval)

(* The input syntax is unchanged *)

type xpr = Const of int
         | Minus of xpr * xpr
         | Times of xpr * xpr
         | Lam of xvar * xpr
         | Var of xvar
         | App of xpr * xpr
         | IfZero of xpr * xpr * xpr

type xenv = Empty
          | Extend of xval * xenv  (* No more variables at run time *)

(* "Bytecodes", which are the output of the compiler: *)

type cxpr = CConst of int
          | CMinus of cxpr * cxpr
          | CTimes of cxpr * cxpr
          | CLam of cxpr          (* No var anymore *)
          | CVar of int           (* Compiled to stack offset *)
          | CApp of cxpr * cxpr
          | CIfZero of cxpr * cxpr * cxpr

(* New environment type for compile-time: *)

type cenv = CEmpty
          | CExtend of xvar * cenv

;;

(* The examples are unchanged *)

let five = Const(5)
let protofac = Lam("f",
		   Lam("n",
		       IfZero(Var("n"),
			      Const(1),
			      Times(Var("n"),
				    App(App(Var("f"), Var("f")),
				        Minus(Var("n"),
					      Const(1)))))))
let fac = App(protofac, protofac)
let onetwenty = App(fac, five)

;;

(* The compiler *)

let rec
comp = function (Const(v), e) -> CConst(v)
              | (Minus(m1, m2), e) -> CMinus(comp(m1, e), comp(m2, e))
              | (Times(m1, m2), e) -> CTimes(comp(m1, e), comp(m2, e))
              | (Lam(var, m), e) -> CLam(comp(m, CExtend(var, e)))
              | (App(m1, m2), e) -> CApp(comp(m1, e), comp(m2, e))
              | (IfZero(m1, m2, m3), e) -> CIfZero(comp(m1, e), comp(m2, e), comp(m3, e))
              | (Var(var), e) -> CVar(offset(var, e))

and
offset = fun (var, CExtend(var2, e)) -> if (var = var2) 
					   then 0 
					   else (1 + offset(var, e))

;;

(* The interpreter, now with index-based lookup *)

let rec
eval = function (CConst(v), e) -> Num(v)
              | (CMinus(m1, m2), e) -> let Num(n1) = eval(m1, e)
                                       and Num(n2) = eval(m2, e)
                                        in Num(n1 - n2)
              | (CTimes(m1, m2), e) -> let Num(n1) = eval(m1, e)
                                       and Num(n2) = eval(m2, e)
                                        in Num(n1 * n2)
              | (CLam(m), e) -> Fun(fun v -> eval(m, Extend(v, e)))
              | (CApp(m1, m2), e) -> let Fun(f) = eval(m1, e)
                                      in f(eval(m2, e))
              | (CIfZero(m1, m2, m3), e) -> let Num(n) = eval(m1, e)
                                             in eval((if (n=0) 
                                                         then m2
                                                         else m3),
                                                     e)
              | (CVar(n), e) -> lookup(n, e)

and
lookup = function (0, Extend(v, e)) -> v
                | (n, Extend(v, e)) -> lookup(n-1, e)

;;

eval(comp(onetwenty,CEmpty), Empty)

;;