(* $Id: oset.ml,v 1.17 2002/08/04 10:32:54 garrigue Exp $ *)

open StdLabels

module Set =
  (* This code is copied from Objective Caml standard library *)
  struct

    type 'a t = Empty | Node of 'a t * 'a * 'a t * int

    (* Sets are represented by balanced binary trees (the heights of the
       children differ by at most 2 *)

    let height = function
        Empty -> 0
      | Node(_, _, _, h) -> h

    (* Creates a new node with left son l, value x and right son r.
       l and r must be balanced and | height l - height r | <= 2.
       Inline expansion of height for better speed. *)

    let create l x r =
      let hl = match l with Empty -> 0 | Node(_,_,_,h) -> h in
      let hr = match r with Empty -> 0 | Node(_,_,_,h) -> h in
      Node(l, x, r, (if hl >= hr then hl + 1 else hr + 1))

    (* Same as create, but performs one step of rebalancing if necessary.
       Assumes l and r balanced.
       Inline expansion of create for better speed in the most frequent case
       where no rebalancing is required. *)

    let bal l x r =
      let hl = match l with Empty -> 0 | Node(_,_,_,h) -> h in
      let hr = match r with Empty -> 0 | Node(_,_,_,h) -> h in
      if hl > hr + 2 then begin
        match l with
          Empty -> invalid_arg "Set.bal"
        | Node(ll, lv, lr, _) ->
            if height ll >= height lr then
              create ll lv (create lr x r)
            else begin
              match lr with
                Empty -> invalid_arg "Set.bal"
              | Node(lrl, lrv, lrr, _)->
                  create (create ll lv lrl) lrv (create lrr x r)
            end
      end else if hr > hl + 2 then begin
        match r with
          Empty -> invalid_arg "Set.bal"
        | Node(rl, rv, rr, _) ->
            if height rr >= height rl then
              create (create l x rl) rv rr
            else begin
              match rl with
                Empty -> invalid_arg "Set.bal"
              | Node(rll, rlv, rlr, _) ->
                  create (create l x rll) rlv (create rlr rv rr)
            end
      end else
        Node(l, x, r, (if hl >= hr then hl + 1 else hr + 1))

    (* Same as bal, but repeat rebalancing until the final result
       is balanced. *)

    let rec join l x r =
      match bal l x r with
        Empty -> invalid_arg "Set.join"
      | Node(l', x', r', _) as t' ->
          let d = height l' - height r' in
          if d < -2 || d > 2 then join l' x' r' else t'

    (* Merge two trees l and r into one.
       All elements of l must precede the elements of r.
       Assumes | height l - height r | <= 2. *)

    let rec merge t1 t2 =
      match (t1, t2) with
        (Empty, t) -> t
      | (t, Empty) -> t
      | (Node(l1, v1, r1, h1), Node(l2, v2, r2, h2)) ->
          bal l1 v1 (bal (merge r1 l2) v2 r2)

    (* Same as merge, but does not assume anything about l and r. *)

    let rec concat t1 t2 =
      match (t1, t2) with
        (Empty, t) -> t
      | (t, Empty) -> t
      | (Node(l1, v1, r1, h1), Node(l2, v2, r2, h2)) ->
          join l1 v1 (join (concat r1 l2) v2 r2)

    (* Splitting *)

    let rec split ~cmp x = function
        Empty ->
          (Empty, None, Empty)
      | Node(l, v, r, _) ->
          let c = compare x v in
          if c = 0 then (l, Some v, r)
          else if c < 0 then
            let (ll, vl, rl) = split ~cmp x l in (ll, vl, join rl v r)
          else
            let (lr, vr, rr) = split ~cmp x r in (join l v lr, vr, rr)

    (* Implementation of the set operations *)

    let empty = Empty

    let is_empty = function Empty -> true | _ -> false

    let rec mem ~cmp x = function
        Empty -> false
      | Node(l, v, r, _) ->
          let c = cmp x v in
          c = 0 || mem ~cmp x (if c < 0 then l else r)

    let rec add ~cmp x = function
        Empty -> Node(Empty, x, Empty, 1)
      | Node(l, v, r, _) as t ->
          let c = cmp x v in
          if c = 0 then t else
          if c < 0 then bal (add ~cmp x l) v r else bal l v (add ~cmp x r)

    let singleton x = Node(Empty, x, Empty, 1)

    let rec remove ~cmp x = function
        Empty -> Empty
      | Node(l, v, r, _) ->
          let c = cmp x v in
          if c = 0 then merge l r else
          if c < 0 then bal (remove ~cmp x l) v r else bal l v (remove ~cmp x r)

    let rec union ~cmp s1 s2 =
      match (s1, s2) with
        (Empty, t2) -> t2
      | (t1, Empty) -> t1
      | (Node(l1, v1, r1, h1), Node(l2, v2, r2, h2)) ->
          if h1 >= h2 then
            if h2 = 1 then add ~cmp v2 s1 else begin
              let (l2, _, r2) = split ~cmp v1 s2 in
              join (union ~cmp l1 l2) v1 (union ~cmp r1 r2)
            end
          else
            if h1 = 1 then add ~cmp v1 s2 else begin
              let (l1, _, r1) = split ~cmp v2 s1 in
              join (union ~cmp l1 l2) v2 (union ~cmp r1 r2)
            end

    let rec inter ~cmp s1 s2 =
      match (s1, s2) with
        (Empty, t2) -> Empty
      | (t1, Empty) -> Empty
      | (Node(l1, v1, r1, _), t2) ->
          match split ~cmp v1 t2 with
            (l2, None, r2) ->
              concat (inter ~cmp l1 l2) (inter ~cmp r1 r2)
          | (l2, Some _, r2) ->
              join (inter ~cmp l1 l2) v1 (inter ~cmp r1 r2)

    let rec diff ~cmp s1 s2 =
      match (s1, s2) with
        (Empty, t2) -> Empty
      | (t1, Empty) -> t1
      | (Node(l1, v1, r1, _), t2) ->
          match split ~cmp v1 t2 with
            (l2, None, r2) ->
              join (diff ~cmp l1 l2) v1 (diff ~cmp r1 r2)
          | (l2, Some _, r2) ->
              concat (diff ~cmp l1 l2) (diff ~cmp r1 r2)

    let rec compare_aux ~cmp l1 l2 =
        match (l1, l2) with
        ([], []) -> 0
      | ([], _)  -> -1
      | (_, []) -> 1
      | (Empty :: t1, Empty :: t2) ->
          compare_aux ~cmp t1 t2
      | (Node(Empty, v1, r1, _) :: t1, Node(Empty, v2, r2, _) :: t2) ->
          let c = cmp v1 v2 in
          if c <> 0 then c else compare_aux ~cmp (r1::t1) (r2::t2)
      | (Node(l1, v1, r1, _) :: t1, t2) ->
          compare_aux ~cmp (l1 :: Node(Empty, v1, r1, 0) :: t1) t2
      | (t1, Node(l2, v2, r2, _) :: t2) ->
          compare_aux ~cmp t1 (l2 :: Node(Empty, v2, r2, 0) :: t2)

    let compare ~cmp s1 s2 =
      compare_aux ~cmp [s1] [s2]

    let equal ~cmp s1 s2 =
      compare ~cmp s1 s2 = 0

    let rec subset ~cmp s1 s2 =
      match (s1, s2) with
        Empty, _ ->
          true
      | _, Empty ->
          false
      | Node (l1, v1, r1, _), (Node (l2, v2, r2, _) as t2) ->
          let c = cmp v1 v2 in
          if c = 0 then
            subset ~cmp l1 l2 && subset ~cmp r1 r2
          else if c < 0 then
            subset ~cmp (Node (l1, v1, Empty, 0)) l2 && subset ~cmp r1 t2
          else
            subset ~cmp (Node (Empty, v1, r1, 0)) r2 && subset ~cmp l1 t2

    let rec iter f = function
        Empty -> ()
      | Node(l, v, r, _) -> iter f l; f v; iter f r

    let rec fold f s accu =
      match s with
        Empty -> accu
      | Node(l, v, r, _) -> fold f l (f v (fold f r accu))

    let rec for_all p = function
        Empty -> true
      | Node(l, v, r, _) -> p v && for_all p l && for_all p r

    let rec exists p = function
        Empty -> false
      | Node(l, v, r, _) -> p v || exists p l || exists p r

    let filter ~cmp p s =
      let rec filt accu = function
        | Empty -> accu
        | Node(l, v, r, _) ->
            filt (filt (if p v then add ~cmp v accu else accu) l) r in
      filt Empty s

    let partition ~cmp p s =
      let rec part (t, f as accu) = function
        | Empty -> accu
        | Node(l, v, r, _) ->
            part (part (if p v then (add ~cmp v t, f)
                        else (t, add ~cmp v f)) l) r in
      part (Empty, Empty) s

    let rec cardinal = function
        Empty -> 0
      | Node(l, v, r, _) -> cardinal l + 1 + cardinal r

    let rec elements_aux accu = function
        Empty -> accu
      | Node(l, v, r, _) -> elements_aux (v :: elements_aux accu r) l

    let elements s =
      elements_aux [] s

    let rec min_elt = function
        Empty -> raise Not_found
      | Node(Empty, v, r, _) -> v
      | Node(l, v, r, _) -> min_elt l

    let rec max_elt = function
        Empty -> raise Not_found
      | Node(l, v, Empty, _) -> v
      | Node(l, v, r, _) -> max_elt r

    let choose = min_elt

  end

type 'a t = 'a Set.t
open Set

class ['a] c ?compare:(cmp=Pervasives.compare) l = object (_ : 'b)
  val mutable set =
    List.fold_left l ~init:Empty ~f:(fun acc x -> add ~cmp x acc)
  method contents : 'a t = set
  method set (s : 'b) = set <- s#contents
  method clear = set <- empty
  method is_empty = is_empty set
  method mem x = mem ~cmp x set
  method add x = set <- add ~cmp x set
  method remove x = set <- remove ~cmp x set
  method union (s : 'b) = {< set = union ~cmp set s#contents >}
  method inter (s : 'b) = {< set = inter ~cmp set s#contents >}
  method diff (s : 'b) = {< set = diff ~cmp set s#contents >}
  method compare (s : 'b) = compare ~cmp set s#contents
  method equal (s : 'b) = equal ~cmp set s#contents
  method subset (s : 'b) = subset ~cmp set s#contents
  method iter ~(f : 'a -> unit) = iter f set
  method fold : 'c. f:('a -> 'c -> 'c) -> 'c -> 'c = fun ~f -> fold f set
  method for_all ~f = for_all f set
  method exists ~f = exists f set
  method filter ~f = set <- filter ~cmp f set
  method partition ~f =
    let s1, s2 = partition ~cmp f set in
    {< set = s1 >}, {< set = s2 >}
  method cardinal = cardinal set
  method elements = elements set
  method min_elt = min_elt set
  method max_elt = max_elt set
  method choose = choose set
end

class ['a] f ?compare:(cmp=Pervasives.compare) l = object (_ : 'b)
  val set =
    List.fold_left l ~init:Empty ~f:(fun acc x -> add ~cmp x acc)
  method contents : 'a t = set
  method is_empty = is_empty set
  method mem x = mem ~cmp x set
  method add x = {< set = add ~cmp x set >}
  method remove x = {< set = remove ~cmp x set >}
  method union (s : 'b) = {< set = union ~cmp set s#contents >}
  method inter (s : 'b) = {< set = inter ~cmp set s#contents >}
  method diff (s : 'b) = {< set = diff ~cmp set s#contents >}
  method compare (s : 'b) = compare ~cmp set s#contents
  method equal (s : 'b) = equal ~cmp set s#contents
  method subset (s : 'b) = subset ~cmp set s#contents
  method iter ~(f : 'a -> unit) = iter f set
  method fold : 'c. f:('a -> 'c -> 'c) -> 'c -> 'c = fun ~f -> fold f set
  method for_all ~f = for_all f set
  method exists ~f = exists f set
  method filter ~f = {< set = filter ~cmp f set >}
  method partition ~f =
    let s1, s2 = partition ~cmp f set in
    {< set = s1 >}, {< set = s2 >}
  method cardinal = cardinal set
  method elements = elements set
  method min_elt = min_elt set
  method max_elt = max_elt set
  method choose = choose set
end