(* This code implements Algorithm M for a language based on lambda calculus and extended with Milner-Let 
 * as mentioned in the paper: Proof about a Folklore Let-Polymorphic Type Inference Algorithm
 * 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 calls 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
;;






type env = (string * types) list ;;


let extendEnv  id ty env = 
  match env with 
      ([]:(string*types)list) -> [(id,ty)]
    | l  -> [(id,ty)]@l
;;



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 t)

;;


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(t1,t2) -> Ap (subst (y,t) t1, subst (y,t) t2)
    | Let(x,t2,t3) -> subst (y,t1) (Ap ((Abs (x,t3)), t2))  
;;
     
let wrap x = "(" ^ x ^ ")" ;;

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 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 ty = 
  match ty 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 rec make_type_scheme ty = 
  let fv_list = List.rev (t_fv ty)
  in match fv_list with 
      [] -> ty 
    | (h::t) -> make_type_scheme (All(h,ty))
;;


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 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_simple_type newer_type)
;;




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))))
;;


     

let rec t_subst (y,t) t1 =
  match t1 with
      Poly a ->( match y with 
                     Poly c -> if a = c then t else t1
                   | _ -> t1           
	       )
    | 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 (t_subst (Poly y, Poly x) ty4) ty3 
  | _ -> false
;;




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 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) -> ty 
			)                         			
;;


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 infer t env  typ = 
  match t  with 
    |  Var x  ->  
	 let ty1 =  try (List.assoc x env) with Not_found -> raise (Failure ("Type of variable "^ x ^" not found in the environment")) in
	   (match ty1 with              
	      | All(x,ty2) -> let ty3 = unroll ty1 in (unify typ ty3)   
	      | ty22  -> unify typ ty22
	   )	  
    | Abs (x,t1)  -> 
	let dummyx = newvar (x) in  
	let typevar1 = Poly (newvar ("tyvar_")) in
        let typevar2 = Poly (newvar ("tyvar_")) in 
        let subst1 = unify typ (Arrow (typevar1,typevar2)) in 
	let newenv = extendEnv dummyx  (apply_subst_types subst1 typevar1) (apply_subst_env subst1 env)   in  
	let subst2 = infer (subst (x, (Var dummyx)) t1) newenv (apply_subst_types subst1 typevar2) in
	  (compose_subst subst2 subst1)
  | Ap (t1,t2)  -> 
      let typevar1 = Poly (newvar ("tyvar_")) in
      let subst1 = (infer t1 env (Arrow(typevar1,typ))) in 
      let newenv = (apply_subst_env  subst1 env) in
      let subst2 = (infer t2  newenv (apply_subst_types subst1  typevar1)) in
       (compose_subst subst2 subst1)
  | Let (x,t1,t2) ->  
      let typevar1 = Poly (newvar ("tyvar_")) in
      let subst1 =  (infer t1 env typevar1)   in     
      let newenvx = (apply_subst_env subst1 (remove_all_env x env))in
      let newtype = (generalize_type (apply_subst_types subst1 typevar1) newenvx) in  
      let newenvx_extended = (extendEnv  x newtype newenvx)in
      let subst2 = (infer t2 newenvx_extended (apply_subst_types subst1 typ)) in 
	 (compose_subst subst2 subst1)  

;;

let print_good_type t = let subst1 = (infer t [] (Poly "jk")) in
        pretty_print (apply_subst_types subst1 (Poly "jk")) ();;


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



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")));; 



(* Infer type for the Y combinator *)
print_good_type (Abs ("h", (Ap ((Abs ("x", (Ap (Var "h", (Ap (Var "x", Var "x")))))),(Abs ("x", (Ap (Var "h", (Ap (Var "x", Var "x"))))))))));;

(* infer type for the compose operator*)
print_good_type (Abs ("f", (Abs ("g" , (Abs ("x", Ap (Var "f" , Ap (Var "g", Var "x"))))))));;

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


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

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

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


 (* this example is taken from some paper *)
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 ("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_good_type (Let ("i" ,Abs("x", Var "x"), Ap (Var "i", Var "i")));; 

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

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"))))),
					   (Ap (Var "x2", Abs ("x", 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 ("x3",(Abs ("z", (Ap (Var "x2", Ap (Var "x2", Var "z"))))), 
					   Ap (Var "x3", (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 ("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")))))));;