fun nth 0 (x::xs) = x
  | nth n (x::xs) = nth (n-1) xs;



fun dpopt n [] ys = <0>
  | dpopt n [0] ys = <0>
  | dpopt n [1] ys = <nth ~(lift n) ~ys>
  | dpopt n [x] ys = < ~(lift x) * (nth ~(lift n) ~ys) >
  | dpopt n (0::xs) ys = <~(dpopt (n+1) xs ys)>
  | dpopt n (1::xs) ys = <(nth ~(lift n) ~ys) + ~(dpopt (n+1) xs ys)>
  | dpopt n (x::xs) ys =
       <(~(lift x) * (nth ~(lift n) ~ys)) + ~(dpopt (n+1) xs ys)>;

fun gen n xs = <fn ys => ~(dpopt n xs <ys>)>;

val ans0 = gen 5 [2,0,1,0,4];

(* rule 1: x+0 = x *)
(* rule 2: 0*x = 0 *)
(* rule 3: 1*x = x *)

fun opt <fn x => ~(g <x>) + 0> = opt <fn y => ~(g <y>)>
  | opt <fn x => ~(g <x>) + 0 + ~(h <x>)> = opt <fn y => ~(g <y>) + ~(h <y>)>
  | opt <fn x => 0 + ~(g <x>)> = opt <fn y => ~(g <y>)>
  | opt <fn x => 0 * ~(g <x>) + ~(h <x>)> = opt <fn y => ~(h <y>)>
  | opt <fn x => ~(e <x>) + 0 * ~(g <x>)> = opt <fn y => ~(e <y>)>
  | opt <fn x => ~(e <x>) + 0 * ~(g <x>) + ~(h <x>)> =
        opt <fn y => ~(e <y>) + ~(h <y>)>
  | opt <fn x => 1 * ~(g <x>) + ~(h <x>)> = opt <fn y => ~(g <y>) + ~(h <y>)>
  | opt <fn x => ~(e <x>) + 1 * ~(g <x>)> = opt <fn y => ~(e <y>) + ~(g <y>)>
  | opt <fn x => ~(e <x>) + 1 * ~(g <x>) + ~(h <x>)> =
        opt <fn y => ~(e <y>) + ~(g <y>) + ~(h <y>)>
  | opt x = x;

(* ************************************************************ *)


datatype Exp =
  Var of string                     
| Add of Exp*Exp                    
| Mult of Exp*Exp                   
| Const of int                      
| Local of string*Exp*Exp;          



fun lookup s [] = error "unknown variable"
  | lookup s ((x,y)::xs) = if s=x then y else lookup s xs;

fun calc term bindings =
case term of
  Var s => lookup s bindings
| Add(x,y) => (calc x bindings) + (calc y bindings)
| Mult(x,y) => (calc x bindings) * (calc y bindings)
| Const n => n
| Local(s,x,y) => calc y ((s,calc x bindings)::bindings);

(***  a staged versions of calc with type:  
      Exp -> (string* <int>)list -> <int>  ***)


fun calc2 term bindings =
case term of
  Var s => lookup s bindings
| Add(x,y) => < ~(calc2 x bindings) + ~(calc2 y bindings) >
| Mult(x,y) => < ~(calc2 x bindings) * ~(calc2 y bindings) >
| Const n => lift n
| Local(s,x,y) => calc2 y ((s,calc2 x bindings)::bindings);

val term = Local("x",Add(Const 3, Const 4),
                 Local("y",Mult(Const 6,Var "x"),Add(Var "y",Const 12)));

val ans2 = calc2 term [];
   
fun calc3 term bindings =
case term of
  Var s => lookup s bindings
| Add(x,y) => < ~(calc3 x bindings) + ~(calc3 y bindings) >
| Mult(x,y) => < ~(calc3 x bindings) * ~(calc3 y bindings) >
| Const n => lift n
| Local(s,x,y) => <let val w = ~(calc3 x bindings)
                   in ~(calc3 y ((s,<w>)::bindings)) end>;

val ans3 = calc3 term [];

(************** Staging an interpreter ********************** *)


