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

(* Desugaring of let is not needed --something to ponder *)
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")  
(*      wrap((print_term (Ap ((Abs(x, t2)),t1)))) *)
;;


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 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) -> 	 
	let new_var =   subst(y,t) (Var x) in 
          (match new_var with    
	       Var b ->	 Let (b, subst(y,t) t2 , subst(y,t) t3) 
	     | _ -> raise (Failure "Unknown term resulting")
	  )
;;


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

(* Prints a list of string *)
let convert_list_to_string l = List.fold_right (fun x -> fun y -> x^y)(List.map (fun(x) ->  x^" ") l) "";;

let count_occurrence l = List.fold_right (fun x -> fun y -> 1 + y)(List.map (fun(x) ->  1) l) 0;;



count_occurrence ["x";"x";"x"];;

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



type pproof = 
  | Emptyp 
  | Letp of env * string * term * term * types * constraint_types * pproof *pproof
  | Appp of env * term * term * types* constraint_types *pproof *pproof
  |  Absp of env * string * term * types *constraint_types * pproof 
  |  Axiomp of env * string * types * constraint_types
;; 



    
let rec convert_proof_string  p = 
  match p with
    |Emptyp -> ""
    |  Letp (gamma, var, t1, t2, ty, cons, p1, p2) -> (convert_env_to_string gamma)^ " |- let " ^ var ^ "=" ^(print_term t1)^" in "^(print_term t2)^ ":"^(convert_type_to_string ty)^"  "^(print_constraint cons) ^"\n"^ (convert_proof_string p1)^ (convert_proof_string p2)
    | Appp(gamma, t1, t2, ty, cons, p1,p2) -> (convert_env_to_string gamma)^ " |- " ^(print_term t1) ^ (print_term t2)^":"^(convert_type_to_string ty)^"  "^(print_constraint cons) ^"\n"^ (convert_proof_string p1)^ (convert_proof_string p2)
    | Absp(gamma,  var, t1, ty, cons,p1) ->  (convert_env_to_string gamma)^ " |- " ^ var ^"."^(print_term t1) ^ ":"^ (convert_type_to_string ty)^"  "^(print_constraint cons) ^"\n"^ (convert_proof_string p1)
    | Axiomp(gamma, var, ty, cons) -> (convert_env_to_string gamma)^" |- " ^ var ^ ":" ^(convert_type_to_string ty)^"  "^ (print_constraint cons)^"\n"
;;


let count_let_var_occurrence l t =
let rec count_number_of_occurrence_term l t =
 match t with 
     Var x -> if x = l then [x] else []
   | Abs(x,t1) -> 
       if x = l 
       then let new_var = newvar x in count_number_of_occurrence_term l (subst (x, Var new_var) t1) else count_number_of_occurrence_term l t1
   | Ap(t1,t2)-> (count_number_of_occurrence_term l t1) @ (count_number_of_occurrence_term l t2)
   | Let(x,t1,t2) -> 
       if x = l 
       then (count_number_of_occurrence_term l t2)
       else (count_number_of_occurrence_term l t1) @ (count_number_of_occurrence_term l t2)
in count_occurrence (count_number_of_occurrence_term l t)
;;


(*count_let_var_occurrence "i" (Let("i",Abs("x", Var "x"), Ap (Var "i", Var "i")));;*)

let count_number_of_occurrence_proof let_var p =
 match p with 
  |Emptyp -> 0
    |  Letp (gamma, var, t1, t2, ty, cons, p1, p2) -> if (var  = let_var) then  (count_let_var_occurrence let_var t2) else (count_let_var_occurrence let_var t1) +  (count_let_var_occurrence let_var t2)
    | Appp(gamma, t1, t2, ty, cons, p1,p2) ->   (count_let_var_occurrence let_var t1) +  (count_let_var_occurrence let_var t2)
    | Absp(gamma,  var, t1, ty, cons,p1) ->  if (var = let_var) then count_let_var_occurrence let_var (subst (var,  Var (newvar var)) t1) else  (count_let_var_occurrence let_var t1) 
    | Axiomp(gamma, var, ty, cons) -> if (var = let_var) then 1 else 0
;;



let filter_polymorphic_lets l p =  List.filter (fun x -> (count_number_of_occurrence_proof x p) > 1) l;;



