(*
 * Konzepte höherer Programmiersprachen (SS 2006)
 *
 * Übungsblatt 9 / Aufgabe 1
 *)

(* ------------------------------------------------------------------------------------------- *)
(* Typdefinitionen                                                                             *)
(* ------------------------------------------------------------------------------------------- *)

type ptype =
    BOOL
  | INT
  | UNIT;;

type ttype =
    TVAR of int
  | TPRIM of ptype
  | TPAIR of ttype * ttype
  | TARROW of ttype * ttype;;


(* ------------------------------------------------------------------------------------------- *)
(* Substitutionen                                                                              *)
(* ------------------------------------------------------------------------------------------- *)

(* Substitution ordnet jeder Typvariable einen ttype zu, z.B. wieder eine TVAR *)
type substitution = int -> ttype;;

(* Ordnet jeder TVAR sich selbst zu *)
let empty_subst = fun n -> TVAR n;;

let mk_subst i tau = fun n -> if n = i then tau else TVAR n;;

let rec apply_to_type (s: substitution) tau =
  match tau with
      TVAR n -> s n
    | TPRIM _ -> tau
    | TPAIR (tau1, tau2) -> TPAIR (apply_to_type s tau1, apply_to_type s tau2)
    | TARROW (tau1, tau2) -> TARROW (apply_to_type s tau1, apply_to_type s tau2);;

let compose_subst (s1: substitution) (s2: substitution) =
  fun n -> apply_to_type s2 (s1 n);;


(* --------------------------------------------------------------------------- *)
(* Anwendung von s auf Typgleichungen                                          *)
(* --------------------------------------------------------------------------- *)

(* Typgleichung: tau1 = tau2 *)
type teqn = ttype * ttype;;

let apply_to_eqn (s: substitution) (tau, tau': teqn): teqn =
  apply_to_type s tau, apply_to_type s tau';;

let apply_to_eqns (s: substitution) = List.map (apply_to_eqn s);;


(* --------------------------------------------------------------------------- *)
(* occur check                                                                 *)
(* --------------------------------------------------------------------------- *)

(* Tested ob TVAR i in tau auftaucht *)
let rec occurs i tau = 
  match tau with 
      TVAR n -> i = n
    | TPRIM _ -> false
    | TPAIR (tau1, tau2) 
    | TARROW (tau1, tau2) -> occurs i tau1 || occurs i tau2;;


(* --------------------------------------------------------------------------- *)
(* Unifikationsalgorithmus                                                     *)
(* --------------------------------------------------------------------------- *)

exception Type_error;;

let rec unify (eqns: teqn list): substitution =
  match eqns with 
      [] -> empty_subst 
    | eqn :: eqns' ->
	match eqn with 
            TVAR n, tau
	  | tau, TVAR n -> 
	      if not (occurs n tau) || tau = TVAR n
	      then 
		let s1 = mk_subst n tau in 
		let s2 = unify (apply_to_eqns s1 eqns')
		in compose_subst s1 s2
	      else 
		raise Type_error
	  | TPAIR (tau1, tau2), TPAIR (tau1', tau2')
	  | TARROW (tau1, tau2), TARROW (tau1', tau2') -> 
	      unify ((tau1, tau1') :: (tau2, tau2') :: eqns')
	  | tau, tau' -> 
	      if tau = tau' 
	      then unify eqns'
	      else raise Type_error;;


(* --------------------------------------------------------------------------- *)
(* Typinferenzalgorithmus                                                      *)
(* --------------------------------------------------------------------------- *)

(* Operatoren: arithmetische und relationale *)
type aoperator = Add | Sub | Mul | Div | Mod;;
type roperator = Eq | Less | Grt | Leq | Grq;;

(* Konstanten *)
type constant =
    Unit
  | Int of int
  | Bool of bool
  | Aop of aoperator
  | Rop of roperator;;

(* Identifier als Strings *)
type identifier = string;;

(* Ausdrücke *)
type expression =
    Const of constant
  | Identifier of identifier
  | If of expression * expression * expression
  | App of expression * expression
  | Abstr of identifier * expression
  | Let of identifier * expression * expression;;

(* Typen für Konstanten *)
let constant_type c =
  match c with
      Unit -> TPRIM (UNIT)
    | Int n -> TPRIM (INT)
    | Bool b -> TPRIM (BOOL)
    | Aop op -> TARROW (TPRIM (INT), TARROW (TPRIM (INT), TPRIM (INT)))
    | Rop op -> TARROW (TPRIM (INT), TARROW (TPRIM (BOOL), TPRIM (BOOL)));;

(* Typumgebung *)
type type_environment = (identifier * ttype) list;;

exception Identifier_error of identifier;;

let extend gamma id tau = (id, tau) :: gamma;;

let rec lookup gamma id =
  match gamma with
      [] -> raise (Identifier_error id)
    | (id', tau') :: gamma' -> if id = id' then tau' else lookup gamma' id;;

let rec apply_to_env (s:substitution) (gamma:type_environment) =
  match gamma with
      (id', t') :: gamma' -> (id', (apply_to_type s t')) :: (apply_to_env s gamma')
    | [] -> [];;

(* neue typ variablen *)
let next_tvar = ref 0;;

let rec solution gamma e tau =
  match e with
      Const c -> unify [tau, constant_type c]
    | Identifier id -> unify [tau, lookup gamma id]
    | If (e0, e1, e2) -> 
	let s0 = solution gamma e0 (TPRIM BOOL) in
	let s1 = solution (apply_to_env s0 gamma) e1 (apply_to_type s0 tau) in
	let s2 = solution
	  (apply_to_env s1 (apply_to_env s0 gamma)) e2
	  (apply_to_type s1 (apply_to_type s0 tau)) in
	  compose_subst s0 (compose_subst s1 s2)
    | App (e1, e2) ->
	let tvar = (next_tvar := !next_tvar + 1; !next_tvar) in
	let s1 = solution gamma e1 (TARROW (TVAR tvar, tau)) in
	let s2 = solution (apply_to_env s1 gamma) e2 (s1 tvar) in
	  compose_subst s1 s2
    | Abstr (id, e1) ->
	let tvar1 = (next_tvar := !next_tvar + 1; !next_tvar) in
	let tvar2 = (next_tvar := !next_tvar + 1; !next_tvar) in
	let s1 = solution (extend gamma id (TVAR tvar1)) e1 (TVAR tvar2) in
	let s2 = unify [(apply_to_type s1 tau), (TARROW (s1 tvar1, s1 tvar2))] in
	  compose_subst s1 s2
    | Let (id, e1, e2) ->
	let tvar = (next_tvar := !next_tvar + 1; !next_tvar) in
	let s1 = solution gamma e1 (TVAR tvar) in
	let s2 = solution (apply_to_env s1 (extend gamma id (s1 tvar))) e2 (apply_to_type s1 tau) in
	  compose_subst s1 s2;;

let infer_type e =
  let tvar = (next_tvar := !next_tvar + 1; !next_tvar) in
  let s = solution [] e (TVAR tvar) in
    s tvar;;

infer_type (Abstr ("id", Let ("x", Const (Int (0)), App (Identifier ("id"), Identifier ("x")))));;