(* 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
 * 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
;;



let rec remove_all x t = List.filter (fun y -> not (x=y)) 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 List.mem h t then unique t else h::(unique t)
;;


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

;;

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 (List.mem 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) ->if x = y then t1 else 
                     if  List.mem x (fv t) then 
                      let v = newvar (x ^ "_") in
			 Let (v, subst (y,t) t2, subst (y,t) (subst (x, Var v) t3))
                     else 
                       Let (x, (subst (y,t) t2),  (subst (y,t) t3))

;;
     
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_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 ->
             let (b,table') = try (List.assoc a table,table) with
                 Not_found -> let c = newchar() in (c, (a,c)::table)
             in
               (Poly (Char.escaped b), 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) ->
          let (b,table') = try (List.assoc x table,table) with
              Not_found -> let c = newchar() in (c, (x,c)::table)
          in
          let (ty,mod_table) = transform_type (t,table') in
            (All(Char.escaped b,ty),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 -> if a = y then t else t1              
    | Arrow(a,b) -> Arrow (t_subst (y,t) a, t_subst (y,t) b)
    | All(x,t2) -> 
	if x = y then t1
	else if List.mem x (t_fv t) then
	  let x' = newvar ("tyvar"  ^ "_") in
	    All (x', (t_subst (y,t) (t_subst (x, Poly x') t2)))
	else
	  All (x, t_subst (y,t) t2)
;;

let rec 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 (y, Poly x) ty4) ty3 
    | _ -> false
;;




let rec subst_subst (ty1,ty2)  s =
  List.map (fun (x,y) -> (x, t_subst (ty1,ty2) y)) s 
;;



let rec subst_constraint_store (ty1,ty2)  s =
  List.map (fun (x,y) -> t_subst (ty1,ty2) x, t_subst (ty1,ty2) y) s 

;;

 
let unify t1 t2 =  
  let rec unifi c s =
    match c   with
        []  -> s 
      | h::t -> (match h with
                      | (Poly a, Poly b)  -> if a = b
                                             then unifi t s
                                             else unifi (subst_constraint_store (a,Poly b) t) ((a, Poly b)::(subst_subst (a, Poly b) s))
                      | (Poly a, ty) ->   if (List.mem a (t_fv ty))
                                                  then raise (Failure "Unable to unify types- occurs check failure")
                                                  else unifi (subst_constraint_store  (a,ty) t) ((a, ty)::(subst_subst (a, ty) s))                                                         
                      | (ty, Poly a) -> unifi ((Poly a,ty)::t) s
                      | (Arrow(t1,t2),Arrow(t3,t4)) -> unifi ((t1, t3)::( (t2,t4)::t)) s
                      | (t1,t2)  -> raise (Failure (" Unknown type " ^ (print_simple_type t1)^(print_simple_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  destruct x =
   match x  with
  Poly a -> a 
     | _ -> raise (Failure ("Substtituion should always map type variables"))
;;  



let   collect_free_terms s1 s2 = 
  let dom_first_list =  (List.map (fun (x,y) ->  x) s1) in 
     List.concat (List.map (fun (x,y) -> if (List.mem x dom_first_list) then [] else [(x,y)]) s2)@s1
;;


  
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 
    | (str,ty2)::tail -> (match ty with 
                            | Poly a -> if a = str then (apply_subst_types tail ty2)  else (apply_subst_types tail ty)
                            | Arrow (ty3,ty4) -> apply_subst_types tail (Arrow (apply_subst_types [(str,ty2)] ty3, apply_subst_types [(str,ty2)] ty4))
                            | All (x,t) -> ty 
			 )                         			
;;

let apply_subst_env s e = List.map (fun ((x:string),y) -> (x, apply_subst_types s y)) e;;



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 (List.mem 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 (y, (Poly (newvar ("tyvar"^"_")))) t1 in (unroll ty2) 
    | _  -> ty 
;;

let remove_all_env id env = List.filter (fun (x,y) -> not (x == id)) env;;




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