open Gui

(* var_t *)
module VarS = struct
  type 'var_t t = V of string

  let map f v = match v with V (s) -> V (s)

  let search _ = false

  let unify f v1 v2 = ()

  let draw v = match v with V(s) -> create_str s
end

module Var  = Fix.F (VarS)

let var n = Var.make_fix (VarS.V (n))

(* tvar_t *)
module TVarS = struct
  type 'tvar_t t = TV of string

  let map f v = match v with TV (s) -> TV (s)

  let search _ = false

  let unify f v1 v2 = ()

  let draw v = match v with TV(s) -> create_str s
end

module TVar = Fix.F (TVarS)

let tv x = TVar.make_fix (TVarS.TV (x))

(* type_t *)
module TypeS = struct
  type 'type_t t =
      TVar of TVar.fix_t
    | Fun of 'type_t * 'type_t
    | UT of TVar.fix_t * 'type_t

  let map f tp = match tp with
      TVar(tv) -> TVar(tv)
    | Fun(tp1, tp2) -> Fun(f tp1, f tp2)
    | UT(tv, tp) -> UT(tv, f tp)

  let search tp = match tp with
      TVar(tv) -> TVar.search tv
    | Fun(tp1, tp2) -> tp1 || tp2
    | UT(tv, tp) -> TVar.search tv || tp

  let unify f tp1 tp2 = match (tp1, tp2) with
      (TVar(tv1), TVar(tv2)) -> TVar.unify tv1 tv2
    | (Fun(tp1, tp2), Fun(tp3, tp4)) ->
	 f tp1 tp3; f tp2 tp4
    | (UT(tv1, tp1), UT(tv2, tp2)) ->
	 TVar.unify tv1 tv2; f tp1 tp2
    | _ -> raise (Unify_Error "of Type")

  let draw tp = match tp with
      TVar(tv) -> TVar.draw tv
    | Fun(tp1, tp2) ->
	combineH [create_str "("; tp1; create_str "→"; tp2; create_str ")"]
    | UT(tv, tp) -> combineH [
	 create_str "∀"; TVar.draw tv;
	 create_str "."; tp]
end

module Type = Fix.F (TypeS)

let tvar tv = Type.make_fix (TypeS.TVar (tv))
let tfun tp1 tp2 = Type.make_fix (TypeS.Fun (tp1, tp2))
let ut tv tp = Type.make_fix (TypeS.UT (tv, tp))

(* term_t *)
module TermS = struct
  type 'term_t t =
      Var of Var.fix_t
    | Abs of Var.fix_t * Type.fix_t * 'term_t
    | App of 'term_t * 'term_t
    | TAbs of TVar.fix_t * 'term_t
    | TApp of 'term_t * Type.fix_t

  let map f tm = match tm with
      Var(v) -> Var(v) (* もう一度書かないと、型エラー *)
    | Abs(v, tp, f1) -> Abs(v, tp, f f1)
    | App(f1, f2) -> App(f f1, f f2)
    | TAbs(tv, f1) -> TAbs(tv, f f1)
    | TApp(f1, tp) -> TApp(f f1, tp)

  let search tm = match tm with
      Var(v) -> Var.search v
    | Abs(v, tp, tm) ->
	   let b1 = Var.search v in
	   let b2 = Type.search tp in
	   b1 || b2 || tm
    | App(tm1, tm2) -> tm1 || tm2
    | TAbs(tv, tm) -> TVar.search tv || tm
    | TApp(tm, tp) -> Type.search tp || tm

  let unify f tm1 tm2 =  match (tm1, tm2) with
      (Var(v1), Var(v2)) -> Var.unify v1 v2
    | (Abs(v1, tp1, f1), Abs(v2, tp2, f2)) ->
	   Var.unify v1 v2;
	   Type.unify tp1 tp2;
	   f f1 f2
    | (App(f1, f2), App(f3, f4)) -> f f1 f3; f f2 f4
    | (TAbs(tv1, f1), TAbs(tv2, f2)) ->
	   TVar.unify tv1 tv2; f f1 f2
    | (TApp(f1, tp1), TApp(f2, tp2)) ->
	   f f1 f2; Type.unify tp1 tp2
    | _ -> raise (Unify_Error "of term")

  let draw tm = match tm with
      Var(v) -> Var.draw v
    | Abs(v, tp, f) ->
	combineH [create_str "(λ"; Var.draw v; create_str ":";
		  Type.draw tp; create_str ". "; f; create_str ")"]
    | App(fix1, fix2) ->
	combineH [create_str "("; fix1; create_str " "; fix2 ; create_str ")"]
    | TAbs(tv, f) ->
	combineH [create_str "(Λ"; TVar.draw tv; create_str ".";
		  f; create_str ")"]
    | TApp(f, tp) ->
	combineH [create_str "("; f;
		  create_str "["; Type.draw tp; create_str "])"]
