typing.ml

open Ast
module H = Hashtbl

let debug = ref false

let dummy_loc = Lexing.dummy_pos, Lexing.dummy_pos

exception Error of Ast.location * string

(* use the following function to signal typing errors, e.g.
      error ~loc "unbound variable %s" id
*)
let error ?(loc=dummy_loc) f =
  Format.kasprintf (fun s -> raise (Error (loc, s)))
    ("@[" ^^ f ^^ "@]")

type fn_env = (string, fn) H.t

let fn_env : fn_env = H.create 16

type lambda_env = (string, tdef) H.t

let lambda_env : lambda_env = H.create 16

let dummy_var = { v_name = "dummy"; v_ofs = -1 }

let () =
  H.add fn_env "len" { fn_name = "len"; fn_params = [dummy_var]; fn_anoymous = false };
  H.add fn_env "list" { fn_name="list"; fn_params = [dummy_var]; fn_anoymous=false }

type var_env = (string, var) H.t

let rec expr (ctx: var_env) (e: Ast.expr) =
  match e with
  | Ecst c -> TEcst c
  | Eident x -> begin
      try
        let v = H.find ctx x.id in
        TEvar v
      with Not_found ->
        try
          let fn = H.find fn_env x.id in
          TEfn fn
        with Not_found ->
          error ~loc:x.loc "Unbound variable %s" x.id
    end
  | Ebinop (op, e1, e2) ->
      TEbinop (op, expr ctx e1, expr ctx e2)
  | Eunop (op, e) ->
      TEunop (op, expr ctx e)
  | Ecall (Eident id, [arg]) when id.id = "range" ->
      let te = expr ctx arg in
      TErange te
  | Ecall (Eident id, [arg]) when id.id = "list" ->
      let te = expr ctx arg in
      begin match te with
      | TElist _ | TErange _ ->
          te
      | _ ->
          error ~loc:id.loc
            "TypeError: list() argument must be a list or range, got %s"
            (match te with
            | TEcst _   -> "constant"
            | TEvar _   -> "variable"
            | TEfn _    -> "function"
            | TEbinop _ -> "expression"
            | TEunop _  -> "expression"
            | TEcall _  -> "expression"
            | TErange _ -> assert false
            | TEget _   -> "expression"
            | _ -> "other")
      end
  | Ecall (fn_expr, args) ->
      let te_fn   = expr ctx fn_expr in
      let te_args = List.map (expr ctx) args in
      let loc =
        match fn_expr with
        | Eident id -> id.loc
        | _         -> dummy_loc
      in
      begin match te_fn with
      | TEfn fn ->
          let expected = List.length fn.fn_params in
          let actual   = List.length te_args in
          if expected <> actual then
            error ~loc
              "arity mismatch: %s expects %d argument(s) but got %d"
              fn.fn_name expected actual
      | _ -> ()
      end;
    TEcall (te_fn, te_args)
  | Elist e ->
      TElist (List.map (expr ctx) e)
  | Eget (e1, e2) ->
      TEget (expr ctx e1, expr ctx e2)
  | Elambda (args, body) ->
      let fn = lambda args body in
      TEfn fn

and stmt (ctx: var_env) (s: Ast.stmt) : Ast.tstmt =
  match s with
  | Sif (e, s1, s2) ->
      TSif (expr ctx e, stmt ctx s1, stmt ctx s2)
  | Sreturn e ->
      TSreturn (expr ctx e)
  | Sassign (x, e) ->
      let v = H.find ctx x.id (* this should never fail *) in
      TSassign (v, expr ctx e)
  | Sprint e ->
      TSprint (expr ctx e)
  | Sblock sl -> 
      TSblock (List.map (stmt ctx) sl)
  | Sfor (x, e, s) -> TSfor(H.find ctx x.id, expr ctx e, stmt ctx s)
  | Swhile (e, s) -> TSwhile (expr ctx e, stmt ctx s)
  | Seval e ->
      TSeval (expr ctx e)
  | Sset (e1, e2, e3) ->
      TSset (expr ctx e1, expr ctx e2, expr ctx e3)

and alloc_var ctx x =
  if not (H.mem ctx x.id) then
    H.add ctx x.id { v_name = x.id; v_ofs = -1 }

and alloc_vars ctx (s: Ast.stmt) =
  match s with
  | Sif (_, s1, s2) ->
      alloc_vars ctx s1; alloc_vars ctx s2
  | Sassign (x, _) ->
      alloc_var ctx x
  | Sblock sl ->
      List.iter (alloc_vars ctx) sl
  | Sfor (x, _, s) ->
      alloc_var ctx x; alloc_vars ctx s
  | Swhile (_, s) ->
      alloc_vars ctx s
  | Sprint _ | Sreturn _ | Seval _ | Sset (_, _, _) -> ()

and lambda args body =
  let id = string_of_int (H.length lambda_env) in
  let mk_var x = { v_name = x.id; v_ofs = -1 } in
  let targs = List.map mk_var args in
  let fn = { fn_name = id; fn_params = targs; fn_anoymous = true } in
  let ctx = H.create 32 in (* empty context *)
  List.iter (fun v -> H.add ctx v.v_name v) targs;
  let tbody = stmt ctx (Sreturn body) in
  H.add lambda_env id (fn, tbody);
  fn

and def (f, args, body) =
  let names = List.map (fun id -> id.id) args in
  let rec check_dups = function
    | [] -> ()
    | x::xs ->
        if List.mem x xs then
          error ~loc:f.loc "Duplicate parameter name %s" x
        else
          check_dups xs
  in
  check_dups names;
  let mk_var x = { v_name = x.id; v_ofs = -1 } in
  let targs = List.map mk_var args in
  let fn = { fn_name = f.id; fn_params = targs; fn_anoymous = false } in
  H.add fn_env f.id fn;
  let ctx = H.create 32 in (* empty context *)
  List.iter (fun v -> H.add ctx v.v_name v) targs;
  alloc_vars ctx body;
  let tbody = stmt ctx body in
  (fn, tbody)

let file ?debug:(b=false) (p: Ast.file) : Ast.tfile =
  debug := b;
  let defs, s = p in
  (* typing all the declared functions *)
  let tdefs = List.map def defs in
  let ctx = H.create 32 in (* empty context *)
  alloc_vars ctx s;
  let ts = stmt ctx s in
  let fn_main = { fn_name = "main"; fn_params = []; fn_anoymous = false } in
  let tmain = (fn_main, ts) in
  let tdefs2 = Hashtbl.fold (fun _ v acc -> v :: acc) lambda_env tdefs in
  tmain :: tdefs2