open Absyn
open Format 

exception TODO
exception RuntimeError of string

let list_find_position (x:'a) (l:'a list) :int option = 
  let rec helper l i = 
    match l with 
        [] -> None
      | (hd::tl) -> if x = hd then Some i else 
          helper tl (i+1)
  in
    helper l 0


module type EVAL = 
  sig 
    type value 
    val eval_program : exp -> unit
  end 


module Ref : EVAL = 
struct 
  type value = 
      Num of int 
    | ProcVal of env * var list * exp 
    | Bool of bool 
    | RefVal of reference 
  and  env = Extend of var list * value array * env 
             | Empty
  and reference = Reference of int * value array
         
  (* --------- values ----------------------------*)
         
  let rec print_value (v:value) : unit = 
    match v with 
        Num i -> print_int i;
      | Bool b -> 
          if b then print_string "true" else 
            print_string "false"
      | ProcVal (env,v,str) -> print_string "<closure>"
      | RefVal r -> print_string "<reference>"
  let closure (ids:var list) (body:exp) (env:env) : value = 
         ProcVal (env, ids, body)
                
  (* ------------ References and environments ------------- *)


  let deref (Reference (i,vec)) : value = vec.(i)   
  let setref (Reference (i,vec)) (v:value) : unit =  vec.(i) <- v
  let extend_env (vars:var list) (vals:value list) (old_env:env) : env = 
         Extend(vars, Array.of_list vals, old_env)
                
  let extend_env_rec (proc_names:var list) (idss:var list list) 
    (bodies:exp list) (old_env:env) : env = 
    let num_procs = List.length proc_names             in 
    let array     = Array.make num_procs (Num 0)       in
    let new_env   = Extend (proc_names, array,old_env) in 
    let rec loop idss bodies n : unit = 
      match idss, bodies with 
          (ids::is), (body::bs)  -> 
            (array.(n) <- closure ids body new_env;
             loop is bs (n+1))
        | [], [] -> () 
        | _,_ -> failwith "BUG: Cannot happen" in
      begin
        loop idss bodies 0;
        new_env
      end 
                
  (* Unlike the old version of apply_env, when the location 
   * of the value is found, that is what is returned as 
   * a reference.  *)
  let rec apply_env_ref (env:env) (sym:var) : reference = 
    match env with 
        Empty -> raise (RuntimeError ("Unbound variable " ^ sym))
      | Extend (vars, vals, old_env) -> 
          (match list_find_position sym vars with 
               Some pos -> Reference (pos,vals)
             | None -> apply_env_ref old_env sym)
                         
  (* Deferencing that reference gives us the value *) 
  let apply_env (e:env) (v:var) = deref (apply_env_ref e v) 


  (* --------- Booleans --------------- *)                                   
  let is_true (v:value) :bool = 
    match v with 
        Bool b -> b
      | _ -> raise (RuntimeError "Invalid value used as bool")
                         
  let true_value :value  = Bool true
  let false_value :value = Bool false
 
(* ----------- Primitives ----------- *)

  let rec apply_prim (p:prim) (args:value list) = 
    match p,args with 
        Plus,  [Num v1;Num v2] -> Num (v1 + v2)
      | Times, [Num v1;Num v2] -> Num (v1 * v2)
      | Minus, [Num v1;Num v2] -> Num (v1 - v2)
      | And,   [Bool b1; Bool b2] -> Bool ( b1 && b2 )
      | Or,    [Bool b1; Bool b2] -> Bool (b1 || b2)
      | Not,   [Bool b1] -> Bool (not b1)
      | Add1, [Num v1] -> Num (v1 + 1)
      | Sub1, [Num v1] -> Num (v1 - 1)
      | Zero, [Num v1] -> if v1 = 0 then true_value else false_value
      | LT, [Num v1; Num v2] -> if v1 < v2 then true_value else false_value
      | GT, [Num v1; Num v2] -> if v1 > v2 then true_value else false_value
      | Equal, [Num v1;Num v2] -> 
          if (v1 = v2) then true_value  else false_value

      (* For problem 3.41 *)
      | Cell, [v] -> RefVal (Reference (0, Array.make 1 v))
      | Contents, [RefVal r] -> deref r
      | Setcell, [RefVal r; v] -> (setref r v; Num 1)

      (* For problem 3.43 *)
      | Deref, [RefVal r] -> deref r
      | Setref, [RefVal r; v] -> (setref r v; Num 1)

      | _, _ -> raise (RuntimeError "Unimplemented  primitive or invalid arguments")


