(*
   CS421 MP6: Unification Algorithm  
*)

open Mp6common

(* Problem 1 - Not graded. Only for exercise. *)
let asExpTy1 () = mk_fun_ty bool_ty (mk_list_ty int_ty)
let asExpTy2 () = mk_fun_ty (fresh()) (mk_fun_ty (fresh()) (mk_fun_ty (fresh()) (fresh())))
let asExpTy3 () = mk_fun_ty (fresh()) (mk_list_ty (mk_pair_ty (fresh()) int_ty))
let asExpTy4 () = mk_pair_ty string_ty (mk_fun_ty (mk_list_ty (fresh())) (fresh()))

(* Problem 2 *)
let to_subst_fun s = 
  fun n -> (try List.assoc n s with _ -> TyVar n)

(* Problem 3 *)
let rec lift_subst s = 
  fun ty ->
  match ty with
    TyVar m -> to_subst_fun s m
  | TyConst(st, typelst) ->  TyConst(st, List.map (fun t -> lift_subst s t) typelst);;

(* Problem 4 *)
let rec occurs n ty =
  match ty with
    TyVar m -> n=m
  | TyConst(st, typelst) ->
     List.fold_left (fun xl x -> if xl then xl else occurs n x) false typelst;;

(* Problem 5 *)
let rec addNewEqs lst1 lst2 acc =
  match lst1,lst2 with
    [],[] -> Some acc
  | t::tl, t'::tl' -> addNewEqs tl tl' ((t,t')::acc)
  | _ -> None

let rec unify eqlst =
  match eqlst with
    [] -> Some([])
    (* Delete *)
  | (s,t)::eqs when s=t -> unify eqs
    (* The case below is just to get a canonical answer *)
  | (TyVar n, TyVar m)::eqs when n > m -> unify ((TyVar m, TyVar n)::eqs)
    (* Eliminate *)
  | (TyVar(n),t)::eqs when not(occurs n t)-> 
      let eqs' = List.map (fun (t1,t2) -> (lift_subst [n,t] t1 , lift_subst [n,t] t2)) eqs
      in (match unify eqs' with
           None -> None
         | Some(phi) -> Some((n, lift_subst phi t):: phi))
    (* Orient *)
  | (TyConst(str, tl), TyVar(m))::eqs -> unify ((TyVar(m), TyConst(str, tl))::eqs)
    (* Decompose *)
  | (TyConst(str, tl), TyConst(str', tl'))::eqs when str=str' -> 
      (match (addNewEqs tl tl' eqs) with
        None -> None
      | Some l -> unify l)
    (* Other *)
  | _ -> None;;

(* Problem 6 *)

let rec first p l =
    match l with [] -> None
    | (x :: xs) -> if p x then Some x else first p xs

let rec cannon_type subst m ty =
   match ty
   with TyVar n ->
     (match first (fun p -> fst p = n) subst
      with Some (j,k) -> (subst, m, TyVar k)
      | None -> ((n,m)::subst), m+1, TyVar m)
   | TyConst (c, tys) ->
     (match
       List.fold_left
       (fun (subst, n, tyl) -> fun ty ->
        (match cannon_type subst n ty
         with (subst', n', ty') -> (subst', n', ty'::tyl)))
       (subst, m, [])
       tys
      with (new_subst, new_m, new_tys) -> (new_subst, new_m, TyConst(c, List.rev new_tys)))

let print_cannon ty = match cannon_type [] 0 ty with (_, _, new_ty) -> print_expType new_ty

let equiv_types ty1 ty2 =
   let (_, _, new_ty1) = cannon_type [] 0 ty1 in
   let (_, _, new_ty2) = cannon_type [] 0 ty2 in new_ty1 = new_ty2;;