(* Code developed by Sunil Kothari.
 * Implements an  extension of Wand's algorithm as described in the paper
 * titled "On Extending Wand's Type Reconstruction Algorithm to Handle Polymorphic Let"
 * 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 function calls tail recursive *)



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



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



type env = Env of (string * types) list ;;

type constraint_types =
  | Equality of types * types (* called as c-constraint in the paper *)
  | MakeScheme of types * types * env  (* called as m-constraint in the paper *)
  | Instance of types * types    (* called as i-constraint in the paper *)
;;


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 ( "let "^x^"= "^(print_term t1)^" in " ^(print_term t2)^"\n")  

;;

let remove_all x t = List.filter (fun h -> not (x = h)) 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 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(t11,t12) -> Ap (subst (y,t) t11, subst (y,t) t12)
    | 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 rec equal_term t1 t2 =
match (t1,t2) with 
   Var z, Var y -> z = y
  | Abs (x,t3), Abs (y,t4) -> if x = y then equal_term t3 t4 else equal_term (subst(y, Var x) t4) t3  
  | Ap (t3,t4), Ap(t5,t6) -> equal_term t3 t5 && equal_term t4 t6
  | Let (x,t3,t4), Let (y,t5,t6) -> equal_term (Ap ((Abs(x, t4)),t3)) (Ap ((Abs(y, t6)),t5))
  | _,_ -> false
   
;;


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


let convert_env_to_string (Env e) = List.fold_right (fun x -> fun y -> x^y)(List.map (fun(x,y) ->  x^":"^(convert_type_to_string y)^" ") e) "";;


let convert_list_to_string l = List.fold_right (fun x -> fun y -> x^y)(List.map (fun(x) ->  x^" ") l) "";;



let  print_constraint c =
  match c with 
    | Equality (t1,t2) -> (convert_type_to_string t1)^" =e "^(convert_type_to_string t2)
    | MakeScheme (t1,t2,gamma) -> (convert_type_to_string t1)^" =s "^"["^ (convert_env_to_string gamma)^"]"^(convert_type_to_string t2)
    | Instance (t1,t2) -> (convert_type_to_string t1)^" =i "^(convert_type_to_string t2)
;;

 
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]
      | TypeSchemePoly a -> [a]
      | Arrow(a,b) ->  (t_fv a) @ (t_fv b)
      | All(x,a) ->  remove_all  x (t_fv a)

;;



                                     
let extendEnv  id ty (Env e) = 
  match e with 
      [] -> Env ([(id,ty)])
    | l  -> Env ((id,ty)::l)
;;
                                       


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




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)
                           | TypeSchemePoly a ->  ty (* We really don't apply substitution to a type scheme variable *)
                           | Arrow (ty3,ty4) -> apply_subst_types tail (Arrow (apply_subst_types [(str,ty2)] ty3, apply_subst_types [(str,ty2)] ty4))
                           | All (x,t) ->  All (x,apply_subst_types (List.filter (fun (y,z) -> not (y = x)) s) t)  
			)

;;


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



let rec print_constraint_list l =
   match l with
     | Equality (t1,t2)::t -> print_string ("Equality("^(convert_type_to_string t1)^","^(convert_type_to_string t2)^")\n"); print_constraint_list t
     | Instance(x,y)::t -> print_string ("Instance("^(convert_type_to_string x)^","^(convert_type_to_string y)^")\n"); print_constraint_list t
     | MakeScheme(x,y,e)::t -> print_string ("Scheme("^ (convert_type_to_string x)^","^(convert_type_to_string y)^ ","^(convert_env_to_string e)^ ")\n" ); print_constraint_list t 
     | [] -> print_string "Empty List \n"
;;
     
let rec transform_constraint_store (str,ty2)  c = 
 match c with 
     [] -> []
   | Equality (ty3,ty4)::t -> 
       (Equality (t_subst (str,ty2) ty3 , t_subst (str,ty2) ty4)):: (transform_constraint_store (str,ty2) t)
   | MakeScheme(ty3,ty4,e)::t -> (MakeScheme(t_subst(str,ty2) ty3, t_subst(str,ty2)ty4, apply_subst_env [(str,ty2)] e))::(transform_constraint_store (str,ty2) t)
   | Instance(ty3,ty4)::t -> Instance(t_subst(str,ty2)ty3,t_subst(str,ty2)ty4)::(transform_constraint_store (str,ty2) t)
;;


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


let rec apply_subst_constraint c s = 
  match  s with
    [] -> c 
    | h::t -> apply_subst_constraint (transform_constraint_store  h c) t 

;;    




 let compose_subst s1 s2 = 
  let comp1 = List.map (fun (x,y) -> (x,apply_subst_types s2 y)) s1 in 
  let comp2 = List.map (fun (x,y) -> x) comp1 in
  let comp3 = List.map (fun (x,y) -> if (List.mem x comp2) then [] else [(x,y)]) s2   in
   comp1 @ (List.concat comp3);;



(* Removed a bug which captures scheme variables in the env for generalization*)

