(* Code developed by Sunil Kothari
 * Implements type inference Algorithm J for a language based on lambda calculus 
 * and extended with Milner-Let.
 * A part of this code is adapted/borrowed from 
 * code written by Prof. James L. Caldwell for COSC 4780.
 * Copyright (C) 2007  Sunil Kothari, University of Wyoming.
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
 *
 *
 *)


(* To do: make recursive functions tail-recursive *)






let counter = ref(96);;

type term =
  | Var of string
  | Abs of string * term
  | Ap of term * term
  | Let of string * term * term 
;;



type  types =
  | Poly of string
  | Arrow of  types *  types
  | All of string * types
;;

let wrap x =
     "(" ^ x ^ ")"
;;

let rec print_term t =
  match t with   
    | Var z -> z
    | Abs (x,t1) -> wrap ("\\" ^ x ^ "." ^ (print_term t1))
    | Ap (x,y) -> wrap ( (print_term x) ^ " " ^ (print_term y))
    | Let(x,t1,t2) -> wrap((print_term (Ap ((Abs(x, t2)),t1))))       
;;

type subs = (types * types) list;;

type env = (string * types) list ;;


let extendEnv  id ty env =  (id,ty)::(List.remove_assoc id env);;


let rec remove_all x t =
  match t with
      [] -> []
    | h::t -> if x = h then remove_all x t else (h::(remove_all x t))
;;

let rec member x t =
  match t with
      [] -> false
    | h::t ->  x = h or member x t
;;

let rec fv t =
  match t with
    | Var (x) -> [x]
    | Abs (x,t) -> remove_all x (fv t)
    | Ap(t1,t2) -> (fv t1)  @  (fv t2)
    | Let (x,t1,t2) -> (remove_all x (fv t2)) @ (fv t1)
;;

let rec unique l =
  match l with
      [] -> []
    | h::t -> if member h t then unique t else h::(unique t)
;;


let t_fv ty =
  let rec fv ty  =
    match ty with
        Poly a -> [a]
      | Arrow(a,b) ->  (fv a) @ fv b
      | All(x,a) ->  remove_all  x (fv a)
  in
    unique (fv ty)
;;



                                       
let newvar  =
  let count = ref (-1) in
  function(s) -> count := !count + 1;
           s ^ (string_of_int !count)
;;




let rec subst (y,t) t1 =
  match t1 with
    |  Var (x) -> if x = y then t else t1
    | Abs (x,t2) ->
        if x = y then t1
        else (if (member x (fv t)) then
                (let v = newvar (x ^ "_") in
                   Abs (v, subst (y,t) (subst (x, Var v) t2)))
              else
                Abs (x, subst (y,t) t2))
    | Ap(t3,t4) -> Ap (subst (y,t) t3, subst (y,t) t4)

   | Let(x,t2,t3) ->
        subst (y,t1) (Ap ((Abs (x,t3)), t2))

;;
    







let newchar =
   function() -> counter := !counter + 1;
          (char_of_int !counter)
;;





let rec search_ST  str  l = 
  match l with
      [] -> (false,'Z')
    | (s,c)::t -> if (s = str) then (true,c) else search_ST str  t
;;



let rec foldr f id l =
  match l with
      [] -> id
    | h::t -> f h (foldr f id t)
;;

let rec print_ST table = match table with
    [] -> ""
  | (s,c)::t -> s^(Char.escaped c)^print_ST t
;;

let rec convert_all_quantifier str l ty =
  match l with
      [] -> str,l
    |(c,s)::t ->  let result_str = if c = str then "All "^(Char.escaped s)^"."  else  "All "^str^"."     in 
        (
          match ty with
              All (x,ty1) ->  let (a,b)=(convert_all_quantifier x l ty1) in (result_str^a,b)
            |  Poly y  -> (result_str^y,l)
            | _  ->  (result_str,l)
        )
;;




let rec print_type t =
  match t with
    |    Poly a -> "'" ^ a
    | Arrow(a,b) -> wrap ((print_type a ) ^ "->" ^ (print_type b))
    | All(x,t) -> "All '" ^ x ^ "." ^ print_type t
 
;;


let rec print_simple_type t =
  match t with
      Poly a -> "'" ^ a
    | Arrow(a,b) -> wrap ((print_simple_type a ) ^ "->" ^ (print_simple_type b))
    | All(x,t) -> "All " ^ x ^ "." ^ print_simple_type t
      
