(* 
 * Implements an  Wand's algorithm for a language based on 
 * pure untyped lambda calculus as given in the paper:
 * "A simple algorithm and proof for type inference".
 * 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 *)

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



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



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






                                     
let extendEnv  id ty (Env e) = 
  match e with 
      [] -> Env ([(id,ty)])
    | l  -> Env ((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 fv t =
  match t with
      Var (x) -> [x]
    | Abs (x,t) -> remove_all x (fv t)
    | Ap(t1,t2) -> (fv t1)  @  (fv t2)

;;





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

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_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 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]
      | Arrow(a,b) ->  (fv a) @ fv b
      | All(x,a) ->  remove_all  x (fv a)
  in
    unique (fv ty)
;;

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 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
    | 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 
  | Arrow (ty1,ty2) -> ((t_member a ty1)|| (t_member a ty2))
  | All (x,ty1) -> List.mem a (t_fv ty1)  
;;


let print_env (Env e) = List.fold_right (fun x -> fun y -> x^y)(List.map (fun(x,y) ->  x^"="^(print_type y)^" ") e) "";;


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 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 (ty1,ty2)  s =
  match s with
      [] -> []
    | (ty3,ty4)::t -> (t_subst (ty1,ty2) ty3 , t_subst (ty1,ty2) ty4):: (transform_constraint_store (ty1,ty2) 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) l  =
( match l with
	[] -> [(ty1,ty2)]
      | (ty3,ty4)::tail34 ->  if (equal_types ty1 ty3) then [] 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 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 ("Applying substitution to a type scheme ")))
;;






let rec t_fv_env  env =
  match env with 
      [] -> []
    |  ((s:string),ty)::t -> (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 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 unifi c s =
  match c   with
      []  -> (s,c)  
    | 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 types1") 
					      else unifi (transform_constraint_store  (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"^ (print_type ty)^" with Poly "^a  ))
					    else unifi (transform_constraint_store  (Poly a,ty) tail) ((transform_subst_list (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
		    | _  -> raise (Failure " Unknown type to be unified")
		 )
    
;;



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





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 
					       |Poly u -> generate_constraint tail_goal ((envty,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 (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 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 
      fun_apply  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 = 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))
		

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

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



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



(* Generate constriant for compose operator *)
print_good_type (Abs ("x", (Abs ("y", (Abs ("z",  (Ap (Var "x", (Ap (Var "y", Var "z"))))))))));;


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

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


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

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