let rec convert_one_bound_occurrence_term let_var t b =
 match t with 
     Var x -> if let_var = x && b = false then let new_var = newvar x in (Var new_var, true, new_var) else (t,b,"")
   | Abs(x,t1) -> 
       if x = let_var then (Abs (x, t1), b,"")
     else let (new_term, new_bool,str) = convert_one_bound_occurrence_term let_var t1 b in (Abs(x, new_term), new_bool, str) 
   | Ap(t1,t2)->  
       let (new_term1, new_bool1,str) = convert_one_bound_occurrence_term let_var t1 b in 
       if (new_bool1 = false) then
       let (new_term2, new_bool2, str) = convert_one_bound_occurrence_term let_var t2 new_bool1 in  
       (Ap (new_term1,new_term2), new_bool2,str)
       else (Ap (new_term1,t2),new_bool1,str)

   | Let(x,t1,t2) -> 
       if x = let_var && b = false
       then let (new_term, new_bool,str) = convert_one_bound_occurrence_term let_var t2 b in (Let (str, t1, Let(x,t1, new_term)), new_bool,str) 
       else  let (new_term1, new_bool1,str) = convert_one_bound_occurrence_term let_var t1 b in 
	 if (new_bool1 = false) 
	 then 
	   let (new_term2, new_bool2,str) = convert_one_bound_occurrence_term let_var t2 new_bool1 in  
	     (Let (x,new_term1,new_term2), new_bool2,str)
	 else (Let (x,new_term1,t2),new_bool1, str)
;;

(*
*let rec convert_one_bound_occurrence_proof let_var p b =
*  match p with 
*    |Emptyp -> (p,b)
*    |  Letp (gamma, var, t1, t2, ty, cons, p1, p2) -> if (var  = let_var && b = false) then  (convert_one_bound_occurrence_term let_var t2 false) 
*       else if (var = let_var && b = true)
*	 | Appp(gamma, t1, t2, ty, cons, p1,p2) ->   (count_let_var_occurrence let_var t1) +  (count_let_var_occurrence let_var t1)
*	 | Absp(gamma,  var, t1, ty, cons,p1) ->  if (var = let_var) then count_let_var_occurrence let_var (subst (var,  Var (newvar var)) t1) else  (count_let_var_occurrence let_var t1) 
*	 | Axiomp(gamma, var, ty, cons) -> if (var = let_var) then 1 else 0
*	    ;;
 *)	   

convert_one_bound_occurrence_term "i" (Let ("i" ,Abs("x", Var "x"),  (Ap (Ap (Var "i", Var "i"), Var "i")))) false;; 
 

(* count_number_of_occurrence  "x" (Ap(Var "y", Var "x"));;*)




   



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

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


let ftv_cons c =
match c with 
  | Equality (ty1,ty2) -> (t_fv ty1)@(t_fv ty2)
  | MakeScheme (ty1, ty2,gamma) -> (t_fv ty1) @ (t_fv ty2)
  | Instance (ty1,ty2) -> (t_fv ty1)@(t_fv ty2)

;;

let extract_free_type_var (Env e) = List.fold_right (fun x -> fun y -> unique ((t_fv x)@y))(List.map (fun(x,y) ->  y) e) [];;

let extract_free_term_var (Env e) = List.fold_right (fun x -> fun y -> x::y)(List.map (fun(x,y) ->  x) e) [];;

let extract_free_type_variables  p =
let rec extract_ftv_proof  p =
  match p with
    |Emptyp -> []
    |  Letp (gamma, var, t1, t2, ty, cons, p1, p2) -> (extract_free_type_var gamma)@(ftv_cons cons) @ (extract_ftv_proof p1)@ (extract_ftv_proof p2)
    | Appp(gamma, t1, t2, ty, cons, p1,p2) ->  (extract_free_type_var gamma)@(extract_ftv_proof p1)@ (extract_ftv_proof p2)
    | Absp(gamma,  var, t1, ty, cons,p1) ->  (extract_free_type_var gamma)@(ftv_cons cons) @ (extract_ftv_proof p1)
    | Axiomp(gamma, var, ty, cons) -> (extract_free_type_var gamma)@(ftv_cons cons)
in unique (extract_ftv_proof p)
;;


(*
 * extract_free_type_var (Env [("x",Poly "a");("y", Arrow (Poly "a", Poly "b"))]);;
 * extract_free_term_var (Env [("x",Poly "a");("y", Arrow (Poly "a", Poly "b"))]);;
*)

let rec collect_let_bound_vars p = 
  match p with
    |Emptyp -> []
    |  Letp (gamma, var, t1, t2, ty, cons, p1, p2) -> [var] @ (collect_let_bound_vars p1)@ (collect_let_bound_vars p2)
    | Appp(gamma, t1, t2, ty, cons, p1,p2) ->  (collect_let_bound_vars p1)@ (collect_let_bound_vars p2)
    | Absp(gamma,  var, t1, ty, cons,p1) ->  (collect_let_bound_vars p1)
    | Axiomp(gamma, var, ty, cons) -> []

;;