(* ------------ Evaluator ---------------------- *)
                  
  let rec apply_procval (v:value) (args:value list)  : value = 
    match v with 
        ProcVal (env, ids, body) -> 
          if (List.length ids = List.length args) then 
            eval body (extend_env ids args env)
          else 
            raise (RuntimeError "Wrong number of arguments to procedure")
      | _ -> raise (RuntimeError "Attempt to apply non-procedure")

  and  eval (exp:exp) (env:env) :value = 
    match exp with 
        LitInt i -> Num i
      | LitBool b -> 
          if b then true_value 
          else false_value
      | Var v -> apply_env env v
      | Prim (p, rands) -> 
          let args = eval_rands rands env in
            apply_prim p args
      | If (test_exp, true_exp, false_exp) -> 
          if (is_true (eval test_exp env))
          then (eval true_exp env)
          else (eval false_exp env)
      | Let (vars,rands,body) -> 
          let args = eval_rands rands env in 
            eval body (extend_env vars args env)
      | Lambda (ids,body) ->            
          closure ids body env
      | App (rator,rands) -> 
          let proc = eval rator env in 
          let args = eval_rands rands env in 
            apply_procval proc args         
      | Letrec (proc_names, idss, bodies, letrec_body) -> 
          eval letrec_body (extend_env_rec proc_names idss bodies env)
      | Varassign (id, rhs) -> 
          begin
            setref 
              (apply_env_ref env id)
              (eval rhs env);
            Num 1
          end
      | Begin(exp, exps) -> 
          let v = eval exp env in
          let rec loop l = 
            match l with 
                [] -> v
              | [e] -> eval e env 
              | (hd :: tl) -> 
                  ignore (eval hd env) ;
                  loop tl in 
            loop exps
      (* ----- For problem 3.43 ----- *)
      | Ref var -> RefVal (apply_env_ref env var)
          
          
  and eval_rands (rands:exp list) (env:env): value list = 
    List.map (fun x -> eval x env) rands

  let eval_program (exp:exp) :unit= 
         print_value (eval exp Empty)
                
end

(******************  Store-passing interpreter ************************** *)