;;

let print_list l  =
   (List.map (fun(x,y) -> print_string (print_simple_type x ^","^ print_simple_type y^ "\n" ) ) l)
  ;;

let print_env l  =
   (List.map (fun(x,y) -> print_string ( x ^" -->"^ print_simple_type y^ "\n" ) ) l)
  ;;




let print_ty ty = print_type ty;;


let rec make_type_scheme ty =
  let fv_list =  t_fv ty
  in match fv_list with
      [] -> ty
    | (h::t) -> make_type_scheme (All(h, ty))
;;




let rec t_member a ty =
  match ty with
    |  Poly b -> a = b
    | Arrow (ty1,ty2) -> ((t_member a ty1)|| (t_member a ty2))
    | All (x,ty1)  -> member a (remove_all x (t_fv ty1))                          
;;

let rec t_subst (y,t) t1 =
  match t1 with
      Poly a ->( match y with
                     Poly c -> if a = c then t else t1
                   | _ -> raise (Failure ("t_subst -- trying to substitute something besides a Poly: " ^ (print_type y)))
               )
    | Arrow(a,b) -> Arrow (t_subst (y,t) a, t_subst (y,t) b)
    | All(x,t2) -> 
        if (equal_types (Poly x) y) then t1
        else if member x (t_fv t ) then
          let x' = newvar ("tyvar"  ^ "_") in
            All (x', (t_subst (y,t) (t_subst (Poly x, Poly x') t2)))
        else
          All (x, t_subst (y,t) t2)
 
and   equal_types  ty1 ty2 =
  match (ty1,ty2) with
    | Poly a, Poly b ->  a = b
    | Arrow (ty3,ty4),Arrow (ty5,ty6) -> (equal_types ty3 ty5) && (equal_types ty4 ty6)
    | All (x,ty3) , All (y,ty4) -> if (x=y) then equal_types ty3 ty4 else equal_types ty3 (t_subst (Poly y, Poly x) ty4) 
    | _ -> false
;;

let rec subst_constraint_store (ty1,ty2)  s =
  match s with
      [] -> []
    | (ty3,ty4)::t -> (t_subst (ty1,ty2) ty3 , t_subst (ty1,ty2) ty4):: (subst_constraint_store (ty1,ty2) t)
;;


let unify t1 t2 =  
  let rec unifi c s =
    match c   with
        ([]:(types*types)list)  -> s 
      | h::tail -> (match h with
                      | (Poly a, ty)  -> (match ty with
                                              Poly b ->
                                                if a = b
                                                then unifi tail s
                                                else
                                                  if (t_member a ty)
                                                  then raise (Failure "Unable to unify types- occurs check failure")
                                                  else unifi (subst_constraint_store  (Poly a,ty) tail) ((subst_constraint_store (Poly a, ty) s)@[(Poly a ,ty)])
                                            | _ ->  if (t_member a ty)
                                              then raise (Failure ("Occurs check failure"^ (print_simple_type (Poly a))^"with"^ (print_simple_type ty)))
                                              else unifi (subst_constraint_store  (Poly a,ty) tail) ((subst_constraint_store (Poly a, ty) s)@[(Poly a ,ty)])
                                         )
                      | (ty, Poly a) -> unifi ((Poly a,ty)::tail) s
                      | (Arrow(t1,t2),Arrow(t3,t4)) -> unifi ((t1, t3)::( (t2,t4)::tail)) s
                      | (t1,t2)  -> raise (Failure (" Unknown type " ^ (print_type t1)^(print_type t2)^" to be unified"))
                   ) in
    unifi [t1,t2] []
;;





let rec t_subst_list (y,t1) l =
  match l with
      [] -> []
    | (ty1,ty2)::tail ->  (ty1,(t_subst (y,t1) ty2)):: (t_subst_list (y,t1) tail) 
;;


let rec  collect_free_terms_list (ty1,ty2) l  =
    match l with
        [] -> [(ty1,ty2)]
      | (ty3,ty4)::tail34 ->  if   (equal_types ty1 ty3) then [] else (collect_free_terms_list (ty1,ty2) tail34)
;;

 
let rec collect_free_terms s1 s2 =
   match s2 with
       [] -> s1
     | (ty1,ty2)::tail12 ->  (collect_free_terms_list (ty1,ty2) s1)@(collect_free_terms  s1  tail12)
;;