let rec cut_proof p t =
  match p with
    |Emptyp -> Emptyp
    |  Letp (gamma, var, t1, t2, ty, cons, p1, p2) -> 
	 if (equal_term t t1)
	 then p1
         else if (equal_term t t2)
         then p2
         else let right_result = (cut_proof p1 t) in 
	   ( match right_result with
		 Emptyp -> (cut_proof p2 t)
               | p -> p
	   )
	     
    | Appp(gamma, t1, t2, ty, cons, p1,p2) ->  
	if (equal_term t t1) 
	then p1 
	else if (equal_term t t2)
	then p2
	else  let right_result = (cut_proof p1 t) in 
	   ( match right_result with
		 Emptyp -> (cut_proof p2 t)
               | p -> p
	   )
    | Absp(gamma,  var, t1, ty, cons,p1) ->  if (equal_term t t1) then p1 else (cut_proof p1 t)
    | Axiomp(gamma, var, ty, cons) -> if (equal_term t (Var var)) then p else Emptyp
;; 


                                     
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 ->( match y with
                     Poly c -> if a = c then t else t1
                   | _ -> raise (Failure ("t_subst -- trying to substitute something besides a Poly: " ^ (convert_type_to_string y)^(convert_type_to_string t1)))
               )
    | TypeSchemePoly a ->  t1 
(*
 *  ( match y with
 *                    TypeSchemePoly 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 List.mem 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
    | 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 (Poly y, Poly x) ty4) 
    | _ -> false
;;



let rec t_member a ty =
  match ty with 
    Poly b -> a = b 
  | TypeSchemePoly b -> a = b   
  | Arrow (ty1,ty2) -> ((t_member a ty1)|| (t_member a ty2))
  | All (x,ty1) -> List.mem a (t_fv ty1)  
;;





let rec apply_subst_env sub (Env e)  =
  let rec apply_subst_env_list (s,ty) l = 
    (match l with 
	 [] ->  (s,ty)
       | (ty1,ty2)::tail -> (apply_subst_env_list (s, t_subst (ty1,ty2) ty ) tail)
    ) in
    match e  with 
	[] ->  [] 
      | (s,ty)::tail_env -> (apply_subst_env_list (s,ty) sub ):: (apply_subst_env sub (Env tail_env))
;;

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 c  s = 
 ( match c with 
	  Equality (ty1,ty2) ->(match s with 
				    [] -> []
                                  
				  | Equality (ty3,ty4)::t -> 
				      (Equality (t_subst (ty1,ty2) ty3 , t_subst (ty1,ty2) ty4)):: (transform_constraint_store c t)
				  | MakeScheme(ty3,ty4,e)::t -> (MakeScheme(t_subst(ty1,ty2) ty3, t_subst(ty1,ty2)ty4,Env (apply_subst_env [(ty1,ty2)] e)))::(transform_constraint_store c t)
				  |Instance(ty3,ty4)::t -> Instance(t_subst(ty1,ty2)ty3,t_subst(ty1,ty2)ty4)::(transform_constraint_store c t)
                               )
	| h -> raise( Failure ("Constraints other than equality not handled2"))
     )
;;

let rec local_transform_constraint_store c s =  
 ( match c with 
	  (ty1,ty2) ->(match s with 
				    [] -> []
                                    
				  | Equality (ty3,ty4)::t -> 
				      (Equality (t_subst (ty1,ty2) ty3 , t_subst (ty1,ty2) ty4)):: (local_transform_constraint_store c t)
				  | MakeScheme(ty3,ty4,e)::t -> (MakeScheme(t_subst(ty1,ty2) ty3, t_subst(ty1,ty2)ty4,Env (apply_subst_env [(ty1,ty2)] e)))::(local_transform_constraint_store c t)
				  |Instance(ty3,ty4)::t -> Instance(t_subst(ty1,ty2)ty3,t_subst(ty1,ty2)ty4)::(local_transform_constraint_store c t)
                               )
	
     )
;;

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

let rec transform_constraint_store_subst_list c s = 
  match  s with
    [] -> c 
    | h::t -> (transform_constraint_store_subst_list (local_transform_constraint_store h c) t )
;;    



 (* Terms are actually types so we have a substitution happening at the type level*)
  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 s1 s2 =
 * let rec  collect_free_terms_list (ty1,ty2) s1  =
 * ( match s1 with
 *	[] -> (ty1,ty2)
 *   | (ty3,ty4)::tail34 ->  if (equal_types ty1 ty3) then (ty3,ty4) else (collect_free_terms_list (ty1,ty2) tail34) 
 *  ) in 
 *  match s2 with 
 *      [] -> s1
 *    | (ty1,ty2)::tail12 ->  (collect_free_terms_list (ty1,ty2) s1)::(collect_free_terms s1  tail12)
 * ;;
 *
*)



 let rec extract_list l = List.fold_right (fun x -> fun y -> x::y)(List.map (fun(x,y) ->  x) l) [];;