module Store : EVAL = 
struct 
  
  (* To emphasize that mutation is *not* allowed, 
   * all data structures are immutable. You should not 
   * change any of these type definitions. *)

  type loc = int

  type value = 
                Num of int 
         | ProcVal of env * var list * exp 
         | Bool of bool 
         | RefVal of loc
  and env = Extend of var list * loc list * env 
                                | Empty
  and store = { data:(loc * value) list; size : int }

  type answer = value * store
         
  (* --------- values ----------------------------*)
         
  let rec print_value (v:value) : unit = 
    match v with 
        Num i -> print_int i;
      | Bool b -> 
          if b then print_string "true" else 
            print_string "false"
      | ProcVal (env,v,str) -> print_string "<closure>"
      | RefVal r -> print_string "<reference>"

  let extend_env (vars:var list) (locs:loc list) (old_env:env) : env = 
    Extend(vars, locs, old_env)
                
  (* Unlike the old version of apply_env, when the location 
   * of the value is found, that is what is returned as 
  * a reference.  *)
  let rec apply_env_ref (env:env) (sym:var) : loc = 
    match env with 
        Empty -> raise (RuntimeError ("Unbound variable " ^ sym))
      | Extend (vars, locs, old_env) -> 
          (match list_find_position sym vars with 
               Some pos -> List.nth locs pos                                      
             | None -> apply_env_ref old_env sym)
                         
  let closure (ids:var list) (body:exp) (env:env) : value = 
    ProcVal (env, ids, body)

  let new_loc (s:store) = 
         s.size
                                                                         
  let empty_store = { data = []; size = 0 }
                                                         
  (* lookup a location in the store *)
  let rec apply_store (s:store) (l:loc) :value = 
    let rec loop data : value = 
      match data with 
          ((l',value)::rest) -> 
            if l=l' then value else loop rest
        | [] -> raise (RuntimeError "unbound variable") in
      loop s.data
        
  (* add a new mapping to the store *)
  let extend_store (l:loc) (value:value) (s:store) : store = 
    if (l > s.size) 
    then failwith ("BUG: store invariant violated:")
    else 
      { data=(l,value)::s.data; size=s.size+1 }
                
  let extends_store (ls:loc list) (values:value list) (s:store) :store = 
    List.fold_left2 (fun s a b -> extend_store a b s) s ls values
         
  (* extend the store with locations for the new values, returning both
   * the store and the list of new locations that were added to the store 
   * for those values. *)
  let extend_vals (vals:value list) (s:store) : loc list * store = 
    List.fold_right
      (fun v (locs, newstore) -> 
         let loc = new_loc newstore in 
           (loc::locs, extend_store loc v newstore))
      vals
      ([],s)


  (* --------- Booleans --------------- *)
                                                                                                
  let is_true (v:value) :bool = 
    match v with 
        Bool b -> b
      | _ -> raise (RuntimeError "Invalid value used as bool")
                         
  let true_value :value  = Bool true
  let false_value :value = Bool false
 
  (* ----------- Primitives ----------- *)

  let rec apply_prim (p:prim) (args:value list) (s:store) : answer = 
    match p,args with 
        Plus,  [Num v1;Num v2] -> Num (v1 + v2), s
      | Times, [Num v1;Num v2] -> Num (v1 * v2), s
      | Minus, [Num v1;Num v2] -> Num (v1 - v2), s
      | And,   [Bool b1; Bool b2] -> Bool (b1 && b2), s 
      | Or,    [Bool b1; Bool b2] -> Bool (b1 || b2), s
      | Not,   [Bool b1] -> Bool (not b1), s
      | Add1, [Num v1] -> Num (v1 + 1), s
      | Sub1, [Num v1] -> Num (v1 - 1), s
      | Zero, [Num v1] -> (if v1 = 0 then true_value else false_value), s
      | LT, [Num v1; Num v2] -> (if v1 < v2 then true_value else false_value), s
      | GT, [Num v1; Num v2] -> (if v1 > v2 then true_value else false_value), s
      | Equal, [Num v1;Num v2] -> 
          (if (v1 = v2) then true_value  else false_value), s

      | Cell, [v] -> let l = new_loc s in (RefVal l, extend_store l v s)
      | Contents, [RefVal l] -> apply_store s l, s
      | Setcell, [RefVal l; v] -> (Num 1, extend_store l v s)

      | Deref, [RefVal l] -> (apply_store s l, s)
      | Setref, [RefVal l; v] -> (Num 1, extend_store l v s)

      | _, _ -> 
          raise (RuntimeError "Unimplemented  primitive or invalid arguments")

  (* The evaluator *)
  let rec apply_procval (v:value) (args:value list) (s:store) : answer = 
    match v with 
        ProcVal (env, ids, body) -> 
          if (List.length ids = List.length args) then 
            let locs,newstore2 = extend_vals args s in 
              eval_store body (extend_env ids locs env) newstore2
          else 
            raise (RuntimeError "Wrong number of arguments to procedure")
      | _ -> raise (RuntimeError "Attempt to apply non-procedure")

  and eval_store (e:exp) (env:env) (s:store) : answer = 
    match e with
        LitInt i -> (Num i,s)
      | LitBool b -> (Bool b, s)
      | Var v -> (apply_store s (apply_env_ref env v), s)
      | Prim (p, rands) -> 
          let (args, newstore) = eval_rands_store rands env s in
            apply_prim p args s
      | If (test_exp, true_exp, false_exp) -> 
          let (newval, newstore) = eval_store test_exp env s
          in if is_true newval
            then eval_store true_exp env newstore  
              (* use newstore instead of store *)
            else eval_store false_exp env newstore
      | Let (vars,rands,body) -> 
          let args,newstore1 = eval_rands_store rands env s in 
          let locs,newstore2 = extend_vals args newstore1 in 
            eval_store body (extend_env vars locs env) newstore2
      | Lambda (ids,body) ->            
          closure ids body env, s
      | App (rator,rands) -> 
          let proc, ns = eval_store rator env s in 
          let args, ns2 = eval_rands_store rands env ns in 
            apply_procval proc args s       
      | Varassign (id, rhs) -> 
          let value,ns = eval_store rhs env s in 
          let loc = apply_env_ref env id in
            (Num 1, extend_store loc value ns)
      | Begin(exp, exps) -> 
          let (val1, store1) = eval_store exp env s in 
          let (vals, store2) = eval_rands_store exps env store1 in 
            (match List.rev vals with 
                 [] -> (val1, store1)
               | (last::rest)  -> (last, store2))
      | Letrec (proc_names, idss, bodies, letrec_body) -> 
          (* allocate the closures *)
          let locs,newstore = extend_vals (List.map (fun _ -> Num 1) proc_names) s in 
          let newenv = extend_env proc_names locs env in 
          let newstore2 = extends_store locs 
                            (List.map2 (fun ids body -> 
                                          closure ids body newenv) idss bodies) 
                            newstore in 
            eval_store letrec_body newenv newstore2
      | Ref v -> (RefVal (apply_env_ref env v), s)


  (* evaluate a list of expressions, using the resulting store from
     each expression for the next one in the list. *)
  and eval_rands_store (rands:exp list) (env:env) (s:store) : value list * store = 
    match rands with 
        [] -> ([], s)
      | (hd::tl) -> 
          let (v, store1) = eval_store hd env s in 
          let (vs, store2) = eval_rands_store tl env store1 in 
            (v::vs, store2)

  let eval_program (exp:exp) : unit = 
    print_value (fst (eval_store exp Empty empty_store))
end