(**************************************************************
 *  CS101: Homework 2 --- Type checker and interpreter        *
 **************************************************************)

open Cs101_hw2_types

module VarMap = Map.Make(String)

type type_environment = lc_type VarMap.t

type value_environment = (value_environment * lc_exp) VarMap.t

(*
 * Lambda-calculus type inference.
 * Raises LcError if the input is not well-typed.
 *)
let rec lc_get_type gamma e =
   match e with
      ExpInt _ ->
         TyInt
    | ExpVar v ->
         begin 
            try VarMap.find v gamma
            with
               Not_found -> raise(LcError(ErrUnknownVariable v))
         end
    | ExpBinop (_, e1, e2) ->
         begin match lc_get_type gamma e1, lc_get_type gamma e2 with
           TyInt, TyInt -> TyInt
         | TyInt, t2 -> raise (LcError(ErrTypeMismatch(e2, t2, TyInt)))
         | t1, _ -> raise (LcError((ErrTypeMismatch(e1, t1, TyInt))))
        end
    | ExpApply (e1, e2) ->
         let t1 = lc_get_type gamma e1 in
         let t2 = lc_get_type gamma e2 in
            begin match t1 with
               TyFun(t_in, t_out) -> 
                  if t2 = t_in then
                     t_out
                  else 
                     raise (LcError((ErrTypeMismatch(e2, t2, t_in))))
             | _ -> raise (LcError((ErrTypeFunApp(e1, t1))))
            end
    | ExpLambda (v, t, e) -> 
        let ty_res = lc_get_type (VarMap.add v t gamma) e in
           TyFun(t, ty_res)

(*
 * Lambda-calculus type checker.
 * Raises LcError if the input is not well-typed.
 *)
let lc_type_check =
   lc_get_type VarMap.empty

let decode_binop = function
   OpPlus -> (+)
 | OpMinus -> (-)
 | OpTimes -> ( * )

let rec lc_get_value gamma e =
   match e with 
      ExpInt _ ->
         VarMap.empty, e
    | ExpLambda _ ->
         gamma, e
    | ExpVar v -> 
         (*
          * The exression have passed the type checker,
          * we can assume that v is in gamma.
          *)
         VarMap.find v gamma
    | ExpBinop (op, e1, e2) ->
         begin match lc_get_value gamma e1, lc_get_value gamma e2 with
            (_, ExpInt i1), (_, ExpInt i2) -> VarMap.empty, ExpInt (decode_binop op i1 i2)
          | _ -> raise (InternalError "Non-int expressions in a binop")
         end
    | ExpApply (e1, e2) ->
         let gamma1, e1 = lc_get_value gamma e1 in
         let closure2 = lc_get_value gamma e2 in
            begin match e1 with
               ExpLambda(v, _, e1) -> 
                  lc_get_value (VarMap.add v closure2 gamma1) e1
             | _ ->
                  raise (InternalError "First exp of application isn't a function.")
            end

(*
 * takes in a closure and an expression, potentially containing free variables,
 * returs a closed expression where all the free variable were replaced
 * with corresponding values.
 *)
let rec expand_closures gamma e =
   match e with
      ExpInt _ ->
         e
    | ExpLambda (v, t, e) ->
         let gamma = VarMap.add v (VarMap.empty, ExpVar v) gamma in
            ExpLambda (v, t, expand_closures gamma e)
    | ExpVar v ->
         let gamma, e = VarMap.find v gamma in
            if VarMap.is_empty gamma then
               e
            else
               expand_closures gamma e
    | ExpBinop (op, e1, e2) ->
         ExpBinop (op, expand_closures gamma e1, expand_closures gamma e2)
    | ExpApply (e1, e2) ->
         ExpApply (expand_closures gamma e1, expand_closures gamma e2)

(*
 * Lambda-calculus interpreter.
 * The returned expression must be a value (a number or a lambda).
 * Raises LcError if the input is not well-typed.
 *)
let lc_evaluate e =
   ignore(lc_type_check e);
   let gamma, e = lc_get_value VarMap.empty e in
      expand_closures gamma e