(*
 * splay_set.ml -- Implementation of splay_set module
 *
 * Ling Li, Caltech
 * Jan. 16, 2001
 *)

module type EltSig =
sig
   type elt
   val compare : elt -> elt -> int
end

module type SetSig =
sig
   type elt
   type t
   val create : unit -> t
   val mem : elt -> t -> bool
   val insert : elt -> t -> t
   val delete : elt -> t -> t
end

module MakeSplaySet (Elt : EltSig) =
struct
   (* Same type of elements *)
   type elt = Elt.elt

   (* Splay trees *)
   type tree =
    | Leaf
    | Node of tree * elt * tree

   (* Have to use a reference cell to save result of splay *)
   type t = tree ref

   (* To 'splay' x to the root (level 1), we will first use splay2 to
    * lift x to the root or to level 2, by only using zig(zag)-zig(zag)
    * rotations. For the purpose of knowing the exact position of x
    * after splay2, we need type splayResult.
    * Root: x is the root;
    * Left, Right: x is the root of the left or right subtree (level 2). *)
   type splayResult = Left | Right | Root

   (* splay2: Lift x by using only zig(zag)-zig(zag) rotations *)
   let rec splay2 x = function
    | Leaf -> Root, Leaf
    | Node (l1, x1, r1) as node ->  (* level 1 node *)
         let i = Elt.compare x x1 in
            if i = 0 then
               Root, node
            else if i < 0 then
               match splay2 x l1 with
                | _, Leaf ->
                     Root, node
                | Left, Node (Node (l3, x3, r3), x2, r2) ->  (* Zig-Zig *)
                     Root, Node (l3, x3, Node (r3, x2, Node (r2, x1, r1)))
                | Right, Node (l2, x2, Node (l3, x3, r3)) -> (* Zag-Zig *)
                     Root, Node (Node (l2, x2, l3), x3, Node (r3, x1, r1))
                | _, l1' -> (* In fact, it is Root, l1' *)
                     Left, Node (l1', x1, r1)
            else
               match splay2 x r1 with
                | _, Leaf ->
                     Root, node
                | Left, Node (Node (l3, x3, r3), x2, r2) ->  (* Zig-Zag *)
                     Root, Node (Node (l1, x1, l3), x3, Node (r3, x2, r2))
                | Right, Node (l2, x2, Node (l3, x3, r3)) -> (* Zag-Zag *)
                     Root, Node (Node (Node (l1, x1, l2), x2, l3), x3, r3)
                | _, r1' -> (* In fact, it is Root, r1' *)
                     Right, Node (l1, x1, r1')

   (* splay: Splay operation *)
   let splay x t =
      match splay2 x t with
       | Left, Node (Node (l2, x2, r2), x1, r1) ->  (* Zig *)
            Node (l2, x2, Node (r2, x1, r1))
       | Right, Node (l1, x1, Node (l2, x2, r2)) -> (* Zag *)
            Node (Node (l1, x1, l2), x2, r2)
       | _, t' -> (* In fact, Root, t' *)
            t'

   (* join: Join two splay trees by [Min(S'); left(S') <- S].
    * Assuming that all elements in s are smaller than x,
    *           and all elements in s' are larger than x. *)
   let join x s s' =
      if s = Leaf then s'
      else if s' = Leaf then s
      else
         match splay x s' with   (* find minimum of s' *)
          | Node (_, x, r) -> Node (s, x, r)
          | _ -> raise (Failure "Splay_set.join")

   (* create: Make an empty tree *)
   let create () = ref Leaf

   (* mem: Membership *)
   let mem x s =
      s := splay x !s;
      match !s with
       | Leaf -> false
       | Node (_, y, _) -> (Elt.compare x y) = 0

   (* insert: Insert a node to the splay tree *)
   let insert x s =
      s := splay x !s;
      match !s with
       | Leaf -> s := Node (Leaf, x, Leaf); s
       | Node (l, y, r) ->
            let i = Elt.compare x y in
               if i = 0 then                 (* x = y *)
                  s
               else if i < 0 then            (* x < y *)
                  (s := Node (l, x, Node (Leaf, y, r)); s)
               else                          (* x > y *)
                  (s := Node (Node (l, y, Leaf), x, r); s)

   (* delete: Delete a node from the splay tree, if it exists *)
   let delete x s =
      s := splay x !s;
      match !s with
       | Leaf -> s;
       | Node (l, y, r) ->
            if (Elt.compare x y) = 0 then
               (s := join x l r; s)
            else s
end