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

open StdLabels

(* module Map = struct *)

type ('a,'b) t =
    Empty
  | Node of ('a,'b) t * 'a * 'b * ('a,'b) t * int

let empty = Empty

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

let create l x d r =
  let hl = height l and hr = height r in
  Node(l, x, d, r, (if hl >= hr then hl + 1 else hr + 1))

let bal l x d 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, ld, lr, _) ->
        if height ll >= height lr then
          create ll lv ld (create lr x d r)
        else begin
          match lr with
            Empty -> invalid_arg "Set.bal"
          | Node(lrl, lrv, lrd, lrr, _)->
              create (create ll lv ld lrl) lrv lrd (create lrr x d r)
        end
  end else if hr > hl + 2 then begin
    match r with
      Empty -> invalid_arg "Set.bal"
    | Node(rl, rv, rd, rr, _) ->
        if height rr >= height rl then
          create (create l x d rl) rv rd rr
        else begin
          match rl with
            Empty -> invalid_arg "Set.bal"
          | Node(rll, rlv, rld, rlr, _) ->
              create (create l x d rll) rlv rld (create rlr rv rd rr)
        end
  end else
    Node(l, x, d, r, (if hl >= hr then hl + 1 else hr + 1))

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

let rec find ~cmp x = function
    Empty ->
      raise Not_found
  | Node(l, v, d, r, _) ->
      let c = compare x v in
      if c = 0 then d
      else find ~cmp x (if c < 0 then l else r)

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

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

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

let rec iter ~f = function
    Empty -> ()
  | Node(l, v, d, r, _) ->
      iter ~f l; (f ~key:v ~data:d : unit); iter ~f r

let rec fold ~f m init =
  match m with
    Empty -> init
  | Node(l, v, d, r, _) ->
      fold ~f l (f ~key:v ~data:d (fold ~f r init))

(* end *)

class ['a,'b] c ?compare:(cmp=Pervasives.compare) l = object
  val mutable map =
    List.fold_left l ~init:Empty
      ~f:(fun acc (x,y : 'a * 'b) -> add ~cmp x y acc)
  method clear = map <- Empty
  method is_empty = map = Empty
  method add ~key ~data = map <- add ~cmp key data map
  method find key = find ~cmp key map
  method mem key = mem ~cmp key map
  method remove key = map <- remove ~cmp key map
  method iter = iter map
  method fold : 'c. f:(key:'a -> data:'b -> 'c -> 'c) -> 'c -> 'c = fold map
end

class ['a,'b] f ?compare:(cmp=Pervasives.compare) l = object
  val map =
    List.fold_left l ~init:Empty
      ~f:(fun acc (x,y : 'a * 'b) -> add ~cmp x y acc)
  method is_empty = map = Empty
  method add ~key ~data = {< map = add ~cmp key data map >}
  method find key = find ~cmp key map
  method mem key = mem ~cmp key map
  method remove key = {< map = remove ~cmp key map >}
  method iter = iter map
  method fold : 'c. f:(key:'a -> data:'b -> 'c -> 'c) -> 'c -> 'c = fold map
end