datatype Exp = 
  Constant of int               (*   5     *)
| Variable of string            (*   x     *)
| Minus of (Exp * Exp)          (*  x - 5  *)
| Greater of (Exp * Exp)        (*  x > 1  *)
| Times of (Exp * Exp) ;        (*  x * 4   *)


datatype Com =
  Assign of (string * Exp)   (*  x := 1                       *)
| Seq of (Com * Com)           (*  { x := 1; y := 2 }           *)
| If of (Exp * Com * Com)    (*  if x then x := 1 else y := 1 *) 
| While of (Exp * Com)         (*  while x>0 do x := x - 1      *)
| Declare of (string * Exp * Com);  
                               (*  Declare x = 1 in x := x - 1  *)




(* ***** environments ****** *)

exception Error;

fun error s = (print s; raise Error);

type environment = string list * int list ;

fun lookup x [] = error ("variable not found: "^x)
  | lookup x ((y,v)::zs) = if x=y then v else lookup x zs;

fun ext nm v zs = (nm,v)::zs;

fun remove (z::zs) = zs;

fun set name v [] = error ("name not found: "^name)
  | set name v ((z as (y,_))::zs) =
     if name=y then (y,v)::zs else z::(set name v zs);



(* ************* interpreters ************** *)

type env = (string * int) list;

fun eval0 exp env  = 
case exp of
  Constant n => n
| Variable s => lookup s env
| Minus(x,y) =>   let val a = eval0 x env 
                      val b = eval0 y env
                  in a - b end
| Greater(x,y) => let val a = eval0 x env
                      val b = eval0 y env
                  in if a '>' b then 1 else 0 end
| Times(x,y) =>   let val a = eval0 x env
                      val b = eval0 y env
                  in a * b end;
				
  
fun interpret0 stmt env =
case stmt of
  Assign(name,e) => 
   let val v = eval0 e env 
   in set name v env end 
| Seq(s1,s2) => 
   let val env1 = interpret0 s1 env 
       val env2 = interpret0 s2 env1 
   in env2 end 
| If(e,s1,s2) =>
   let val x = eval0 e env
   in if x=1 
	     then interpret0 s1 env
	     else interpret0 s2 env
   end
| While(e,body) => 
   let val v = eval0 e env
   in if v=0 
         then env
         else interpret0 (While(e,body)) 
                 (interpret0 body env)
   end
| Declare(nm,e,stmt) =>
   let val v = eval0 e env
       val env1 = ext nm v env 
   in remove(interpret0 stmt env1) end;   


(********** Split the enviroment into index plus stack *************** *)

type location = int;

type index = string list;

type value = int;

type stack = value list;

fun get 1 (x::xs) = x
  | get 0 _ = error "No value at index 0."
  | get n (x::xs) = get (n-1) xs
  | get n [] = error "Stack is empty";
fun put 1 v (x::xs) = (v::xs)
  | put 0 v _ = error "No value at index 0."
  | put n v (x::xs) = x :: (put (n-1) v xs)
  | put n v [] = error "Stack is empty";

(* ******************************************************** *)


fun pos name index = 
    let fun p n (x::xs) = if x=name then n else p (n+1) xs
    in p 0 index end


fun eval1 exp index stack = 
case exp of
  Constant n => n
| Variable s => get (pos s index) stack
| Minus(x,y) =>   let val a = eval1 x index stack 
                      val b = eval1 y index stack
                  in a - b end
| Greater(x,y) => let val a = eval1 x index stack
                      val b = eval1 y index stack
                  in if a '>' b then 1 else 0 end
| Times(x,y) =>   let val a = eval1 x index stack
                      val b = eval1 y index stack
                  in a * b end;
				


fun interp1 stmt index stack =
case stmt of
  Assign(name,e) => 
   let val v = eval1 e index stack 
       val loc = pos name index
   in put loc v stack end 
| Seq(s1,s2) => 
   let val stack1 = interp1 s1 index stack 
       val stack2 = interp1 s2 index stack1 
   in stack2 end 
| If(e,s1,s2) =>
   let val x = eval1 e index stack
   in if x=1 
	     then interp1 s1 index stack
	     else interp1 s2 index stack
   end
