(** 
  * the Nominal Unification algorithm by Urban, Pitts & Gabbay 
  * June, 2003
*)

open List;;
exception Fail;;                               (* exception for unification failure *) 

(* terms *)
type term = 
    Unit                                       (* units *)
  | Atm   of string                            (* atoms *)
  | Susp  of (string * string) list * string   (* suspensions *)          
  | Fun   of string * term                     (* function symbols *)
  | Abst  of string * term                     (* abstracted terms *)
  | Pair  of term * term;;                     (* pairs *)  

(* swapping operation on atoms *)
let swap (a,b) c =                             
  if a=c then b else if b=c then a else c;;
let swaps pi c = fold_right swap pi c;;

(* permutation on terms *)
let rec perm pi t =       
  match t with	
      Unit              -> Unit
    | Atm(a)            -> Atm(swaps pi a)
    | Susp(pi',x)       -> Susp(pi@pi',x)
    | Fun(name,t')      -> Fun(name, perm pi t') 
    | Abst(a,t')        -> Abst(swaps pi a,perm pi t')
    | Pair(t1,t2)       -> Pair(perm pi t1,perm pi t2);;	

(* substitution operation 
 * - implements the simple operation of "hole filling" or "context substitution"
 *)
let rec subst t sigma =     
  match t with 
      Unit           -> Unit
    | Atm(a)         -> Atm(a)
    | Susp(pi,y)     -> (try (perm pi (assoc y sigma)) with Not_found -> Susp(pi,y))
    | Fun(name,t')   -> Fun(name,subst t' sigma)
    | Abst(a,t')     -> Abst(a,subst t' sigma)
    | Pair(t1,t2)    -> Pair(subst t1 sigma,subst t2 sigma);;
	
(* substitution composition *)
let rec subst_compose sigma = 
  match sigma with
      [] -> []
    | h::tail -> h::(map (fun (v,t) -> (v,subst t [h])) (subst_compose tail));;

(* occurs check *)
let rec occurs x t =
  match t with 
      Unit                   -> false
    | Atm(a)                 -> false
    | Susp(pi,y)             -> if x=y then true else false
    | Fun(_,t') | Abst(_,t') -> occurs x t' 
    | Pair(t1,t2)            -> (occurs x t1) || (occurs x t2);;

(* deletes duplicates from a list 
 *  - used for calculating disagreement sets
 *)
let rec delete_dups l =             
  match l with                               
      []   -> []
    | h::t -> let t' = delete_dups t in if mem h t' then t' else (h::t');;

(* disagreement set 
 * - takes two permutation (lists of pairs of atoms) as arguments
 *) 