end

module Term = Fix.F (TermS)

let term_var v = Term.make_fix (TermS.Var (v))
let term_abs x t e = Term.make_fix (TermS.Abs (x, t, e))
let term_app e1 e2 = Term.make_fix (TermS.App (e1, e2))
let term_tapp tm tp = Term.make_fix (TermS.TApp (tm, tp))
let term_tabs tv tm = Term.make_fix (TermS.TAbs (tv, tm))

(* env_t *)
module EnvS = struct
  type elem_t =
      E of Var.fix_t * Type.fix_t
    | T of TVar.fix_t

  type 'env_t t =
      Empty
    | Cons of elem_t * 'env_t

  let map f env = match env with
      Empty -> Empty
    | Cons(e, env) -> Cons(e, f env)

  let search_e elm = match elm with
      E(v, tp) ->
	let b1 = Var.search v in
	let b2 = Type.search tp in
	b1 || b2
    | T(tv) -> TVar.search tv

  let search env = match env with
      Empty -> false
    | Cons(elm, env) -> search_e elm || env

  let unify_e e1 e2 = match (e1, e2) with
      (E(v1, tp1), E(v2, tp2)) -> Var.unify v1 v2; Type.unify tp1 tp2
    | (T(tv1), T(tv2)) -> TVar.unify tv1 tv2
    | _ -> raise (Unify_Error "of Env_e")

  let unify f env1 env2 = match (env1, env2) with
      (Empty, Empty) -> ()
    | (Cons(e1, env1), Cons(e2, env2)) -> unify_e e1 e2; f env1 env2
    | _ -> raise (Unify_Error "of Env")

  let draw_e e = match e with
      E(v, tp) -> combineH [Var.draw v; create_str ":"; Type.draw tp]
    | T(tv) -> TVar.draw tv

  let draw env = match env with
      Empty -> create_str "Φ"
    | Cons(e, env) -> combineH [ draw_e e; create_str ", "; env]
end

module Env = Fix.F (EnvS)

let empty = Env.make_fix (EnvS.Empty)
let cons ev e = Env.make_fix (EnvS.Cons(ev,e))
let env_e v tp = EnvS.E(v, tp)
let env_t tv = EnvS.T(tv)

(* judge_t *)
module JudgeS = struct
  type 'judge_t t =
      Judge of Env.fix_t * Term.fix_t * Type.fix_t
    | Equal of Type.fix_t * Type.fix_t

  let map _ j = match j with
      Judge(e, tm, tp) -> Judge(e, tm, tp)
    | Equal(tp1, tp2) -> Equal(tp1, tp2)

  let search j = match j with
      Judge(env, tm, tp) ->
	let b1 = Env.search env in
	let b2 = Term.search tm in
	let b3 = Type.search tp in
	b1 || b2 || b3
    | Equal(tp1, tp2) ->
	let b1 = Type.search tp1 in
	let b2 = Type.search tp2 in
	b1 || b2

  let unify f j1 j2 = match (j1, j2) with
      (Judge(env1,tm1,tp1), Judge(env2,tm2,tp2)) ->
	Env.unify env1 env2;
	Term.unify tm1 tm2;
	Type.unify tp1 tp2
    | (Equal(tp1,tp2), Equal(tp3,tp4)) ->
	Type.unify tp1 tp3;
	Type.unify tp2 tp4
    | _ -> raise (Unify_Error "of Judge")

  let draw j = match j with
      Judge(env, tm, tp) ->
	combineH [Env.draw env; create_str "├ ";
		  Term.draw tm; create_str " : "; Type.draw tp]
    | Equal(tp1, tp2) ->
	combineH[Type.draw tp1; create_str " = "; Type.draw tp2]
end

