type atm = int
type cst = int
type var = int

(*********************************************)
(* Helpers on Arrays *)


let mkArr n = Array.init n (fun x -> x)

let permuteArr f arr =
 let arr2 = Array.copy arr in
 (for i = 0 to (Array.length arr2) - 1
   do arr2.(f i) <- arr.(i) done ;
  arr2
 )

let array'diff arr l =
  let arr2 = Array.copy arr in
  (List.iter (fun (i,e) -> arr2.(i) <- e) l ;
   arr2
  )

let maxAtoms = 30
let allAtoms = mkArr maxAtoms


(**************************************************)
(* Perms *)

type perm = Perm of int array
let unPerm (Perm p) = p

let idPerm = Perm allAtoms
let image   (Perm p) a = p.(a)
let inverse       p    = Perm (permuteArr (image p) allAtoms)

let compose (Perm p1) (Perm p2) =
     Perm (Array.map (fun x -> p1.(p2.(x))) allAtoms)

let swap a b = function
 | i when i = a -> b
 | i when i = b -> a
 | i            -> i

let swapping a b = Perm (Array.map (swap a b) allAtoms)
let swapPermL a b p = compose (swapping a b) p
                                              

(***************************************************************************)
(* FrsSet *)

type frsSet = FrsSet of bool array
let unFrsSet (FrsSet fs) = fs

let emptyFrsSet = FrsSet (Array.map (fun _ -> false) allAtoms)
let isInFrsSet (FrsSet fs) a = fs.(a)

(* permute p fs = p(fs) *)
let permute p (FrsSet fs) = FrsSet (permuteArr (image p) fs)

let combineFrsSet op (FrsSet fs1) (FrsSet fs2) =
  FrsSet (Array.map (fun i -> op fs1.(i) fs2.(i)) allAtoms)
let unionFrsSet = combineFrsSet (or)