let rec compose_subst s1 s2 =
  match (s1,s2) with 
      ([],[]) -> []
    | ([],s2) -> s2
    | (s1,[]) -> s1
    | (s1,(ty21,ty22)::t2) -> (collect_free_terms ( compose_subst (t_subst_list (ty21,ty22) s1) t2) s2)
;;  




let rec apply_subst_types s ty =
 match s with
     [] -> ty
   | (ty1,ty2)::tail -> (match ty with
                            
                           | Poly a -> ( match  ty1 with
                                             Poly b -> if a = b then (apply_subst_types tail ty2)  else (apply_subst_types tail ty)
                                           | Arrow (b,c) -> (apply_subst_types tail ty )
                                           | _ -> raise (Failure "Failure in applying substitution to a type"))
                           | Arrow (ty3,ty4) -> apply_subst_types tail (Arrow (apply_subst_types [(ty1,ty2)] ty3, apply_subst_types [(ty1,ty2)] ty4))
                           | All (x,t) -> raise (Failure ("Do we reach here")))
;;


let rec apply_subst_env_list (s,ty) l =
  match l with
      [] ->  ((s:string),ty)
    | (ty1,ty2)::tail -> (apply_subst_env_list (s, t_subst (ty1,ty2) ty ) tail)
;;  



let rec apply_subst_env sub env  =
 match env  with
     [] -> []
   | (s,ty)::tail_env ->  (apply_subst_env_list (s,ty) sub ):: (apply_subst_env sub tail_env)
      
      
;;

let rec special_append l1 l2 =
  match l1 with
       [] -> l2
    | h::t -> if (List.mem h l2) then special_append t l2 else h::(special_append t l2);;

let rec t_fv_env  env =
  match env with
      [] -> []
    |  ((s:string),ty)::t -> special_append (t_fv ty) (t_fv_env t)
;;



let rec fv_diff l1 l2 =
  match l1 with
      [] -> []
    |  h::t -> if (member h l2) then (fv_diff t l2)  else [h]@ (fv_diff t l2)
;;


let rec generate_type_with_forall  l ty =
  match l with
      [] -> ty
    | h::t -> All (h, (generate_type_with_forall t ty))
;;

let generalize_type ty env =
  let fv_type = t_fv ty in
  let fv_env = t_fv_env env in
  let fv_type_modulo_env = fv_diff fv_type fv_env in
    generate_type_with_forall  fv_type_modulo_env  ty  
;;

let rec unroll ty =
  match ty with
      All(y,t1) -> let ty2 = t_subst ((Poly y) , (Poly (newvar ("tyvar"^"_")))) t1 in (unroll ty2)
    | _  -> ty
;;

let rec remove_all_env id env =
  match env with
      ([]:(string*types)list) -> []
    |(s,ty)::tail -> if (id = s) then remove_all_env id tail else (s,ty)::remove_all_env id tail
;;



let rec transform_type (t,table) =
match t with
    Poly a -> (match  search_ST a  table with
                   (true,b) -> (Poly (Char.escaped b), table)
                    
                 | (false,_) ->
                    
                     let c = newchar() in
                       (Poly (Char.escaped c),((a,c)::table))

              )
  | Arrow(a,b) ->
      let (ty1,mod_table1) = (transform_type (a,table)) in
      let (ty2,mod_table2) = (transform_type (b,mod_table1)) in
        (Arrow(ty1,ty2) ,mod_table2)
  | All(x,t) -> (match  search_ST x  table with
                     (true,b) -> let (ty,mod_table) = transform_type (t,table) in (All(Char.escaped b,ty), mod_table)
                   | (false,_) ->
                       let c = newchar() in
                       let mod_table = [(x,c)]@table in
                       let (str,new_mod_table) = transform_type (t,mod_table) in
                         (All ("'"^Char.escaped c,str), new_mod_table)
                )
;;



