(* 
 * 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".
 * 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) =  Env ((id,ty)::(List.remove_assoc id e));;


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 newvar  =
  let count = ref (-1) in
  function(s) -> count := !count + 1;
           s ^ (string_of_int !count)
;;



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


let rec print_simple_type t =
  match t 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 apply_subst_types s ty =
 match s with
     [] -> ty
   | (str,ty2)::tail -> (match ty with
                           | Poly a -> if a = str then  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) -> raise (Failure ("Applying substitution to a type scheme ")))
;;


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 apply_subst_env s (Env e) = List.map (fun ((x:string),y) -> (x, apply_subst_types s y)) e;;


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



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 subst_constraint_store (ty1,ty2)  s =
  List.map (fun (x,y) -> t_subst (ty1,ty2) x, t_subst (ty1,ty2) y) s 


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


let unify c =  
  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 c []
;;





      
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) 
                                            ((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 a="jk" then b else fun_apply t
;; 



let type_infer t  =
  let initial_const_store = generate_constraint [(Env [],t,Poly "jk")] [] in
  let subst_list  =  unify initial_const_store  in 
      fun_apply  subst_list;;



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




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