(* persistent implementation using a map giving the class id 
 * the updating of the map is implemented with a customized remap function
 * which preserves the AVL structure
 *)

module M = struct

  type t =
      Empty
    | Node of t * int * int ref * t * int
	
  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 "Map.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 "Map.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 "Map.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 "Map.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 empty = Empty

  let rec add x data = function
      Empty ->
        Node(Empty, x, data, Empty, 1)
    | Node(l, v, d, r, h) ->
        let c = compare x v in
        if c = 0 then
          Node(l, x, data, r, h)
        else if c < 0 then
          bal (add x data l) v d r
        else
          bal l v d (add x data r)
	    
  let rec find_outside node x m = 
    let rec find0_outside node x = function
      | Node(l, v, d, r, _) as n when n != node ->
          let c = compare x v in
          if c = 0 then begin
	    let vd = !d in
	    assert (x <> vd);
	    try 
	      let rd = 
		find0_outside node vd (if compare vd x < 0 then l else r) 
	      in
	      d := rd;
	      rd
	    with Not_found -> 
	      find_outside n vd m
          end else 
	    find0_outside node x (if c < 0 then l else r)
      | Node (_,_,_,_,h) ->
	  (*Format.printf "[%d]@?" h; *)raise Not_found
      | Empty ->
          raise Not_found
    in
    try find0_outside node x m with Not_found -> x
  and find x m =
    let rec find0 x = function
	Empty ->
          raise Not_found
      | Node(l, v, d, r, _) as n ->
          let c = compare x v in
          if c = 0 then begin
	    let vd = !d in
	    assert (x <> vd);
	      try 
		let rd = find0 vd (if compare vd x < 0 then l else r) in
		d := rd;
		rd
	      with Not_found -> 
		find_outside n vd m
          end else 
	    find0 x (if c < 0 then l else r)
    in
    try find0 x m with Not_found -> x

(***
  let rec find m x =
    let rec find0 x = function
      | Empty ->
          raise Not_found
      | Node(l, v, d, r, _) ->
          let c = compare x v in
          if c = 0 then begin
	    let vd = !d in
	    assert (x <> vd);
	    try 
	      let rd = find0 vd (if compare vd x < 0 then l else r) in
	      d := rd;
	      rd
	    with Not_found -> 
	      find m vd
          end else 
	    find0 x (if c < 0 then l else r)
    in
    try find0 x m with Not_found -> x
***)

end

type t = M.t

let create n = M.empty

let find m x = M.find x m (*try M.find m x with Not_found -> x*)
    
let union m i j =
  let ri = find m i in
  let rj = find m j in
  if ri <> rj then M.add ri (ref rj) m else m