open Genlex
open List


(* ------------- Data structures ------------------- *)
(* Representation of terms:
   Example f(x g(c())) becomes App ("f", [Var "x"; App ("g", [App ("c", [])])])
*)

type term = 
    Var of string
  | App of string * term list


(* -------------   small parser for terms ------------------- *)

let lexer = make_lexer ["("; ")"; ","; "%"; "+"; "-"; "."]

let rec parse_term = parser
    [< fx = parse_ident; (args, is_app) = parse_rest >] ->
      if is_app then App(fx, args) else Var fx
and parse_rest = parser
    [< 'Kwd "("; args = parse_args; 'Kwd ")" >] -> (args, true)
  | [< >] -> ([], false)
and parse_ident = parser
    [<'Ident s>] -> s
and parse_args = parser
    [< a = parse_term ; args = parse_args >] -> a::args
  | [<  >] -> [] 


let read_term s = parse_term (lexer(Stream.of_string s)) 

(* -------------   small printer for terms ------------------- *)

let print_term term =
  let rec print_term_aux term =
    match term with
      Var x -> print_string x
    | App (f, args) -> 
	print_string f; 
	print_string "(";  
	map (fun a -> print_term_aux  a; print_string " ") args ; 
	print_string ")"
  in print_term_aux term ; print_newline ()


(* -------------   substitution ------------------- *)

let rec subst_var var subst =
  match subst with
    [] -> var
  | (hdvar, hdterm) :: tl -> if (hdvar = var) then hdterm else subst_var var tl

let rec substitution term subst =
  match term with
    Var x -> subst_var term subst
  | App (f, args) -> App (f, (map (fun t -> substitution t subst) args))


(* -------------   occurs_check ------------------- *)

let rec occurs_p var term =
  match term with
    Var x -> var = term
  | App (f, args) -> exists (fun t -> occurs_p var t) args

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

(* propagate substitution t1 := t2 in subst *)
let make_new_substitution var term subst = 
  (var,term) :: map (fun (v, s) -> (v, substitution s [(var,term)])) subst

let rec unify t1 t2 subst =
  match t1 with
    Var x1 ->
      (match t2 with
	Var x2 -> 
	  if x1 = x2
	  then (true, subst)
	  else (true, make_new_substitution t1 t2 subst)
      | App (f2, args2) ->
	  if occurs_p t1 t2
	  then (false, [])
	  else (true, make_new_substitution t1 t2 subst))
  | App (f1, args1) ->
      (match t2 with
	Var x2 -> unify t2 t1 subst
      | App (f2, args2) -> if f1 = f2 then unify_args args1 args2 subst else (false, []))
and unify_args args1 args2 subst =
  match args1 with
    [] -> (match args2 with [] -> (true, subst) | hd :: tl -> (false, []))
  | hd1 :: tl1 -> 
      (match args2 with
	[] -> (false, [])
      | hd2 :: tl2 -> 
	  (let (success, new_subst) =
	    unify (substitution hd1 subst) (substitution hd2 subst) subst in
	  if success then unify_args tl1 tl2 new_subst else (false, [])))

let print_unif_result (success, subst) =
  if success 
  then (print_string "Unification succeeded with substitution:" ; print_newline ();
      	map (fun (v, s) -> print_term v; print_string " := "; print_term s; print_newline ()) subst;
	())
  else (print_string "Unification failed." ; print_newline ())

let unification t1 t2 = (print_unif_result (unify (read_term t1) (read_term t2) []))

(* An optimized version of unification. Idea: carry out substitutions only on variables *)
let rec unify_opt t1 t2 subst =
  match t1 with
    Var x1 ->
      (match t2 with
	Var x2 ->
	  if t1 = t2 
	  then (true, subst) 
      (* difference wrt unify *)
	  else 
	    (let t1subst = (substitution t1 subst) and t2subst = (substitution t2 subst) in
	    if (t1 = t1subst) & (t2 = t2subst)
      (* t1 <> t2, t1 = (subst t1)  and t2 = (subst t2) --> both t1 and t2 are variables*)
	    then (true, make_new_substitution t1 t2 subst)
      (* t1 <> t2, t1 <> (subst t1)  or t2 <> (subst t2) --> try unif with substitutions *)
	    else unify_opt t1subst t2subst subst)
      | App (f2, args2) -> 
	  (let t1subst = (substitution t1 subst) in 
	  (if (t1 = t1subst)
	      (* no change when substituting *)
	  then (let t2subst = (substitution t2 subst) in
	  (if (occurs_p t1 t2subst)
	  then (false, [])
	  else (true, make_new_substitution t1 t2subst subst)))
	      (* something changed when substituting *)
	  else (unify_opt t1subst t2 subst))))
  | App (f1, args1) ->
      (match t2 with
   	Var x2 -> unify_opt t2 t1 subst
      | App (f2, args2) -> if f1 = f2 then unify_args_opt args1 args2 subst else (false, []))
and unify_args_opt args1 args2 subst =
  match args1 with
    [] -> (match args2 with [] -> (true, subst) | hd :: tl -> (false, []))
  | hd1 :: tl1 -> 
      (match args2 with
   	[] -> (false, [])
      | hd2 :: tl2 -> 
	   (* difference wrt unify              vv  vv  *)
	  (let (success, new_subst) = unify_opt hd1 hd2 subst in
	  if success then unify_args_opt tl1 tl2 new_subst else (false, [])))

let unification_opt t1 t2 = (print_unif_result (unify_opt (read_term t1) (read_term t2) []))



(* Difference between unification and unification_opt:

# trace substitution ;;

-- Relativ grosser Unterschied bei:
unification_opt "f(x g(x) f(g(x)))" "f(a() g(a()) f(g(a())))";;
unification "f(x g(x) f(g(x)))" "f(a() g(a()) f(g(a())))";;


-- Kein Unterschied bei:
unification_opt "f(x1 x2 x3)" "f(g(a() a()) g(x1 x1) g(x2 x2))" ;;
unification "f(x1 x2 x3)" "f(g(a() a()) g(x1 x1) g(x2 x2))" ;;



*)