(* Random trees (c) Copyright 2004 Pierre Corbineau *)
(* published under the Library General Public License *) 

module Make(M:Set.OrderedType) : Set.S with type elt = M.t =
struct

  type elt = M.t (* elements *)

  type level = int

  let prng_state = Random.State.make_self_init ()

  let gen () = 
    let tmp = (Random.State.bits prng_state) in 
      if tmp = 0 then max_int else tmp land (-tmp)

  type t = 
      Node of level * t * elt * t
    | Empty

  let empty = Empty

  let is_empty s = s = Empty 
			 
  let rec mem x = function
      Empty -> false
    | Node (_,left,a,right) ->
	let cmp = M.compare x a in
	  cmp = 0 || mem x (if cmp < 0 then left else right)

  (* splits a set A in subsets < x and > x , 
	present is true if x is in A, false otherwise*) 
	
  let rec split x = function
      Empty -> (Empty,false,Empty) 
    | Node (lvl,left,a,right) ->
	let cmp = M.compare x a in
	  if cmp = 0 then 
	    (left,true,right)
	  else if cmp > 0 then
	    let l,p,r = split x right in 
	      (Node(lvl,left,a,l),p,r)
	  else 
	    let l,p,r = split x left in 
	      (l,p,Node(lvl,r,a,right))
	    
  let rec insert lvlx x set =
    match set with
	Empty -> Node(lvlx,Empty,x,Empty)
      | Node(lvla,left,a,right) ->
	    if lvlx >= lvla then
	      let l,p,r = split x set in
		if p then set else Node(lvlx,l,x,r) 
	    else 
	      let cmp = M.compare x a in
		if cmp = 0 then set
		else if cmp < 0 then
		  Node(lvla,insert lvlx x left,a,right)		    
		else (* cmp > 0 *)
		  Node(lvla,left,a,insert lvlx x right)
  	      
  let add x s = insert (gen ()) x s
		  
  let singleton x = add x Empty
			  
  let rec merge lset rset = 
    (* merges two sets A and B such that for all x in A and y in B, y > x *)  
    match lset,rset with
	Empty,  set  -> set
      |  set , Empty -> set
      | Node (llvl,lleft,lx,lright), Node (rlvl,rleft,rx,rright) ->
	  if llvl < rlvl then
	    Node (rlvl,merge lset rleft,rx,rright)
	  else
	    Node (llvl,lleft,lx,merge lright rset) 
		    
  let rec remove x = function
      Empty -> Empty
    | Node (lvla,left,a,right) ->
	let cmp = M.compare x a in
	  if cmp = 0 then 
	    merge left right
	  else if cmp < 0 then 
	    Node (lvla,remove x left,a,right)
	  else
	    Node (lvla,left,a,remove x right)

  let rec union s1 s2 =
    match s1,s2 with
	Empty, set -> set
      | set, Empty -> set
      | Node (lvl1,left1,x1,right1), Node (lvl2,left2,x2,right2) ->
	  if lvl1 >= lvl2 then
	    let l2,p2,r2 = split x1 s2 in
              if p2 then	
                (merge (union left1 l2) (add x1 (union right1 r2)))
	      else	
	        Node(lvl1,union left1 l2,x1,union right1 r2)
	  else (* lvl1 < lvl2 *)
	    let l1,p1,r1 = split x2 s1 in
	      if p1 then
		(merge (union l1 left2) (add x2 (union r1 right2)))
	      else 
		Node(lvl2,union l1 left2,x2,union r1 right2)

  let rec inter s1 s2 =
    match s1,s2 with
	Empty ,   _   -> Empty
      |   _   , Empty -> Empty
      | Node (lvl1,left1,x1,right1), s2 ->
	    let l2,p2,r2 = split x1 s2 in
	      if p2 then
		Node(lvl1,inter left1 l2,x1,inter right1 r2)
	      else 
		merge (inter left1 l2) (inter right1 r2)

  let rec diff s1 s2 =
    match s1,s2 with
	Empty ,   _   -> Empty
      |   _   , Empty -> Empty
      | Node (lvl1,left1,x1,right1), s2 ->
	    let l2,p2,r2 = split x1 s2 in
	      if p2 then
		merge (diff left1 l2) (diff right1 r2)
	      else 
		Node(lvl1,diff left1 l2,x1,diff right1 r2)

  type zipper = Root | Head of elt * t * zipper

  let rec zip z = function
      Empty -> z
    | Node(lvlx,left,x,right) -> 
	zip (Head(x,right,z)) left 

  let rec compare_zip z1 z2 = 
    match z1,z2 with
	Root,Root -> 0
      | Root,Head(_,_,_) -> -1
      | Head(_,_,_),Root -> 1
      | Head(x1,r1,z1),Head(x2,r2,z2) ->
	  let cmp = M.compare x1 x2 in
	    if cmp = 0 then
	      compare_zip (zip z1 r1) (zip z2 r2)
	    else cmp

  let compare s1 s2 = compare_zip (zip Root s1) (zip Root s2)

  let equal s1 s2 = compare s1 s2 = 0

  let rec subset s1 s2 =
    match s1,s2 with
	Empty ,   _   -> true
      |   _   , Empty -> false
      | Node (lvl1,left1,x1,right1), Node (lvl2,left2,x2,right2) ->
	  let cmp = M.compare x1 x2 in
	    if cmp = 0 then
	      subset left1 left2 && subset right1 right2
	    else if cmp < 0 then
	      subset (Node (lvl1,left1,x1,Empty)) left2 && subset right1 s2
	    else (* cmp > 0 *)
	      subset (Node (lvl1,Empty,x1,right1)) right2 && subset left1 s2
   
  let filter p s = 
    let rec f_aux accu = function
	Empty -> accu
      | Node (lvl,left,x,right) ->
	  f_aux (f_aux (if p x then add x accu else accu) left) right in
      f_aux Empty s

  let partition p s =
    let rec p_aux (s_true, s_false as accu) = function
	Empty -> accu
      | Node (lvl,left,x,right) ->
	  p_aux (p_aux (if p x then add x s_true,s_false 
			else s_true,add x s_false) left) right in
      p_aux (Empty,Empty) s

  let rec iter f = function
      Empty -> ()
    | Node(_,left,a,right) -> iter f left; f a; iter f right
  	  
  let rec fold f set accu = 
    match set with
	Empty -> accu
      | Node(_,left,a,right) ->
	  fold f right (f a (fold f left accu))
	  
  let rec for_all p = function
      Empty -> true
    | Node(_,left,a,right) ->
	for_all p left && p a && for_all p right

  let rec exists p = function
      Empty -> false
    | Node(_,left,a,right) ->
	exists p left || p a || exists p right

  let cardinal s = fold (fun _ n->n+1) s 0

  let elements s = fold (fun x f l-> f (x::l)) s (fun x -> x) []

  let rec max_elt = function
      Empty -> raise Not_found
    | Node(_,_,a,Empty) -> a
    | Node(_,_,_,right) -> max_elt right

  let rec min_elt = function
      Empty -> raise Not_found
    | Node(_,Empty,a,_) -> a
    | Node(_,left,_,_) -> min_elt left 

  let choose = max_elt

end