let rec t_fv_env  env =
  match env with 
      [] -> []
    | ((s:string),ty)::t ->  
  if (String.get s 0) =='s'then  (t_fv_env t) else  (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 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 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 check_for_empty_equality_predicates constraint_store = 
  match constraint_store with 
      [] -> true 
    | h::tail -> ( match h with 
                       
		     | Instance(ty1,ty2) -> check_for_empty_equality_predicates tail
                     |  Equality(ty1,ty2) -> false 
                     |  MakeScheme(ty1,ty2,_) ->  check_for_empty_equality_predicates tail
		 )
;;			    
			     


let rec check_for_empty_scheme_predicates constraint_store = 
  match constraint_store with 
      [] -> true 
    | h::tail -> ( match h with                        
		     | Instance(ty1,ty2) -> check_for_empty_scheme_predicates tail
                     |  MakeScheme(ty1,ty2,_) ->  false 
                     | Equality(ty1,ty2) -> check_for_empty_scheme_predicates tail
		 )
;;



let rec check_for_empty_instance_predicates constraint_store = 
  match constraint_store with 
      [] -> true 
    | h::tail -> ( match h with 
		     | Instance(ty1,ty2) -> false
                     |  MakeScheme(ty1,ty2,_) ->  check_for_empty_instance_predicates tail
                     | Equality(ty1,ty2) -> check_for_empty_instance_predicates tail
		 )
;;

 


let rec unifi c s =
  if (check_for_empty_equality_predicates c) then (s,c) 
  else 
    (match c   with
	([]:constraint_types list)  -> (s,c)  
      | h::tail -> (match h with
		      | Equality (Poly a, Poly b)  ->  
			  if a = b 
			  then unifi tail s 
			  else unifi (transform_constraint_store  (a,Poly b) tail)  ((transform_subst_list  (a, Poly b) s)@[(a ,Poly b)])
		      | Equality (Poly a, ty) ->  
			  if (List.mem a (t_fv ty)) 
			  then raise (Failure ("Occur check failure: trying to unify"^ (convert_type_to_string ty)^" with Poly "^a  ))
			  else unifi (transform_constraint_store  (a,ty) tail)  ((transform_subst_list  (a, ty) s)@[(a ,ty)])
		      | Equality (ty, Poly a) -> unifi ((Equality (Poly a,ty))::tail) s
		      | Equality (Arrow(t1,t2),Arrow(t3,t4)) -> unifi ((Equality (t1, t3))::( (Equality (t2,t4))::tail)) s
		      | MakeScheme (ty1,ty2, env) ->  unifi (tail@[MakeScheme(ty1,ty2,env)]) s 		  
                      | Instance (ty1, ty2) -> unifi (tail@[Instance(ty1,ty2)]) s
		      | _  -> raise (Failure " Types cannot be unified")
		   )
    )
;;



let rec replace_sch_var  (ty1,ty2,env) tail new_constraint_store = 
  match tail with 
     [] -> new_constraint_store 
    | h::t -> ( match h with 
		    Instance (ty3, ty4) ->  
		      (if (equal_types ty2 ty3) 
		       then 
			 let new_type = generalize_type ty1 env  in 
			 let ultra_new_type = unroll new_type in 
			 let updated_store_1= (Equality (ty4,ultra_new_type)::new_constraint_store) in
		           replace_sch_var (ty1,ty2,env) t updated_store_1 
		       else  
			   replace_sch_var (ty1,ty2,env) t new_constraint_store )
		  | Equality(ty4,ty5) ->  replace_sch_var (ty1,ty2,env) t new_constraint_store
		  | MakeScheme(ty3,ty4,env) -> replace_sch_var (ty1,ty2,env) t new_constraint_store
	      )
;;



      
let rec generate_constraint  goal constraint_store = 
  match goal with 
      []  -> constraint_store 
    | (Env e,t, ty)::tail_goal  -> (match t  with 
					Var x  ->  let envty = try (List.assoc x e) with Not_found -> raise (Failure "generate_constraint exception") in
					  (match envty with 
					      
                                           | All(x,ty1) -> generate_constraint tail_goal ((Equality (ty,unroll envty))::constraint_store) 
					     | somety -> 
						 (match somety with 
						    |Poly u -> generate_constraint tail_goal ((Equality (somety,ty))::constraint_store)
						    |TypeSchemePoly a -> generate_constraint tail_goal ((Instance (somety,ty))::constraint_store)
						      
						    |_ -> raise (Failure "We should be dealing with type variables")
						 )
					  )
					    
				      | Abs (x,t1)  -> 
					  let newtype_var1 = Poly (newvar ("tyvar"^"_")) in
					  let newtype_var2 = Poly (newvar ("tyvar"^"_")) in 
                                          let new_envx = remove_all_env x e in 
					  let new_env = extendEnv x  newtype_var1 (Env new_envx)  in  
					  generate_constraint  ((new_env,t1,newtype_var2)::tail_goal) 
                                             ((Equality (ty, (Arrow (newtype_var1, newtype_var2))))::constraint_store)

				      | Ap (t1,t2)  -> 
					  let new_type_var = Poly (newvar ("tyvar" ^ "_")) in 
					  let new_goal = ((Env e),t1,Arrow (new_type_var, ty))::tail_goal in 
                                          let newest_goal = ((Env e),t2, new_type_var)::new_goal in 
					    generate_constraint  newest_goal constraint_store 
					      
				      | Let (x,t1,t2) -> 
				
      	 				  let new_type_var1  =  Poly (newvar ("tyvar" ^ "_")) in
					  let new_store_left = generate_constraint [((Env e),t1,new_type_var1)][] in 
					  let new_type_var2 = TypeSchemePoly (newvar ("tyvar" ^ "_")) in 
                                          let new_envx = (remove_all_env x e) in 
                                          let new_env = extendEnv x  new_type_var2 (Env new_envx)  in  
                                          let new_store_right = generate_constraint [(new_env,t2,ty)] [] in 
                                         let new_constraint_store1 = 
					     [(MakeScheme (new_type_var1,new_type_var2, (Env e)))]@ constraint_store in
                                         generate_constraint  tail_goal (new_store_left @ new_constraint_store1 @ new_store_right) 
        
       
					      
					      
				   )
	
;;





let rec unifi2 c   =
  if (check_for_empty_scheme_predicates c) then ([]:(string* types) list)
  else 
    match c   with
	[]  -> ([]:(string*types) list)
      | h::tail -> (match h with 			 			 
		      | MakeScheme (ty1,ty2, env) ->  
			  let const_store = replace_sch_var (ty1,ty2,env) tail  [] in 
                          let new_subst = fst (unifi const_store []) in 
			    compose_subst  new_subst (unifi2 (apply_subst_constraint tail new_subst))
                      | Instance (ty1, ty2) -> unifi2 (tail@[Instance(ty1,ty2)])  			   			   
		      | Equality(ty1,ty2) -> raise (Failure "No equality constraints should be found in Phase II")
		   )
	  
;;







let rec fun_apply subs_list =  
    match subs_list with 
	[] -> raise (Failure "Type jk not found")
      | (a,b)::t  ->  if a = "jk" then b else fun_apply t
;; 



let type_infer t  =
  let initial_const_store = generate_constraint  [((Env [("fix",All ("sigma", (Arrow  ((Arrow (Poly "sigma" , Poly "sigma"),Poly "sigma")))))]),t,Poly "jk")] [] in
  let (subst_list,const_store)  =  unifi initial_const_store [] in 
    if not (check_for_empty_scheme_predicates const_store)
    then             
      let new_subst_list = unifi2 const_store  in 
	fun_apply (compose_subst subst_list new_subst_list)
    else  
      fun_apply  subst_list;;



let type_infer_with_constraint c = 
 let (subst_list,const_store)  =  unifi c [] in 
    if not (check_for_empty_scheme_predicates const_store)
    then             
      let new_subst_list = unifi2 const_store  in 
	fun_apply (compose_subst subst_list new_subst_list)
    else  
      fun_apply  subst_list;;


let type_infer_with_subst t  =
  let initial_const_store = generate_constraint  [((Env [("fix",All ("sigma", (Arrow  ((Arrow (Poly "sigma" , Poly "sigma"),Poly "sigma")))))]),t,Poly "jk")] [] in
  let (subst_list,const_store)  =  unifi initial_const_store [] in 
    if not (check_for_empty_scheme_predicates const_store)
    then             
      let new_subst_list = unifi2 const_store in 
	((fun_apply (compose_subst subst_list new_subst_list)), compose_subst subst_list new_subst_list)
    else  

      ((fun_apply  subst_list), subst_list)
;;




let type_infer_constraint t = generate_constraint [((Env []),t,Poly "jk")] [];;

let type_infer_phase1 t = snd (unifi (generate_constraint [((Env []),t,Poly "jk")] [])[]);;

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 counter = ref(96);;

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



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 
      (convert_type_to_string newer_type)
;;

let print_good_type t = pretty_print (type_infer t)();;




(* Example 1 in the paper *)
print_good_type (Let ("i" ,Abs("x", Var "x"), Ap (Var "i", Var "i")));; 



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

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


(* double application *)
print_good_type (Abs ("f", (Abs ("x" ,  Ap (Var "f" , Ap (Var "f", Var "x"))))));;

(* Example 2 in the tech report*)
print_good_type (Abs ("x", Let ("f", Let ("g", Var "x", Abs("a", Var "g")), Let ("h", Var "f", Ap (Var "h", Var "x")))));;
 type_infer_constraint (Abs ("x", Let ("f", Let ("g", Var "x", Abs("a", Var "g")), Let ("h", Var "f", Ap (Var "h", Var "x")))));;



(* Example 2 with lets moved to the righmost sub-tree*)
print_good_type (Abs ("x", Let ("g", Var "x", Let ("f", Abs("a", Var "g"), Let ("h", Var "f", Ap (Var "h", Var "x"))))));;

(* Print ocnstraints for the above example *)
 type_infer_constraint (Abs ("x", Let ("g", Var "x", Let ("f", Abs("a", Var "g"), Let ("h", Var "f", Ap (Var "h", Var "x"))))));; 



(* Example 3 *)

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



(* Example 4 - Additional Example Taken from Reynolds *) 

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