let frsSetSet b l (FrsSet fs) =
     FrsSet (array'diff fs (List.map (fun i -> (i,b)) l))

let frsSetAdd = frsSetSet true
let frsSetDel = frsSetSet false

let dsPerm p1 p2 = FrsSet (Array.map
        (fun a -> (image p1 a) <> (image p2 a)) allAtoms)


(******************************************************************)
(* DAG *)


type node = { (* Data not changed by the algorithm *)
                      head        : head                            ;
              mutable fathers     : node list                       ;

              (* Data changed by the algorithm *)
              mutable deleted     : bool                            ;
              mutable edges       : ((perm * node) * del_edge) list ;
              mutable pointer     : (perm * node) option            ;

              (* Nominal data (changed also by the alogorithm) *)
              mutable fresh_cond  : frsSet
            }

and head = Atm  of atm
         | Cst  of cst
         | Var  of var
         | Abs  of atm  * node
         | Pair of node * node
         | Perm of perm * node

and del_edge = { mutable edge_deleted : bool }
                     

(***************************************************************)
(* Node management *)

let is_same_node (n1 : node) (n2 : node) = n1 == n2

let list_node_iter f = List.iter (function n -> if n.deleted then ()
                                                else f n
                                 )

let addFrsCond r frs  = r.fresh_cond <- unionFrsSet r.fresh_cond frs
let addFrsAtms r b    = r.fresh_cond <- frsSetAdd b r.fresh_cond

(*********************************************************************)
(* Edge management *)

let rec reduceHead r pi s = match (r.head , s.head) with
 | (Perm (pi2,r1), _ ) -> reduceHead r1 (compose (inverse pi2) pi) s
 | (_ , Perm (pi2,s2)) -> reduceHead r (compose pi pi2) s2
 | (_ , _ )            -> (r,pi,s)


let add_undirected_edge r pi s =
 let (r1, pi1, s1) = reduceHead r pi s in
 let pi1inv    = inverse pi1
 and del_egde  = { edge_deleted = false }
 in  (r1.edges <- ((pi1   ,s1) , del_egde) :: r1.edges ;
      s1.edges <- ((pi1inv,r1) , del_egde) :: s1.edges
     )


let list_edges_iter f =
   List.iter (function ((pi,n),d) -> if n.deleted or d.edge_deleted
                                     then ()
                                     else f ((pi,n),d)
             )

let delete_edge (_,_,d) = d.edge_deleted <- true



(************************************************************)
(* DAG *)



type dag = { nonvar     : node list ;
             variables  : node list ;
             suspended  : node list
           }


let reset_node n =
 ( n.deleted     <- false     ;
   n.edges       <- []        ;
   n.pointer     <- None
 )

let reset_dag d =
 ( List.iter reset_node d.nonvar    ;
   List.iter reset_node d.variables ;
   List.iter reset_node d.suspended
 )


(*************************************************)
(* Substitution *)

let maxVars = 30
let sigma = Array.create maxVars None




(************************************************)
(* Unification Algorithm *)



let push stack n = (stack := n :: !stack)
let pop  stack   = match !stack with
 | []        -> failwith "Error: empty stack"
 | (n :: st) -> (stack := st ; n)


let check_clash r pi s =
 let (r2,pi2,s2) = reduceHead r pi s in
 match (r2.head,s2.head) with
 | (Var _      , _          )
 | (_          , Var _      )
 | (Abs (_,_)  , Abs (_,_)  )
 | (Pair (_,_) , Pair (_,_) )                                     -> ()
 | (Atm a      , Atm b      ) when a  = image pi b
                               && not (isInFrsSet r.fresh_cond a)
                               && not (isInFrsSet s.fresh_cond b) -> ()
 | (Cst c1     , Cst c2     ) when c1 = c2                        -> ()
 | (_          , _          )                  -> failwith "FAIL-CLASH"


let rec propagate_unification_edges r pi s =
 let (r2,pi2,s2) = reduceHead r pi s in
 match (r2.head,s2.head) with
 | (Var _        , _            )
 | (_            , Var _        ) -> failwith "propagation failed: var"
 | (Atm _        , Atm _        )
 | (Cst _        , Cst _        ) -> ()
 | (Pair (r1,r2) , Pair (s1,s2) ) -> (add_undirected_edge r1 pi s1 ;
                                      add_undirected_edge r2 pi s2 ;
                                     )
 | (Abs (a,r1)   , Abs  (b1,s1) ) -> let b2  = image pi b1                in
                                     let pi2 = compose (swapping a b2) pi in
                                     (add_undirected_edge r1 pi2 s1 ;
                                      addFrsAtms r [b2]
                                     )
 | (_            , _            ) -> assert false



let rec propagate_freshness_perm fc r = match r.head with
 | Perm (p,r1) -> propagate_freshness_perm (permute (inverse p) fc) r1
 | _           -> addFrsCond r fc
 


let propagate_freshness r =
  let rec aux fc r = match r.head with
    | Cst _
    | Atm _ -> ()
    | Var x -> print_string "r.fresh_cond # x"
    | Pair (r1,r2) -> (propagate_freshness_perm fc r1;
                       propagate_freshness_perm fc r2
                      )
    | Abs (a,r1)   -> propagate_freshness_perm (frsSetDel [a] fc) r1
    | Perm (p,r1)  -> aux (permute (inverse p) fc) r1
  in aux r.fresh_cond r


let rec do_fathers s =
    list_node_iter (fun t -> match t.head with
                              | Perm (_,_) -> do_fathers t
                               | _          -> finish t
                   )
                   s.fathers

and finish r =
 let stack = ref [] in
 (push stack (idPerm,r) ;
  while !stack <> [] do
     let (pi,s) = pop stack in
     ((match s.pointer with
        | Some (pi2,r2) when is_same_node r2 r -> ()
        | Some _    -> failwith "FAIL-LOOP"
        | None      -> s.pointer <- Some ((inverse pi),r)
      ) ;
      do_fathers s ;
      check_clash r pi s;
      (match s.head with
        | Var x        -> sigma.(x) <- Some ((inverse pi),r)
        | _            -> propagate_unification_edges r pi s
      ) ;
      list_edges_iter (function (e,d) ->
                         push stack e ;
                         d.edge_deleted <- true
                      ) s.edges;
      if is_same_node r s
       then addFrsCond r (dsPerm idPerm pi)
       else (addFrsCond r (permute (inverse pi) s.fresh_cond) ;
             s.deleted <- true
            )
    ) ;
    propagate_freshness r ;
    r.deleted <- true
  done
 )




(***************************************************)
(* Main function *)


(* a dag where u = v *)
let unifiy dag s t =
 (reset_dag dag ;
  add_undirected_edge s idPerm t ;
  list_node_iter finish dag.nonvar ;
  list_node_iter finish dag.variables
 )