(**********************************************************************)
(* (c) Greg Morrisett, Neal Glew, Dave Walker, Steve Zdancewic        *)
(*     June 1998, all rights reserved.                                *)
(**********************************************************************)

(* Sets *)

type 'a tree =
   LEAF
 | NODE of ('a * 'a t * 'a t) ref
and 'a t =
    'a tree * int
;;

type 'a set = 'a t * ('a->'a->int);;

type direction = LEFT | RIGHT ;;

let rotate_left node =
  match !node with
   (key,
    (NODE({ contents = left_key,left_left,((_,slr) as left_right) }),_),
    ((_,sr) as right)) -> 
      node := (left_key,left_left,(NODE(ref(key,left_right,right)),slr+sr+1))
  | _ -> invalid_arg "Set.rotate_left"
;;

let rotate_right node =
  match !node with
    (key,
     ((_,sl) as left),
     (NODE({ contents = right_key,((_,srl) as right_left),right_right }),_)) ->
       node :=
	  (right_key,(NODE(ref(key,left,right_left)),1+sl+srl),right_right)
  | _ -> invalid_arg "Set.rotate_right"
;;	       

let rec lift = function
    [] -> ()
  | [(LEFT, parent)] ->
      rotate_left parent
  | [(RIGHT, parent)] ->
      rotate_right parent
  | (LEFT, parent) :: (LEFT, grandparent) :: ancestors ->
      rotate_left grandparent;
      rotate_left grandparent;
      (* parent has moved into grandparent's position *)
      lift ancestors
  | (RIGHT, parent) :: (RIGHT, grandparent) :: ancestors ->
      rotate_right grandparent;
      rotate_right grandparent;
      (* parent has moved into grandparent's position *)
      lift ancestors
  | (LEFT, parent) :: (RIGHT, grandparent) :: ancestors ->
      rotate_left parent;
      rotate_right grandparent;
      lift ancestors
  | (RIGHT, parent) :: (LEFT, grandparent) :: ancestors ->
      rotate_right parent;
      rotate_left grandparent;
      lift ancestors
;;

let splay reln t =
  let rec aux path = function 
      ((NODE ({ contents = (key, left, right) } as node)),_) ->
        let comp = reln key in
        if comp = 0 then
          (lift path;
           true)
        else if comp < 0 then
          (* left *)
          aux ((LEFT, node) :: path) left
        else
          (* right *)
          aux ((RIGHT, node) :: path) right
    | (LEAF,_) ->
        (match path with
          [] -> false
        | _ :: path' ->
            (lift path';
             false)) in
  aux [] t
;;

let empty_t = (LEAF,0)
let empty comp = (empty_t,comp);;
let singleton comp i = ((NODE (ref (i,empty_t,empty_t)),1),comp);;
let cardinality ((_,size),_) = size;;

let member (t,comp) key =
  splay (comp key) t
;;

let insert (t,comp as s) key =
   if splay (comp key) t then
     s
   else
     let t = 
       match t with
	 (NODE {contents = (key', ((_,sl) as left), ((_,sr) as right))}, _) ->
           if comp key key' < 0 then
             (* left *)
             (NODE (ref (key,
			 left,
			 (NODE (ref (key', empty_t, right)), sr+1))),
	      2+sl+sr)
           else
             (NODE (ref (key,
			 (NODE (ref (key', left, empty_t)), sl+1),
			 right)),
	      2+sl+sr)
       | (LEAF,_) -> (NODE (ref (key, empty_t, empty_t)),1) in
     (t,comp)
;;

let delete (t,comp as s) key =
  if splay (comp key) t then
    let t =
      match t with
      	(NODE {contents = (_, (LEAF,_), right)}, _) -> right
      | (NODE {contents = (_, left, (LEAF,_))}, _) -> left
      | (NODE {contents = (_, left, ((_,sr) as right))}, _) ->
          splay (fun _ -> 1) left;
          (match left with
            (NODE {contents = (left_key,
			       ((_,sll) as left_left),
			       (LEAF,_))},_) ->
	      (NODE (ref (left_key, left_left, right)), 1+sll+sr)
          | _ -> failwith "Set.remove")
      | _ -> failwith "Set.remove" in
    (t,comp)
  else
    s
;;

let union (t1,comp as s1) (t2,_ as s2) =
  let rec aux t1 t2 =
    match (t1,t2) with
      ((LEAF,_), _) -> (t2,comp)
    | (_, (LEAF,_)) -> (t1,comp)
    | ((NODE n1, sz1), (NODE n2, sz2)) -> 
      	if (sz1>=sz2) then
          if sz2=1 then
            let (x2,_,_)=(!n2) in insert (t1,comp) x2
          else
            let (x,l,r)=(!n1) in
            if splay (comp x) t2 then
              let (_,l2,r2) = (!n2) in
              let ((_,sll as ll),_) = aux l l2 in
              let ((_,srr as rr),_) = aux r r2 in
              ((NODE (ref (x, ll, rr)), 1+sll+srr),comp)
            else
              let (x2,((_,sl2) as l2),((_,sr2) as r2)) = (!n2) in
              if comp x x2 < 0 then 
              	let ((_,sll as ll),_) = aux l l2 in
              	let ((_,srr as rr),_) =
		  aux r (NODE(ref(x2,empty_t,r2)), succ sr2) in
              	((NODE (ref (x, ll, rr)), 1+sll+srr),comp)
              else
              	let ((_,sll as ll),_) =
		  aux l (NODE(ref(x2,l2,empty_t)), succ sl2) in
              	let ((_,srr as rr),_) = aux r r2 in
              	((NODE (ref (x, ll, rr)), 1+sll+srr),comp)
      	else
          if sz1=1 then
            let (x1,_,_)=(!n1) in insert (t2,comp) x1
          else
            let (x,l,r)=(!n2) in
            if (splay (comp x) t1) then
              let (_,l1,r1) = (!n1) in
              let ((_,sll as ll),_) = aux l l1 in
              let ((_,srr as rr),_) = aux r r1 in
              ((NODE (ref (x, ll, rr)), 1+sll+srr),comp)
            else
              let (x1,((_,sl1) as l1),((_,sr1) as r1)) = (!n1) in
              if comp x x1 < 0 then 
              	let ((_,sll as ll),_) = aux l l1 in
              	let ((_,srr as rr),_) =
		  aux r (NODE(ref(x1,empty_t,r1)), succ sr1) in
              	((NODE (ref (x, ll, rr)), 1+sll+srr),comp)
              else
              	let ((_,sll as ll),_) =
		  aux l (NODE(ref(x1,l1,empty_t)), succ sl1) in
              	let ((_,srr as rr),_) = aux r r1 in
              	((NODE ( ref (x, ll, rr)), 1+sll+srr),comp) in
  aux t1 t2
;;

let rec elements_aux coll t =
  match t with
    (LEAF,_) -> coll
  | (NODE {contents=(x,l,r)},_) ->
      x::(elements_aux (elements_aux coll l) r)
;;

let elements (t,_) = elements_aux [] t;;

let fold f (t,_) b =
  let rec fold_aux b t =
    match t with
      (LEAF, _) -> b
    | (NODE {contents=(x,l,r)}, _) ->
      	let b' = fold_aux b l in
      	f x (fold_aux b' r)
  in fold_aux b t
;;

let intersect ((_, sz1), comp as s1) ((_, sz2), _ as s2) =
  if sz1 > sz2 then
    fold (fun x s -> if member s1 x then insert s x else s) s2 (empty comp)
  else
    fold (fun x s -> if member s2 x then insert s x else s) s1 (empty comp)
;;

let is_empty ((_, sz), _) = sz = 0
;;

let from_list comp =
  List.fold_left (fun s x -> insert s x) (empty comp)
;;

let app f s =
  fold (fun x _ -> begin f x; () end) s ()
;;

exception Not_member
let subset s1 s2 =
    try begin
      fold (fun x _ -> if member s2 x then () else raise Not_member) s1 ();
      true
    end
    with Not_member -> false
;;

let diff s1 s2 =
  fold (fun x s -> delete s x) s2 s1
;;

let equals s1 s2 = (subset s1 s2) & (subset s2 s1)
;;

let choose (t,_) =
  match t with
    (LEAF, _) -> failwith "set.ml: choose called on empty set"
  | (NODE {contents=(x,_,_)}, _) -> x
;;

open Format;;

let print prn_elt f (t,_) =
  let rec aux (t,_) first =
    match t with
      LEAF -> ()
    | NODE {contents=(elt,left,right)} ->
	if not first then fprintf f ",@ ";
	prn_elt f elt; aux left false; aux right false in
  fprintf f "{@[<hov>"; aux t true; fprintf f "@]}"
;;

(* EOF: set.ml *)