extract_list [(Poly "a", Poly "b"); (Poly "c",Poly "d")];;

  let rec collect_free_terms s1 s2 =
 match s2 with 
       [] -> s1
     | (ty1,ty2)::tail12 ->  if (List.mem ty1 (extract_list s1)) then (collect_free_terms s1 tail12) else (ty1,ty2):: (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
;;   

compose_subst [][];;
collect_free_terms [(Poly "x",Poly "y")] [(Poly "x", Poly "z")];;
compose_subst [(Poly "x",Poly "y")] [];;

  
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"))
                           | TypeSchemePoly a ->  raise (Failure "Type scheme variable should not be substituted for")
                           | 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 ("Applying substitution to a type scheme ")))
;;




(* 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 -> (t_fv ty)@ (t_fv_env t) *)
    | ((s:string),ty)::t ->  
  if (String.get s 0) =='s'then  (t_fv_env t) else  (t_fv ty)@ (t_fv_env t)

(*
     (match ty with 
	 Poly a -> (t_fv ty)@ (t_fv_env t)
       | TypeSchemePoly a -> (t_fv_env t)
       |  _ -> raise (Failure (" Type Env cannot contain binding for copound types"))
     )
*)
;;

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 ((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 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 
			(*		      (Poly a ,Poly b)  -> if a = b then unifi tail s *)
		      | Equality (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 types1") 
						 else unifi (transform_constraint_store  (Equality (Poly a,ty)) tail)  ((transform_subst_list  (Poly a, ty) s)@[(Poly a ,ty)])
					
						     | _ ->  if (t_member a ty) 
						       then raise (Failure ("Occur check failure: trying to unify"^ (convert_type_to_string ty)^" with Poly "^a  ))
						 else unifi (transform_constraint_store  (Equality (Poly a,ty)) tail)  ((transform_subst_list  (Poly a, ty) s)@[(Poly 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 " Unknown type to be unified")
		   )
    )
;;



let print_bool t = if t then "true" else "false";;

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  
			 let updated_store_2= new_constraint_store in 
			   replace_sch_var (ty1,ty2,env) t updated_store_2 )
		  | Equality(ty4,ty5) -> let updated_store_3 = new_constraint_store in ( replace_sch_var (ty1,ty2,env) t updated_store_3 )
		  | MakeScheme(ty3,ty4,env) -> let updated_store_4 = new_constraint_store in ( replace_sch_var (ty1,ty2,env) t updated_store_4)
	      )
;;







let addStoreFront const store = const::store;;

let addStoreBack const store = store@[const];;


      
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 (addStoreFront (Equality (ty,unroll envty)) constraint_store) 
					     | somety -> 
						 (match somety with 
						    |Poly u -> generate_constraint tail_goal (addStoreFront (Equality (somety,ty)) constraint_store)
						    |TypeSchemePoly a -> generate_constraint tail_goal (addStoreFront (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) 
                                            (addStoreFront (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_goal1 = ((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_constraint_store1 = 
	*				    addStoreBack (MakeScheme (new_type_var1,  
	*							       new_type_var2, 
	*							       (Env e))) constraint_store in
        *                                  				   generate_constraint  (new_goal1::((new_env,t2,ty)::tail_goal)) new_constraint_store1  
        *
        *)
		
			      
					      
	 				  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 = 
					    addStoreBack (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 ([]:(types* types) list)
  else 
    match c   with
	[]  -> ([]:(types*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 (transform_constraint_store_subst_list 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 (equal_types (Poly "jk") a ) 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)
;;

(* todo: use exceptions to check if  a branch return null *)

let rec extract_let_bound_term  let_var t =
match t with
  |  Let(x,t1,t2) -> 
       if x = let_var 
       then (t1,true) 
       else  (let (term,result1) = extract_let_bound_term let_var t1 in 
                    if (result1 = false) then  (extract_let_bound_term let_var t2)   else (term,result1))
  |  Abs (x,t1)   -> extract_let_bound_term let_var t1
  |  Ap (t1,t2) -> let (term,result1) = extract_let_bound_term let_var t1 in 
                    if (result1 = false) then  (extract_let_bound_term let_var t2)   else (term,result1)
  |  Var x ->  (Var "notfound", false)
  
  ;;




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 = 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 c = 
	try (List.assoc  a  table)  
	with Not_found -> newchar()	   
      in if List.mem (a,c) table
        then (Poly (Char.escaped c), table) else  (Poly (Char.escaped c),((a,c)::table))
  | TypeSchemePoly a -> (t,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 c = try (List.assoc x table) with Not_found -> newchar()
      in if List.mem (x,c) table
        then let (ty,mod_table) = transform_type (t,table) in (All(Char.escaped c,ty), mod_table)
        else let (str,new_mod_table) = transform_type (t,(x,c)::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 
      (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")));; 



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





(* Example 2 *)
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"))))));;
 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")))))));;
print_good_type  (Abs ("f", Abs("x", Ap (Var "f", Ap (Var "f", Var "x")))));;