| While(e,body) => 
   let val v = eval1 e index stack
   in if v=0 
         then stack
         else interp1 (While(e,body)) index
                 (interp1 body index stack)
   end
| Declare(nm,e,stmt) =>
   let val v = eval1 e index stack
       val stack1 = v :: stack 
   in tl (interp1 stmt (nm::index) stack1) end; 

(******** Add staging annotations ************* *)

fun eval2 exp index stack = 
case exp of
  Constant n => lift n
| Variable s => <get ~(lift (pos s index)) ~stack>
| Minus(x,y) => <let val a = ~(eval2 x index stack) 
                     val b = ~(eval2 y index stack)
                 in a - b end>
| Greater(x,y) => <let val a = ~(eval2 x index stack)
                       val b = ~(eval2 y index stack)
                   in if a '>' b then 1 else 0 end>
| Times(x,y) => <let val a = ~(eval2 x index stack)
                     val b = ~(eval2 y index stack)
                 in a * b end>;
				




fun interp2 stmt index stack =
case stmt of
  Assign(name,e) => 
   <let val v = ~(eval2 e index stack) 
    in put ~(lift (pos name index)) v ~stack end>
| Seq(s1,s2) => 
   <let val stack1 = ~(interp2 s1 index stack) 
        val stack2 = ~(interp2 s2 index <stack1>) 
    in stack2 end> 
| If(e,s1,s2) =>
   <let val x = ~(eval2 e index stack)
    in if x=1 
	  then ~(interp2 s1 index stack)
	  else ~(interp2 s2 index stack)
   end>
| While(e,body) => 
   <let val v = ~(eval2 e index stack)
    in if v=0 
          then ~stack
          else ~(interp2 (While(e,body)) index
                  (interp2 body index stack))
    end>
| Declare(nm,e,stmt) =>
   <let val v = ~(eval2 e index stack)
        val stack1 = v :: ~stack 
    in tl ~(interp2 stmt (nm::index) <stack1>) end>; 




val s0 =
Declare("x",Constant 150,
  Declare("y",Constant 200,
    Seq(Assign("x",Minus(Variable "x",Constant 1)),
        Assign("y",Minus(Variable "y",Constant 1)))));

val ans2 = <fn stack => ~(interp2 s0 [] <stack>)>;


val s1 =
Declare("x",Constant 150,
  Declare("y",Constant 200,
    While(Greater(Variable "x",Constant 0),
          Seq(Assign("x",Minus(Variable "x",Constant 1)),
              Assign("y",Minus(Variable "y",Constant 1))))));

(* causes infinite loop *)
(* val ans3 = <fn stack => ~(interp2 s1 [] <stack>)>; *)



fun interp2 stmt index stack =
case stmt of
  Assign(name,e) => 
   <let val v = ~(eval2 e index stack) 
    in put ~(lift (pos name index)) v ~stack end>
| Seq(s1,s2) => 
   <let val stack1 = ~(interp2 s1 index stack) 
        val stack2 = ~(interp2 s2 index <stack1>) 
    in stack2 end> 
| If(e,s1,s2) =>
   <let val x = ~(eval2 e index stack)
    in if x=1 
	  then ~(interp2 s1 index stack)
	  else ~(interp2 s2 index stack)
   end>
| While(e,body) => 
   <let fun loop stk0 =
           let val v = ~(eval2 e index <stk0>)
           in if v=0 
                 then stk0
                 else let val stk1 = ~(interp2 body index <stk0>)
                      in loop stk1 end
           end
    in loop ~stack end>
| While(e,body) => 
   <let val v = ~(eval2 e index stack)
    in if v=0 
          then ~stack
          else ~(interp2 (While(e,body)) index
                  (interp2 body index stack))
    end>
| Declare(nm,e,stmt) =>
   <let val v = ~(eval2 e index stack)
        val stack1 = v :: ~stack 
    in tl ~(interp2 stmt (nm::index) <stack1>) end>; 

val ans3 = <fn stack => ~(interp2 s1 [] <stack>)>;