Untitled

open SmallCTypes
open Utils

exception DeclareError
exception TypeError
exception DivByZeroError

(* Lookup a variable's value in the environment *)
let rec lookup env x =
  match env with
  | [] -> raise DeclareError
  | (k, v)::t -> if k = x then v else lookup t x

(* Evaluate expressions *)
let rec eval_expr env = function
  | Int n -> Int_Val n
  | Bool b -> Bool_Val b
  | ID x -> lookup env x

  | Add (a, b) ->
      (match eval_expr env a, eval_expr env b with
       | Int_Val x, Int_Val y -> Int_Val (x + y)
       | _ -> raise TypeError)

  | Sub (a, b) ->
      (match eval_expr env a, eval_expr env b with
       | Int_Val x, Int_Val y -> Int_Val (x - y)
       | _ -> raise TypeError)

  | Mult (a, b) ->
      (match eval_expr env a, eval_expr env b with
       | Int_Val x, Int_Val y -> Int_Val (x * y)
       | _ -> raise TypeError)

  | Div (a, b) ->
      (match eval_expr env a, eval_expr env b with
       | Int_Val _, Int_Val 0 -> raise DivByZeroError
       | Int_Val x, Int_Val y -> Int_Val (x / y)
       | _ -> raise TypeError)

  | Pow (a, b) ->
      (match eval_expr env a, eval_expr env b with
       | Int_Val x, Int_Val y -> Int_Val (int_of_float (float_of_int x ** float_of_int y))
       | _ -> raise TypeError)

  | Greater (a, b) ->
      (match eval_expr env a, eval_expr env b with
       | Int_Val x, Int_Val y -> Bool_Val (x > y)
       | _ -> raise TypeError)

  | Less (a, b) ->
      (match eval_expr env a, eval_expr env b with
       | Int_Val x, Int_Val y -> Bool_Val (x < y)
       | _ -> raise TypeError)

  | GreaterEqual (a, b) ->
      (match eval_expr env a, eval_expr env b with
       | Int_Val x, Int_Val y -> Bool_Val (x >= y)
       | _ -> raise TypeError)

  | LessEqual (a, b) ->
      (match eval_expr env a, eval_expr env b with
       | Int_Val x, Int_Val y -> Bool_Val (x <= y)
       | _ -> raise TypeError)

  | Equal (a, b) ->
      (match eval_expr env a, eval_expr env b with
       | Int_Val x, Int_Val y -> Bool_Val (x = y)
       | Bool_Val x, Bool_Val y -> Bool_Val (x = y)
       | _ -> raise TypeError)

  | NotEqual (a, b) ->
      (match eval_expr env a, eval_expr env b with
       | Int_Val x, Int_Val y -> Bool_Val (x <> y)
       | Bool_Val x, Bool_Val y -> Bool_Val (x <> y)
       | _ -> raise TypeError)

  | And (a, b) ->
      (match eval_expr env a, eval_expr env b with
       | Bool_Val x, Bool_Val y -> Bool_Val (x && y)
       | _ -> raise TypeError)

  | Or (a, b) ->
      (match eval_expr env a, eval_expr env b with
       | Bool_Val x, Bool_Val y -> Bool_Val (x || y)
       | _ -> raise TypeError)

  | Not a ->
      (match eval_expr env a with
       | Bool_Val x -> Bool_Val (not x)
       | _ -> raise TypeError)

(* Evaluate statements *)
let rec eval_stmt env = function
  | NoOp -> env

  | Seq (s1, s2) ->
      let env' = eval_stmt env s1 in
      eval_stmt env' s2

  | Declare (typ, name) ->
      if List.mem_assoc name env then raise DeclareError
      else
        let default =
          match typ with
          | Int_Type -> Int_Val 0
          | Bool_Type -> Bool_Val false
        in
        (name, default) :: env

  | Assign (name, e) ->
      let value = eval_expr env e in
      let old_val =
        try Some (lookup env name)
        with DeclareError -> None
      in
      match old_val, value with
      | None, _ -> raise DeclareError
      | Some (Int_Val _), Int_Val _ -> (name, value) :: List.remove_assoc name env
      | Some (Bool_Val _), Bool_Val _ -> (name, value) :: List.remove_assoc name env
      | _ -> raise TypeError

  | Print e ->
      let v = eval_expr env e in
      (match v with
       | Int_Val n -> print_output_int n
       | Bool_Val b -> print_output_bool b);
      print_output_newline ();
      env

  | If (cond, s_then, s_else) ->
      (match eval_expr env cond with
       | Bool_Val true -> eval_stmt env s_then
       | Bool_Val false -> eval_stmt env s_else
       | _ -> raise TypeError)

  | While (cond, body) ->
      let rec loop env =
        match eval_expr env cond with
        | Bool_Val true -> loop (eval_stmt env body)
        | Bool_Val false -> env
        | _ -> raise TypeError
      in
      loop env

  | For (name, start_expr, end_expr, body) ->
      let start_val = eval_expr env start_expr in
      let end_val = eval_expr env end_expr in
      match start_val, end_val with
      | Int_Val s, Int_Val e ->
          let rec loop i env =
            if i > e then env
            else
              let env_with_i = (name, Int_Val i) :: List.remove_assoc name env in
              let env_after_body = eval_stmt env_with_i body in
              loop (i + 1) env_after_body
          in
          loop s env
      | _ -> raise TypeError