let rec infer t env subs =
  match t  with
    | Var x  -> 

        let ty = try (List.assoc x env) with Not_found -> raise (Failure "infer: not found.") in
          (match ty with             
               All(x,ty1) -> (unroll ty, subs)
             | _  -> (ty, subs))
           
    | Abs (x,t1)  ->

        let newtype = Poly (newvar ("tyvar_")) in             
        let complicated = (List.mem_assoc x env) in
        let dummyx = if complicated then (newvar x) else x  in
        let newenv = extendEnv dummyx newtype env  in  
        let new_t1 = if complicated then (subst (x, Var dummyx) t1) else t1 in
        let (type2,new_subs) = (infer  new_t1  newenv subs) in
          (Arrow (newtype, type2), new_subs)

    | Ap (t1,t2)  ->
        let new_type_var = (Poly (newvar ("tyvar_"))) in
        let (ty1,subs1) = infer t1 env subs in
        let (ty2,subs2) =  infer t2 env subs1 in
        let unification_result  = unify (apply_subst_types subs2 ty1) (apply_subst_types subs2 (Arrow (ty2, new_type_var))) in
        let final_subst = (compose_subst  unification_result subs2) in
          (new_type_var, final_subst)
           
    | Let (x,e1,e2) -> 

        let (ty1,subs) =  infer e1 env subs   in    
        let newenv = apply_subst_env subs env in
       let newtype = apply_subst_types subs ty1 in
       let envx = remove_all_env x env in  
       let type_scheme = generalize_type newtype  newenv in
       let newenv_extended = extendEnv   x  type_scheme envx in
          infer e2  newenv_extended subs

;;






let pretty_print ty = 
  function() -> counter  := 96 ;
    let new_ty,new_store = transform_type (ty,[]) in
    let newer_type = make_type_scheme new_ty in
      (print_type newer_type)
;;

let print_good_type t =
  let (k1,k2) = (infer t [] []) in
    pretty_print (apply_subst_types k2 k1) ();;



print_good_type  (Abs ("x", Abs ("y", (Ap (Var "x", (Abs ("z", Var "z")))))));;


print_good_type(Abs ("x" ,Abs ("y", (Ap (Var "x", Var "y")))));;

print_good_type (Abs ("f", Abs ("g", Abs ("x", (Ap ((Ap (Var "f" , Var "x")), (Ap (Var "g", Var "x"))))))));;

print_good_type (Let ("i" ,Abs("x", Var "x"), Ap (Var "i", Var "i")));;

print_good_type (Abs ("f", (Abs ("g" , (Abs ("x", Ap (Var "f" , Ap (Var "g", Var "x"))))))));;

print_good_type (Abs ("f", (Abs ("g" , (Abs ("x", (Ap ((Ap (Var "f", Var "x")) , (Ap (Var "g", Var "x"))))))))));;


print_good_type ( Let ("x1",(Abs ("y", (Abs ("z", (Ap ((Ap (Var "z", Var "y")), Var "y")))))),
                            Let ("x2",(Abs ("z", (Ap (Var "x1", Ap (Var "x1", Var "z"))))),
                                      Let ("x4",(Abs ("z", (Ap (Var "x2", Ap (Var "x2", Var "z"))))),
                                           Ap (Var "x4", (Abs ("x", Var "x")))))));;

print_good_type ( Let ("x1",(Abs ("y", (Abs ("z", (Ap ((Ap (Var "z", Var "y")), Var "y")))))),
                                      Let ("x4",(Abs ("z", (Ap (Var "x1", Ap (Var "x1", Var "z"))))), Ap (Var "x4", (Abs ("x", Var "x"))))));;


print_good_type (Let ("i" ,(Abs ("y",  Abs("x", Ap (Var "x", Var "y")))), Ap (Var "i", Var "i")));;



let t1 =  ( Let ("x1",(Abs ("y", (Abs ("z", (Ap ((Ap (Var "z", Var "y")), Var "y")))))),
                           Let ("x2",(Abs ("z", (Ap (Var "x1", Ap (Var "x1", Var "z"))))),
                                 Let ("x3",(Abs ("z", (Ap (Var "x2", Ap (Var "x2", Var "z"))))),
                                           Ap (Var "x3", Abs ("x", Var "x"))))));;




let t2 =  ( Let ("x1",(Abs ("y", (Abs ("z", (Ap ((Ap (Var "z", Var "y")), Var "y")))))),
                       Let ("x2",(Abs ("z", (Ap (Var "x1", Ap (Var "x1", Var "z"))))),
                            Let ("x3",(Abs ("z", (Ap (Var "x2", Ap (Var "x2", Var "z"))))),
                                 Let ("x4",(Abs ("z", (Ap (Var "x3", Ap (Var "x3", Var "z"))))),
                                      (Ap (Var "x4", Abs("x", Var "x"))))))));;





print_string (print_term t2);;


print_good_type t2;;