let (<) : int -> int -> bool = (<)
let (>) : int -> int -> bool = (>)
let (=) : int -> int -> bool = (=)


module type E =
sig
  (* module V : sig type t end *)
  type elem
  include Custom.T

  val empty : t
  val full  : t
  val cup   : t -> t -> t
  val cap   : t -> t -> t
  val diff  : t -> t -> t
  val atom  : elem -> t

end

module type S =
sig
  type elem
  include Custom.T

  val get: t -> (elem list * elem list) list
  val get': t -> (elem list * (elem list) list) list

  val empty : t
  val full  : t
  val cup   : t -> t -> t
  val cap   : t -> t -> t
  val diff  : t -> t -> t
  val atom  : elem -> t

  val iter: (elem-> unit) -> t -> unit

  val compute: empty:'b -> full:'b -> cup:('b -> 'b -> 'b) 
    -> cap:('b -> 'b -> 'b) -> diff:('b -> 'b -> 'b) ->
    atom:(elem -> 'b) -> t -> 'b

  val is_empty : t -> bool

  val splitvars : t -> t * t

  val print: string -> t -> (Format.formatter -> unit) list

  val trivially_disjoint: t -> t -> bool
end

(* ternary BDD
 * where the nodes are Atm of X.t | Var of String.t
 * Variables are always before Values
 * All the leaves are then base types 
 *
 * we add a third case when two leaves of the bdd are of the same
 * kind, that's it Val of t1 , Val of t2
 *
 * This representation can be used for all kinds.
 * Intervals, Atoms and Chars can be always merged (for union and intersection)
 * Products can be merged for intersections
 * Arrows can be never merged
 *
 * extract_var : copy the orginal tree and on one copy put to zero all 
 * leaves that have an Atm on the other all leaves that have a Var
 *
 * *)

module type MAKE = functor (T : E) -> S with type elem = T.t Custom.pairvar

module Make(T : E) =
struct
  (* ternary decision trees . cf section 11.3.3 Frish PhD *)
  (* plus variables *)
  (* Custom.Atm are containers (Atoms, Chars, Intervals, Pairs ... )
   * Custom.Var are String
   *)
  module X = Custom.Var(T)
  type elem = T.t Custom.pairvar
  type t =
    | True
    | False
    | Split of int * elem * t * t * t

  let rec equal a b =
    (a == b) ||
    match (a,b) with
      | Split (h1,x1,p1,i1,n1), Split (h2,x2,p2,i2,n2) ->
	  (h1 == h2) &&
	  (equal p1 p2) && (equal i1 i2) &&
	  (equal n1 n2) && (X.equal x1 x2)
      | _ -> false

(* Idea: add a mutable "unique" identifier and set it to
   the minimum of the two when egality ... *)

  let rec compare a b =
    if (a == b) then 0 
    else match (a,b) with
      | Split (h1,x1, p1,i1,n1), Split (h2,x2, p2,i2,n2) ->
	  if h1 < h2 then -1 else if h1 > h2 then 1 
	  else let c = X.compare x1 x2 in if c <> 0 then c
	  else let c = compare p1 p2 in if c <> 0 then c
	  else let c = compare i1 i2 in if c <> 0 then c 
	  else compare n1 n2
      | True,_  -> -1
      | _, True -> 1
      | False,_ -> -1
      | _,False -> 1

  let rec hash = function
    | True -> 1
    | False -> 0
    | Split(h, _,_,_,_) -> h

  let compute_hash x p i n = 
	(X.hash x) + 17 * (hash p) + 257 * (hash i) + 16637 * (hash n)

  let rec check = function
    | True | False -> ()
    | Split (h,x,p,i,n) ->
	assert (h = compute_hash x p i n);
	(match p with Split (_,y,_,_,_) -> assert (X.compare x y < 0) | _ -> ());
	(match i with Split (_,y,_,_,_) -> assert (X.compare x y < 0) | _ -> ());
	(match n with Split (_,y,_,_,_) -> assert (X.compare x y < 0) | _ -> ());
	X.check x; check p; check i; check n

  let atom x =
    let h = X.hash x + 17 in (* partial evaluation of compute_hash... *)
    Split (h, x,True,False,False)
 
  let neg_atom x =
    let h = X.hash x + 16637 in (* partial evaluation of compute_hash... *)
    Split (h, x,False,False,True)

  let rec iter f = function
    | Split (_, x, p,i,n) -> f x; iter f p; iter f i; iter f n
    | _ -> ()

  let rec dump ppf = function
    | True -> Format.fprintf ppf "+"
    | False -> Format.fprintf ppf "-"
    | Split (_,x, p,i,n) -> 
	Format.fprintf ppf "%i(@[%a,%a,%a@])" 
	(* X.dump x *) (X.hash x) dump p dump i dump n

  let rec print f ppf = function
    | True -> Format.fprintf ppf "Any"
    | False -> Format.fprintf ppf "Empty"
    | Split (_, x, p,i, n) ->
	let flag = ref false in
	let b () = if !flag then Format.fprintf ppf " | " else flag := true in
	(match p with 
	   | True -> b(); Format.fprintf ppf "%a" f x
	   | False -> ()
	   | _ -> b (); Format.fprintf ppf "%a & @[(%a)@]" f x (print f) p );
	(match i with 
	   | True -> assert false;
	   | False -> ()
	   | _ -> b(); print f ppf i);
	(match n with 
	   | True -> b (); Format.fprintf ppf "@[~%a@]" f x
	   | False -> ()
	   | _ -> b (); Format.fprintf ppf "@[~%a@] & @[(%a)@]" f x (print f) n)
	
  let print a = function
    | True -> [ fun ppf -> Format.fprintf ppf "%s" a ]
    | False -> []
    | c -> [ fun ppf -> print X.dump ppf c ]

  (* XXX : since every path contains 1 Atm, I should be able to
   * descend on the first path and get a sample from the leaf *)    
  let rec sample = function
    | Split (_,Custom.Var _, p,i,n) ->
        begin match sample p with
        |Some x -> Some x
        |None ->
            begin match sample i with
            |Some x -> Some x
            |None ->
                begin match sample n with
                |Some x -> Some x
                |None -> None
                end
            end
        end
    | Split (_,Custom.Atm x, _,_,_) -> Some x
    | _ -> None

  let rec contains y x =
    match x,y with
    |True,_ |False,_ -> false
    |Split (_,Custom.Var a, p,i,n),Custom.Var b ->
        (a == b) || (contains y p) || (contains y i) || (contains y n)
    |Split (_,Custom.Atm a, p,i,n),Custom.Atm b ->
        ((T.cap a b) == T.empty) || (contains y p) || (contains y i) || (contains y n)
    |Split (_,_, p,i,n),_ ->
        (contains y p) || (contains y i) || (contains y n)

  let rec get accu pos neg = function
    | True -> (pos,neg) :: accu
    | False -> accu
    | Split (_,x, p,i,n) ->
	(*OPT: can avoid creating this list cell when pos or neg =False *)
	let accu = get accu (x::pos) neg p in
	let accu = get accu pos (x::neg) n in
	let accu = get accu pos neg i in
	accu
	  
  let get x = get [] [] [] x

  let rec get' accu pos neg = function
    | True -> (pos,neg) :: accu
    | False -> accu
    | Split (_,x,p,i,n) ->
	let accu = get' accu (x::pos) neg p in
	let rec aux l = function
	  | Split (_,x,False,i,n') when equal n n' ->
	      aux (x :: l) i
	  | i ->
	      let accu = get' accu pos (l :: neg) n in
	      get' accu pos neg i
	in
	aux [x] i

  let get' x = get' [] [] [] x

  let compute ~empty ~full ~cup ~cap ~diff ~atom b =
    let rec aux = function
      | True -> full
      | False -> empty
      | Split(_,x, p,i,n) ->
	  let p = cap (atom x) (aux p)
	  and i = aux i
	  and n = diff (aux n) (atom x) in
	  cup (cup p i) n
    in
    aux b
      
(* Invariant: correct hash value *)

  let split0 x pos ign neg =
    Split (compute_hash x pos ign neg, x, pos, ign, neg)

  let empty = False
  let full = True

  let is_empty t = (t == empty)

(* Invariants:
     Split (x, pos,ign,neg) ==>  (ign <> True), (pos <> neg)
*)

  let rec has_true = function
    | [] -> false
    | True :: _ -> true
    | _ :: l -> has_true l

  let rec has_same a = function
    | [] -> false
    | b :: l -> (equal a b) || (has_same a l)

  let rec split x p i n =
    if X.equal x (Custom.Atm T.empty) then False
    else if i == True then True 
    else if equal p n then p ++ i
    else let p = simplify p [i] and n = simplify n [i] in
    if equal p n then p ++ i
    else split0 x p i n

  and simplify a l =
    match a with
      | False -> False
      | True -> if has_true l then False else True
      | Split (_,Custom.Atm x, True,False,False) -> split (Custom.Atm(T.diff T.full x)) True False False
      | Split (_,Custom.Atm x, False,False,True) -> split (Custom.Atm(T.diff T.full x)) True False False
      | Split (_,x,p,i,n) ->
	  if (has_true l) || (has_same a l) then False
	  else s_aux2 a x p i n [] [] [] l
  and s_aux2 a x p i n ap ai an = function
    | [] -> 
	let p = simplify p ap 
	and n = simplify n an
	and i = simplify i ai in
	if equal p n then p ++ i else split0 x p i n
    | b :: l -> s_aux3 a x p i n ap ai an l b 
  and s_aux3 a x p i n ap ai an l = function
    | False -> s_aux2 a x p i n ap ai an l
    | True -> assert false
    | Split (_,x2,p2,i2,n2) as b ->
	if equal a b then False 
	else let c = X.compare x2 x in
	if c < 0 then s_aux3 a x p i n ap ai an l i2
	else if c > 0 then s_aux2 a x p i n (b :: ap) (b :: ai) (b :: an) l
	else s_aux2 a x p i n (p2 :: i2 :: ap) (i2 :: ai) (n2 :: i2 :: an) l

  (* Inv : all leafs are of type Val are always merged *)
  (* union *)
  and ( ++ ) a b = if a == b then a
  else match (a,b) with
    | True, _ | _, True -> True
    | False, a | a, False -> a
    
    | Split (_,Custom.Atm x1, True,False,False), Split (_,Custom.Atm x2, True,False,False) ->
        split (Custom.Atm (T.cup x1 x2)) True False False

    | Split (_,Custom.Atm x1, False,False,True), Split (_,Custom.Atm x2, False,False,True) ->
        split (Custom.Atm (T.cup (T.diff T.full x1) (T.diff T.full x2))) True False False

    | Split (_,Custom.Atm x1, True,False,False), Split (_,Custom.Atm x2, False,False,True) ->
        split (Custom.Atm (T.cup x1 (T.diff T.full x2))) True False False

    | Split (_,Custom.Atm x1, False,False,True), Split (_,Custom.Atm x2, True,False,False) ->
        split (Custom.Atm (T.cup (T.diff T.full x1) x2)) True False False

    | Split (_,x1, p1,i1,n1), Split (_,x2, p2,i2,n2) ->
	let c = X.compare x1 x2 in
	if c = 0 then split x1 (p1 ++ p2) (i1 ++ i2) (n1 ++ n2)
	else if c < 0 then split x1 p1 (i1 ++ b) n1
	else split x2 p2 (i2 ++ a) n2

(* seems better not to make ++ and this split mutually recursive;
   is the invariant still inforced ? *)

  (* intersection *)
  let rec ( ** ) a b = if a == b then a else match (a,b) with
    | True, a | a, True -> a
    | False, _ | _, False -> False

    | Split (_,Custom.Atm x1, True,False,False), Split (_,Custom.Atm x2, True,False,False) ->
        split (Custom.Atm(T.cap x1 x2)) True False False

    | Split (_,Custom.Atm x1, False,False,True), Split (_,Custom.Atm x2, False,False,True) ->
        split (Custom.Atm(T.cap (T.diff T.full x1) (T.diff T.full x2))) True False False

    | Split (_,Custom.Atm x1, True,False,False), Split (_,Custom.Atm x2, False,False,True) ->
        split (Custom.Atm(T.cap x1 (T.diff T.full x2))) True False False

    | Split (_,Custom.Atm x1, False,False,True), Split (_,Custom.Atm x2, True,False,False) ->
        split (Custom.Atm(T.cap (T.diff T.full x1) x2)) True False False

    | Split (_,x1, p1,i1,n1), Split (_,x2, p2,i2,n2) ->
	let c = X.compare x1 x2 in
	if c = 0 then
	  split x1 
	    (p1 ** (p2 ++ i2) ++ (p2 ** i1))
	    (i1 ** i2)
	    (n1 ** (n2 ++ i2) ++ (n2 ** i1))  
	else if c < 0 then split x1 (p1 ** b) (i1 ** b) (n1 ** b)
	else split x2 (p2 ** a) (i2 ** a) (n2 ** a)

  let rec trivially_disjoint a b =
    if a == b then a == False
    else match (a,b) with
      | True, a | a, True -> a == False
      | False, _ | _, False -> true
      | Split (_,x1, p1,i1,n1), Split (_,x2, p2,i2,n2) ->
	  let c = X.compare x1 x2 in
	  if c = 0 then
(* try expanding -> p1 p2; p1 i2; i1 p2; i1 i2 ... *)
	    trivially_disjoint (p1 ++ i1) (p2 ++ i2) &&
	    trivially_disjoint (n1 ++ i1) (n2 ++ i2)
	  else if c < 0 then
	    trivially_disjoint p1 b &&
	    trivially_disjoint i1 b &&
	    trivially_disjoint n1 b
	  else
	    trivially_disjoint p2 a &&
	    trivially_disjoint i2 a &&
	    trivially_disjoint n2 a

  let rec neg = function
    | True -> False
    | False -> True
    | Split (_,Custom.Atm x, True,False,False) -> split (Custom.Atm(T.diff T.full x)) True False False
    | Split (_,Custom.Atm x, False,False,True) -> split (Custom.Atm(T.diff T.full x)) True False False
    | Split (_,x, p,i,False) -> split x False (neg (i ++ p)) (neg i)
    | Split (_,x, False,i,n) -> split x (neg i) (neg (i ++ n)) False 
    | Split (_,x, p,False,n) -> split x (neg p) (neg (p ++ n)) (neg n)  
    | Split (_,x, p,i,n) -> split x (neg (i ++ p)) False (neg (i ++ n))
	      
  let rec ( // ) a b =
    let negatm = T.diff T.full in
    if a == b then False 
    else match (a,b) with
      | False,_ | _, True -> False
      | a, False -> a
      | True, b -> neg b

      | Split (_,Custom.Atm x1, True,False,False), Split (_,Custom.Atm x2, True,False,False) ->
          split (Custom.Atm(T.diff x1 x2)) True False False

      | Split (_,Custom.Atm x1, False,False,True), Split (_,Custom.Atm x2, False,False,True) ->
          split (Custom.Atm(T.diff (negatm x1) (negatm x2))) True False False

      | Split (_,Custom.Atm x1, True,False,False), Split (_,Custom.Atm x2, False,False,True) ->
          split (Custom.Atm(T.diff x1 (negatm x2))) True False False

      | Split (_,Custom.Atm x1, False,False,True), Split (_,Custom.Atm x2, True,False,False) ->
          split (Custom.Atm(T.diff (negatm x1) x2)) True False False

      | Split (_,x1, p1,i1,n1), Split (_,x2, p2,i2,n2) ->
	  let c = X.compare x1 x2 in
	  if c = 0 then
	    if (i2 == False) && (n2 == False) 
	    then split x1 (p1 // p2) (i1 // p2) (n1 ++ i1)
	    else 
	      split x1 ((p1++i1) // (p2 ++ i2)) False ((n1++i1) // (n2 ++ i2))
	  else if c < 0 then
	    split x1 (p1 // b) (i1 // b) (n1 // b) 
	  else
	    split x2 (a // (i2 ++ p2)) False (a // (i2 ++ n2))
	      
  let cup = ( ++ )
  let cap = ( ** )
  let diff = ( // )

  (* return a couple of trees (v,a), the second where all variables
   * v = only variables as leaves
   * a = only atoms as leaves
   *)
  let rec splitvars = function
    | True -> True,True
    | False -> False,False
    | Split (_,Custom.Var _, True,False,False) as x -> x, False
    | Split (_,Custom.Atm _, True,False,False) as x -> True, x
    | Split (_,x, p,i,n) ->
        let l,p' = splitvars p in
        let c,i' = splitvars i in
        let r,n' = splitvars n in 
        match split x l c r, split x p' i' n' with
        |t,Split (_,Custom.Var _, True,False,False) -> t,False
        |(_,_) as t -> t

  let splitvars t = splitvars t
  
end