module Judge = Fix.F(JudgeS)
let judge e tm tp = Judge.make_fix (JudgeS.Judge (e, tm, tp))
let equal tp1 tp2 = Judge.make_fix (JudgeS.Equal (tp1, tp2))

(* 推論規則を定義する *)
(* 環境から変数に当たる型を取ってくる *)
let rec get env var = try
  (match Env.take env with
      EnvS.Empty -> raise Not_found
    | EnvS.Cons(e, rest) ->
	(match e with
	    EnvS.E(var2, tp) ->
	      if Var.take var = Var.take var2
	      then tp
	      else get rest var
	  | EnvS.T(ut) -> get rest var))
with _ -> Type.make_code
  (fun () -> combineH [create_str "get("; Env.draw env; create_str ";";
		       Var.draw var; create_str ")"])
  (fun () -> get env var)

(* TVAR *)
let t_var () =
  let x = Var.make_meta "x" in
  let t = Type.make_meta "tp" in
  let e = Env.make_meta "e" in
  Infer(judge e (term_var x) t,
	[Infer(equal t (get e x), [])])

(* TABS *)
let t_abs () =
  let e = Env.make_meta "e" in
  let x = Var.make_meta "x" in
  let tm = Term.make_meta "t" in
  let tp1 = Type.make_meta "T" in
  let tp2 = Type.make_meta "T" in
  Infer(judge e (term_abs x tp1 tm) (tfun tp1 tp2),
	[Infer(judge (cons (env_e x tp1) e) tm tp2,[])])

(* TAPP *)
let t_app () =
  let e = Env.make_meta "e" in
  let tm1 = Term.make_meta "t" in
  let tm2 = Term.make_meta "t" in
  let tp1 = Type.make_meta "T" in
  let tp2 = Type.make_meta "T" in
  Infer(judge e (term_app tm1 tm2) tp2,
	[Infer(judge e tm1 (tfun tp1 tp2), []);
	 Infer(judge e tm2 tp1,[])])

(* TEQUAL *)
let t_equal () =
  let tp = Type.make_meta "T" in
  Axiom (equal tp tp)

(* 置換関数 つづきを書く T[X → A] *)
let rec subst origin target sub = try
  (match TVar.take target with
      TVarS.TV(str1) ->
	(match Type.take origin with
	    TypeS.TVar(tv) ->
	      let TVarS.TV(str2) = TVar.take tv in
	      if str1 = str2
	      then sub (* T = X *)
	      else origin (* T != X *)
	  | TypeS.Fun(tp1, tp2) ->
	      Type.make_fix (TypeS.Fun (subst tp1 target sub,
					subst tp2 target sub))
	  | TypeS.UT(tv, tp) ->
	      let TVarS.TV(str2) = TVar.take tv in
	      if str1 = str2
	      then origin
	      else Type.make_fix(TypeS.UT(tv, subst tp target sub))))
with _ -> Type.make_code
  (fun () -> combineH [create_str "subst("; Type.draw origin;
		       create_str ";"; TVar.draw target;
		       create_str ";"; Type.draw sub;
		       create_str ")"])
  (fun () -> subst origin target sub)

(* TTAPP *)
let t_tapp () =
  let env = Env.make_meta "e" in
  let tm = Term.make_meta "t" in
  let tp = Type.make_meta "T" in
  let tp1 = Type.make_meta "T" in
  let tp2 = Type.make_meta "T" in
  let x = TVar.make_meta "x" in
  let c = Type.make_meta "C" in
  Infer(judge env (term_tapp tm tp1) c,
	[Infer(judge env tm tp, []);
	 Infer(equal tp (ut x tp2),[]);
	 Infer(equal c (subst tp2 x tp1), [])])

(* TTABS *)
let t_tabs () =
  let env = Env.make_meta "e" in
  let x = TVar.make_meta "x" in
  let t2 = Term.make_meta "t" in
  let tp2 = Type.make_meta "T" in
  Infer(judge env (term_tabs x t2) (ut x tp2),
	[Infer(judge (cons (env_t x) env) t2 tp2, [])])

(* 登録される推論規則 *)
let infer_rule_list =
  [("TVAR", t_var); ("TABS", t_abs); ("TTABS", t_tabs);
   ("TAPP", t_app); ("TTAPP", t_tapp); ("TEQUAL", t_equal)]