let rec ds pi pi' =    
  let (l1,l2) = split (pi@pi') 
  in filter (fun a -> (swaps pi a)!=(swaps pi' a)) (delete_dups (l1@l2))

(* eliminates a solved equation *) 
let eliminate (v,t) (eprobs,fprobs) =  
  if occurs v t
  then raise Fail
  else (map (fun (t1,t2) -> (subst t1 [(v,t)],subst t2 [(v,t)])) eprobs,
        map (fun (a,t')  -> (a,subst t' [(v,t)])) fprobs);;

(***************)
(* unification *)
(***************)

(* checks and solves all freshness problems *)
let rec check fprobs =                        
  match fprobs with
      [] -> []
    | (a,Unit)::tail        -> check tail
    | (a,Atm(b))::tail      -> if a=b then raise Fail else check tail 
    | (a,Susp(pi,x))::tail  -> (swaps (rev pi) a,x)::(check tail)
    | (a,Fun(_,t))::tail    -> check ((a,t)::tail)
    | (a,Abst(b,t))::tail   -> if a=b then (check tail) else (check ((a,t)::tail))                  
    | (a,Pair(t1,t2))::tail -> check ((a,t1)::(a,t2)::tail);;  

(* solves all equational problems *)
let rec solve eprobs fprobs = solve_aux eprobs fprobs [] 
and solve_aux eprobs fprobs sigma =
  match eprobs with 
      [] -> (subst_compose sigma, delete_dups (check fprobs))
    | (Unit,Unit)::tail -> solve_aux tail fprobs sigma
    | (Atm(a),Atm(b))::tail -> if a=b then (solve_aux tail fprobs sigma) else raise Fail
    | (Susp(pi,x),Susp(pi',x'))::tail when x=x' ->                     
        let new_fps = map (fun a -> (a,Susp([],x))) (ds pi pi') 
	in solve_aux tail (new_fps @ fprobs) sigma 
    | (Susp(pi,x),t)::tail | (t,Susp(pi,x))::tail ->
	let new_sigma = (x,perm (rev pi) t) in
	let (new_eprobs,new_fprobs) = eliminate new_sigma (tail,fprobs) 
	in solve_aux new_eprobs new_fprobs (new_sigma::sigma)
    | (Fun(n1,t1),Fun(n2,t2))::tail -> 
	if  n1 = n2 
	then solve_aux ((t1,t2)::tail) fprobs sigma
	else raise Fail
    | (Abst(a1,t1),Abst(a2,t2))::tail ->
	if a1 = a2 
	then solve_aux ((t1,t2)::tail) fprobs sigma
	else solve_aux ((t1,perm [(a1,a2)] t2)::tail) ((a1,t2)::fprobs) sigma
    | (Pair(t1,t2),Pair(s1,s2))::tail ->
	solve_aux ((t1,s1)::(t2,s2)::tail) fprobs sigma
    | _ -> raise Fail;;

   
(************)
(* Examples *)
(************)

(* a few variables*)
let x = Susp([],"X")
and y = Susp([],"Y")
and z = Susp([],"Z");;

(* lam a.(X a) =? lam b.(c b)     --> [X:=c] *)
let t1 = Abst("a",Pair(x,Atm("a")));;
let t2 = Abst("b",Pair(Atm("c"),Atm("b")));;
solve [(t1,t2)] [];; 

(* lam a.(X a) =? lam b.(a b)     --> fails      *)
let t1 = Abst("a",Pair(x,Atm("a")));;
let t2 = Abst("b",Pair(Atm("a"),Atm("b")));;
solve [(t1,t2)] [];;    

(* lam a.(X a) =? lam b.(X b)     --> a#X, b#X   *)
let t1 = Abst("a",Pair(x,Atm("a")));;
let t2 = Abst("b",Pair(x,Atm("b")));;
solve [(t1,t2)] [];;    

(* lam a.(X a) =? lam b.(Y b)     --> [X:=(a b)oY]  a#Y *)
(*                                --> [Y:=(b a)oX]  b#X *)
let t1 = Abst("a",Pair(x,Atm("a")));;
let t2 = Abst("b",Pair(y,Atm("b")));;
solve [(t1,t2)] [];;   
solve [(t2,t1)] [];;

(* quiz-questions from the paper *)
let m1 = Susp([],"M1")
and m2 = Susp([],"M2")
and m3 = Susp([],"M3")
and m4 = Susp([],"M4")
and m5 = Susp([],"M5")
and m6 = Susp([],"M6")
and m7 = Susp([],"M7");;

(* 1 --> fail *)
let t1 = Abst("a",Abst("b",Pair(m1,Atm("b"))));;
let t2 = Abst("b",Abst("a",Pair(Atm("a"),m1)));;
solve [(t1,t2)] [];;

(* 2 --> [M2:=b ,M3:=a] *)
let t1 = Abst("a",Abst("b",Pair(m2,Atm("b"))));;
let t2 = Abst("b",Abst("a",Pair(Atm("a"),m3)));;
solve [(t1,t2)] [];;

(* 3 --> [M4:=(a b)o M5] *)
let t1 = Abst("a",Abst("b",Pair(Atm("b"),m4)));;
let t2 = Abst("b",Abst("a",Pair(Atm("a"),m5)));;
solve [(t1,t2)] [];;

(* 4 --> [M6:=(b a)oM7] b#M7 *)
let t1 = Abst("a",Abst("b",Pair(Atm("b"),m6)));;
let t2 = Abst("a",Abst("a",Pair(Atm("a"),m7)));;
solve [(t1,t2)] [];;