
(*Semantic checking for the Pie-Num compiler*)

open Ast

module StringMap = Map.Make(String)

(* Semantic checking of a program. Returns void if successful,
   throws an exception if something is wrong.

   Check each global variable, then check each function *)

let check (globals, functions) =

(* Raise an exception if the given list has a duplicate *)
  let report_duplicate exceptf list =
    let rec helper = function
        n1 :: n2 :: _ when n1 = n2 -> raise (Failure (exceptf n1))
      | _ :: t -> helper t
      | [] -> ()
    in helper (List.sort compare list)
  in

  (* Raise an exception if a given binding is to a void type *)
  let check_not_void exceptf = function
      (Void, n) -> raise (Failure (exceptf n))
    | _ -> ()
  in  (*identical up to this function*)

  (* Raise an exception of the given rvalue type cannot be assigned to
     the given lvalue type *)
  let check_assign lvaluet rvaluet err =
     match (lvaluet, rvaluet) with
    (Int, Int) -> lvaluet
    | (Float, Float) -> lvaluet
    | (String, String) -> lvaluet
    | (Bool, Bool) -> lvaluet
    | (Void, Void) -> lvaluet
    | (Img, Img) -> lvaluet
    | (ArrayTyp(Int, l1), ArrayTyp(Int, l2)) -> if l1 == l2 then lvaluet else if l1 == 0 then lvaluet else raise err
    | (ArrayTyp(Float, l1), ArrayTyp(Float, l2)) -> if l1 == l2 then lvaluet else if l1 == 0 then lvaluet else raise err
    | (MatrixTyp(Int, r1, c1), MatrixTyp(Int, r2, c2)) -> if r1 == r2 && c1 == c2 then lvaluet else raise err
    | (MatrixTyp(Float, r1, c1), MatrixTyp(Float, r2, c2)) -> if r1 == r2 && c1 == c2 then lvaluet else raise err
    | _ -> raise err
  in

  (**** Checking Global Variables ****)

  List.iter (check_not_void (fun n -> "illegal void global " ^ n)) globals;

  report_duplicate (fun n -> "duplicate global " ^ n) (List.map snd globals);

  (**** Checking Functions. Add all newly declared functions here ****)


  if List.mem "printi" (List.map (fun fd -> fd.fname) functions)
  then raise (Failure ("Function printi may not be defined")) else ();

  if List.mem "prints" (List.map (fun fd -> fd.fname) functions)
  then raise (Failure ("Function prints may not be defined")) else ();

  if List.mem "printf" (List.map (fun fd -> fd.fname) functions)
  then raise (Failure ("Function printf may not be defined")) else ();

  if List.mem "read_image" (List.map (fun fd -> fd.fname) functions)
  then raise (Failure ("Function read_image may not be defined")) else ();

  if List.mem "leni3" (List.map (fun fd -> fd.fname) functions)
  then raise (Failure ("Function leni3 may not be defined")) else ();

  if List.mem "lenf3" (List.map (fun fd -> fd.fname) functions)
  then raise (Failure ("Function lenf3 may not be defined")) else ();

  if List.mem "leni33" (List.map (fun fd -> fd.fname) functions)
  then raise (Failure ("Function leni33 may not be defined")) else ();

  if List.mem "lenf33" (List.map (fun fd -> fd.fname) functions)
  then raise (Failure ("Function lenf33 may not be defined")) else ();

  if List.mem "to_int" (List.map (fun fd -> fd.fname) functions)
  then raise (Failure ("Function to_float may not be defined")) else ();

  if List.mem "to_float" (List.map (fun fd -> fd.fname) functions)
  then raise (Failure ("Function to_float may not be defined")) else ();

  if List.mem "write_image" (List.map (fun fd -> fd.fname) functions)
  then raise (Failure ("Function to_float may not be defined")) else ();

  if List.mem "printfil" (List.map (fun fd -> fd.fname) functions)
  then raise (Failure ("Function lenf33 may not be defined")) else ();

  if List.mem "printiil" (List.map (fun fd -> fd.fname) functions)
  then raise (Failure ("Function to_float may not be defined")) else ();

  if List.mem "printsil" (List.map (fun fd -> fd.fname) functions)
  then raise (Failure ("Function to_float may not be defined")) else ();

  if List.mem "sleep" (List.map (fun fd -> fd.fname) functions)
  then raise (Failure ("Function to_float may not be defined")) else ();


  report_duplicate (fun n -> "Duplicate function " ^ n)
    (List.map (fun fd -> fd.fname) functions);




    let built_in_decls =
        StringMap.add "printi"
        { typ = Void; fname = "printi"; formals = [(Int, "x")];
        locals = []; body = [] } (StringMap.add "prints"
        { typ = Void; fname = "prints"; formals = [(String, "x")];
        locals = []; body = [] } (StringMap.add "leni3"
        { typ = Int; fname = "leni3"; formals =[(ArrayTyp(Int,3), "x")];
        locals = []; body = [] }(StringMap.add "lenf3"
        { typ = Int; fname = "lenf3"; formals =[(ArrayTyp(Float,3), "x")];
        locals = []; body = [] }(StringMap.add "leni33"
        { typ = Int; fname = "leni33"; formals =[(MatrixTyp(Int,3,3), "x")];
        locals = []; body = [] }(StringMap.add "lenf33"
        { typ = Int; fname = "lenf33"; formals =[(MatrixTyp(Float,3,3), "x")];
        locals = []; body = [] }(StringMap.add "read_image"
        { typ = Img; fname = "read_image"; formals =[(String, "x")];
        locals = []; body = [] }(StringMap.add "to_float"
        { typ = Float; fname = "to_float"; formals =[(Int, "x")];
        locals = []; body = [] } (StringMap.add "to_int"
        { typ = Int; fname = "to_int"; formals =[(Float, "x")];
        locals = []; body = [] } (StringMap.add "write_image"
        { typ = Int; fname = "write_image"; formals =[(Img, "x");(Int, "x");(Int, "x");(String, "x")];
        locals = []; body = [] } (StringMap.add "printsil"
        { typ = Void; fname = "printsil"; formals =[(String, "x")];
        locals = []; body = [] } (StringMap.add "printiil"
        { typ = Void; fname = "printiil"; formals =[(Int, "x")];
        locals = []; body = [] }(StringMap.add "printfil"
        { typ = Void; fname = "printfil"; formals =[(Float, "x")];
        locals = []; body = [] }(StringMap.add "sleep"
        { typ = Int; fname = "sleep"; formals =[(Int, "x")];
        locals = []; body = [] }(StringMap.singleton "printf"
        { typ = Void; fname = "printf"; formals = [(Float,"x")];
        locals = []; body = [] }))))))))))))))
    in



    let function_decls =
        (List.fold_left (fun m fd -> StringMap.add fd.fname fd m) built_in_decls functions)
    in
    let function_decl s = try StringMap.find s function_decls
    	with Not_found -> raise (Failure ("Unrecognized function " ^ s))
    in

    let _ = function_decl "main" in


    let check_function func =
        List.iter (check_not_void (fun n -> "illegal void formal " ^ n ^
          " in " ^ func.fname)) func.formals;
        report_duplicate (fun n -> "duplicate formal " ^ n ^ " in " ^ func.fname)
          (List.map snd func.formals);
        List.iter (check_not_void (fun n -> "illegal void local " ^ n ^
          " in " ^ func.fname)) func.locals;
        report_duplicate (fun n -> "duplicate local " ^ n ^ " in " ^ func.fname)
          (List.map snd func.locals);

    (* Type of each variable (global, formal, or local *)(*Check variables*)
    let symbols = List.fold_left (fun m (t, n) -> StringMap.add n t m)
        StringMap.empty (globals @ func.formals @ func.locals )
    in

    let symbols = ref symbols in


let find_rowtyp name m =
        let m = StringMap.find m !symbols in
        let typ = match m with
                MatrixTyp(Int, _, _) -> Int
          | MatrixTyp(Float, _, _) -> Float
          | _ -> raise (Failure ("Illegal matrix type")) in
        let cols = match m with
                MatrixTyp(_, _, c) -> c
          | _ -> raise (Failure ("Illegal matrix type")) in
        symbols := StringMap.add name (ArrayTyp(typ, cols)) !symbols in

 let type_of_identifier s =
      try StringMap.find s !symbols
    with Not_found -> raise (Failure ("Undeclared identifier " ^ s))
   in

  let row_access_type = function
      ArrayTyp(r, _) -> r
    | _ -> raise (Failure ("Illegal row access")) in

  let matrix_access_type = function
      MatrixTyp(t, _, _) -> t
    | _ -> raise (Failure ("Illegal matrix access") ) in

  (*this may not be necessary*)
  let mrow_access_type = function
            MatrixTyp(t, _, c) -> ArrayTyp(t, c)
          | _ -> raise (Failure ("Illegal matrix access") ) in

  let type_of_row r l =
        match (List.hd r) with
        IntLiteral _ -> ArrayTyp(Int, l)
      | FloatLiteral _ -> ArrayTyp(Float, l)
      | _ -> raise (Failure ("Illegal row type"))
  in

  let type_of_matrix m r c =
    match (List.hd (List.hd m)) with
        IntLiteral _ -> MatrixTyp(Int, r, c)
      | FloatLiteral _ -> MatrixTyp(Float, r, c)
      | _ -> raise (Failure ("Illegal matrix type"))
  in

    let matrix_type s = match (List.hd s) with
    | IntLiteral _ -> ArrayTyp(Int, List.length s)
    | FloatLiteral _ -> ArrayTyp(Float, List.length s)
    | BoolLit _ -> ArrayTyp(Bool, List.length s)
    | _ -> raise ( Failure ("Cannot instantiate a matrix of that type")) in

    let rec check_all_matrix_literal m ty idx =
      let length = List.length m in
        match (ty, List.nth m idx) with
      (ArrayTyp(Int, _), IntLiteral _) -> if idx == length - 1 then ArrayTyp(Int, length) else check_all_matrix_literal m (ArrayTyp(Int, length)) (succ idx)
    | (ArrayTyp(Float, _), FloatLiteral _) -> if idx == length - 1 then ArrayTyp(Float, length) else check_all_matrix_literal m (ArrayTyp(Float, length)) (succ idx)
    | (ArrayTyp(Bool, _), BoolLit _) -> if idx == length - 1 then ArrayTyp(Bool, length) else check_all_matrix_literal m (ArrayTyp(Bool, length)) (succ idx)
    | _ -> raise (Failure ("Illegal matrix literal"))
  in


    let rec expr = function
    IntLiteral _ -> Int
  | FloatLiteral _ -> Float
  | StringLiteral _ -> String
  | BoolLit _ -> Bool
  | ID s -> type_of_identifier s
  | Null ->Void
  | ArrayLit r -> type_of_row r (List.length r)
  | MatrixLit m -> type_of_matrix m (List.length m) (List.length (List.hd m))
  | ArrayAccess(s, e) -> let _ = (match (expr e) with
                                    Int -> Int
                                  | _ -> raise (Failure ("Attempting to access with non-integer type"))) in
                            row_access_type (type_of_identifier s)
  | MatrixAccess(s, e1, e2) -> let _ = (match (expr e1) with
                                          Int -> Int
                                        | _ -> raise (Failure ("Attempting to access with a non-integer type")))
                               and _ = (match (expr e2) with
                                          Int -> Int
                                        | _ -> raise (Failure ("Attempting to access with a non-integer type"))) in
                               matrix_access_type (type_of_identifier s)

  | Binop(e1, op, e2) as e -> let t1 = expr e1 and t2 = expr e2 in
   (match op with
    Add -> (match t1,t2 with Int,Int -> Int
      | Float,Float -> Float
      | ArrayTyp(Int,l1),ArrayTyp(Int,l2) when l1=l2 -> ArrayTyp(Int,l1)
      | ArrayTyp(Int,l1), Int -> ArrayTyp(Int, l1)
      | Int, ArrayTyp(Int,l1) -> ArrayTyp(Int, l1)
      | ArrayTyp(Float,l1),ArrayTyp(Float,l2) when l1=l2 -> ArrayTyp(Float,l1)
      | ArrayTyp(Float,l1), Float -> ArrayTyp(Float, l1)
      | Float, ArrayTyp(Float,l1) -> ArrayTyp(Float, l1)
      | MatrixTyp(Int,r1,c1),MatrixTyp(Int,r2,c2) when r1=r2 && c1=c2 -> MatrixTyp(Int,r1,c1)
      | MatrixTyp(Int,r1,c1), Int -> MatrixTyp(Int,r1,c1)
      | Int, MatrixTyp(Int,r1,c1) -> MatrixTyp(Int,r1,c1)
      | MatrixTyp(Float,r1,c1),MatrixTyp(Float,r2,c2) when r1=r2 && c1=c2 -> MatrixTyp(Float,r1,c1)
      | MatrixTyp(Float,r1,c1), Float -> MatrixTyp(Float,r1,c1)
      | Float, MatrixTyp(Float,r1,c1) -> MatrixTyp(Float,r1,c1)
      | _,_ -> raise (Failure("Illegal addition operator in " ^ string_of_expr e1 ^ " + " ^ string_of_expr e2)))
    | Sub -> (match t1,t2 with Int,Int -> Int
      | Float,Float -> Float
      | ArrayTyp(Int,l1),ArrayTyp(Int,l2) when l1=l2 -> ArrayTyp(Int,l1)
      | ArrayTyp(Int,l1), Int -> ArrayTyp(Int, l1)
      | Int, ArrayTyp(Int,l1) -> ArrayTyp(Int, l1)
      | ArrayTyp(Float,l1),ArrayTyp(Float,l2) when l1=l2 -> ArrayTyp(Float,l1)
      | ArrayTyp(Float, l1), Float -> ArrayTyp(Float, l1)
      | Float, ArrayTyp(Float,l1) -> ArrayTyp(Float, l1)
      | MatrixTyp(Int,r1,c1),MatrixTyp(Int,r2,c2) when r1=r2 && c1=c2 -> MatrixTyp(Int,r1,c1)
      | MatrixTyp(Int,r1,c1), Int -> MatrixTyp(Int,r1,c1)
      | Int, MatrixTyp(Int,r1,c1) -> MatrixTyp(Int,r1,c1)
      | MatrixTyp(Float,r1,c1),MatrixTyp(Float,r2,c2) when r1=r2 && c1=c2 -> MatrixTyp(Float,r1,c1)
      | MatrixTyp(Float,r1,c1), Float -> MatrixTyp(Float,r1,c1)
      | Float, MatrixTyp(Float,r1,c1) -> MatrixTyp(Float,r1,c1)
      | _,_ -> raise (Failure("Illegal subtraction operator in "^ string_of_expr e1 ^ " - " ^ string_of_expr e2)))
    | Mult -> (match t1,t2 with Int,Int -> Int
      | Float,Float -> Float
      | ArrayTyp(Int,l1),ArrayTyp(Int,l2) when l1=l2 -> ArrayTyp(Int,l1)
      | ArrayTyp(Int,l1), Int -> ArrayTyp(Int, l1)
      | Int, ArrayTyp(Int,l1) -> ArrayTyp(Int, l1)
      | ArrayTyp(Float,l1),ArrayTyp(Float,l2) when l1=l2 -> ArrayTyp(Float,l1)
      | ArrayTyp(Float, l1), Float -> ArrayTyp(Float, l1)
      | Float, ArrayTyp(Float,l1) -> ArrayTyp(Float, l1)
      | Int, MatrixTyp(Int,r1,c1) -> MatrixTyp(Int,r1,c1)
      | MatrixTyp(Int,r1,c1), Int -> MatrixTyp(Int,r1,c1)
      | MatrixTyp(Int,r1,c1), MatrixTyp(Int,r2,c2) when r1=r2 && c1=c2 -> MatrixTyp(Int,r1,c1)
      | Float, MatrixTyp(Float,r1,c1) -> MatrixTyp(Float,r1,c1)
      | MatrixTyp(Float,r1,c1), Float -> MatrixTyp(Float,r1,c1)
      | MatrixTyp(Float,r1,c1), MatrixTyp(Float,r2,c2) when r1=r2 && c1=c2 -> MatrixTyp(Float,r1,c1)
      | _,_ -> raise (Failure("Illegal multiplication operator in " ^ string_of_expr e1 ^ " * " ^ string_of_expr e2)))
    | Div -> (match t1,t2 with Int,Int -> Int
      | Float,Float -> Float
      | ArrayTyp(Int,l1),ArrayTyp(Int,l2) when l1=l2 -> ArrayTyp(Int,l1)
      | ArrayTyp(Int,l1), Int -> ArrayTyp(Int, l1)
      | Int, ArrayTyp(Int,l1) -> ArrayTyp(Int, l1)
      | ArrayTyp(Float,l1),ArrayTyp(Float,l2) when l1=l2 -> ArrayTyp(Float,l1)
      | ArrayTyp(Float, l1), Float -> ArrayTyp(Float, l1)
      | Float, ArrayTyp(Float,l1) -> ArrayTyp(Float, l1)
      | Int, MatrixTyp(Int,r1,c1) -> MatrixTyp(Int,r1,c1)
      | MatrixTyp(Int,r1,c1), Int -> MatrixTyp(Int,r1,c1)
      | MatrixTyp(Int,r1,c1), MatrixTyp(Int,r2,c2) when r1=r2 && c1=c2 -> MatrixTyp(Int,r1,c1)
      | Float, MatrixTyp(Float,r1,c1) -> MatrixTyp(Float,r1,c1)
      | MatrixTyp(Float,r1,c1), Float -> MatrixTyp(Float,r1,c1)
      | MatrixTyp(Float,r1,c1), MatrixTyp(Float,r2,c2) when r1=r2 && c1=c2 -> MatrixTyp(Float,r1,c1)
      | _,_ -> raise (Failure("Illegal division operator in " ^ string_of_expr e1 ^ " / " ^ string_of_expr e2)))
    | Equal | Neq when t1 = t2 -> Bool
    | Less | Leq | Greater | Geq when t1 = Int && t2 = Int -> Bool
    | Less | Leq | Greater | Geq when t1 = Float && t2 = Float -> Float
    | And | Or when t1 = Bool && t2 = Bool -> Bool
    | _ -> raise (Failure ("Illegal binary operator " ^
      string_of_typ t1 ^ " " ^ string_of_op op ^ " " ^
      string_of_typ t2 ^ " in " ^ string_of_expr e))
  )
  | Unop(op, e) as ex -> let t = expr e in
  (match op with
      Neg when t = Int -> Int
    | Neg when t = Float -> Float
    | Not when t = Bool -> Bool
    | _ -> raise (Failure ("Illegal unary operator " ^ string_of_uop op ^
      string_of_typ t ^ " in " ^ string_of_expr ex)))
  | Noexpr -> Void
  | Assign(e1, e2) as ex -> let lt = (match e1 with
                                      | ArrayAccess(s, _) -> (match (type_of_identifier s) with
                                                                     ArrayTyp(t, _) -> (match t with
                                                                                                    Int -> Int
                                                                                                  | Float -> Float


                                                                                                  | _ -> raise ( Failure ("Illegal row") )
                                                                                                )
                                                                      | _ -> raise ( Failure ("Cannot access a primitive") )
                                                                   )
                                      | MatrixAccess(s, _, _) -> (match (type_of_identifier s) with
                                                                      MatrixTyp(t, _, _) -> (match t with
                                                                                                    Int -> Int
                                                                                                  | Float -> Float
                                                                                                  | _ -> raise ( Failure ("Illegal matrix of matrices") )
                                                                                                )
                                                                      | _ -> raise ( Failure ("Cannot access a primitive") )
                                                                   )
                                      | _ -> expr e1)
                            and rt = (match e2 with
                                      | ArrayAccess(s, _) -> (match (type_of_identifier s) with
                                                                     ArrayTyp(t, _) -> (match t with
                                                                                                    Int -> Int
                                                                                                  | Float -> Float
                                                                                                  | _ -> raise ( Failure ("Illegal row") )
                                                                                                )
                                                                      | _ -> raise ( Failure ("Cannot access a primitive") )
                                                                   )
                                      | MatrixAccess(s, _, _) -> (match (type_of_identifier s) with
                                                                      MatrixTyp(t, _, _) -> (match t with

                                                                                                  Int -> Int
                                                                                                | Float -> Float
                                                                                                | _ -> raise ( Failure ("Illegal matrix of matrices") )
                                                                                              )
                                                                      | _ -> raise ( Failure ("Cannot access a primitive") )
                                                                   )
                                      | _ -> expr e2) in
  check_assign lt rt (Failure ("Illegal assignment " ^ string_of_typ lt ^
    " = " ^ string_of_typ rt ^ " in " ^ string_of_expr ex))
  | Call(fname, actuals) as call -> let fd = function_decl fname in
  if List.length actuals != List.length fd.formals then
   raise (Failure ("Expecting " ^ string_of_int
     (List.length fd.formals) ^ " arguments in " ^ string_of_expr call))
  else
   List.iter2 (fun (ft, _) e -> let et = expr e in
    ignore (check_assign ft et
      (Failure ("Illegal actual argument found " ^ string_of_typ et ^
        " expected " ^ string_of_typ ft ^ " in " ^ string_of_expr e))))
   fd.formals actuals;
   fd.typ
  | Reference(s) -> (match (type_of_identifier s) with
                   ArrayTyp(_, _) -> Img
                  |MatrixTyp(_,_,_) -> Img
                  | _ -> raise ( Failure ("Cannot reference a type that is not an array"))
                  )
  | Dereference(s) -> (match (type_of_identifier s) with
                    Img(_) -> Int
                  | _ -> raise ( Failure ("Cannot dereference a type that is not an array pointer"))
                  )
  | MovePointer(s) -> (match (type_of_identifier s) with
                      Img -> Img
                    | _ -> raise ( Failure ("Cannot move a type that is not an array pointer"))
                  )

  | _ -> raise (Failure ("Unexpected type of expression"))

  in



  let check_bool_expr e =
    match (expr e) with
      Bool -> ()
    | _ -> raise (Failure ("Expected Boolean expression in " ^ string_of_expr e))
  in

  (* Verify a statement or throw an exception *)
  let rec stmt = function
  Block sl -> let rec check_block = function
  [Return _ as s] -> stmt s
  | Return _ :: _ -> raise (Failure "Nothing may follow a return")
  | Block sl :: ss -> check_block (sl @ ss)
  | s :: ss -> stmt s ; check_block ss
  | [] -> ()
  in check_block sl
  | Expr e -> ignore (expr e)
  | Return e -> let t = expr e in if t = func.typ then () else
  raise (Failure ("Return gives " ^ string_of_typ t ^ " expected " ^
   string_of_typ func.typ ^ " in " ^ string_of_expr e))

  |If(p, b1, b2) -> check_bool_expr p; stmt b1; stmt b2
  |For(e1, e2, e3, st) -> ignore (expr e1); check_bool_expr e2;
  ignore (expr e3); stmt st
  | While(p, s) -> check_bool_expr p; stmt s
  in

  stmt (Block func.body)

  in
  List.iter check_function functions
