(******************************************************************************
 ** RedBlackTable.sml
 ** sml
 **
 ** Guy Blelloch 
 ** Implementation of TABLE using Red Black Trees
 **   The red black code is based on Chris Okasaki's code from his thesis
 **   although it had to be changed to deal with the fact that
 **   an update can be either a delete or insert.
 **   This implementation is not efficient for unions and intersections.
 ******************************************************************************)

functor RedBlackTable (structure Index : TOTAL_ORD) 
    : TABLE where type index = Index.t =
struct
    exception NoEntry

    type index = Index.t

    exception notImplemented

    datatype Color = R | B

    datatype 'a Tree = 
	E 
      | T of Color * 'a Tree * (index * 'a) * 'a Tree 
      | BB of 'a Tree

    type 'a table = 'a Tree

    val empty = E

    fun isEmpty(t) = case t of E => true | _ => false

    fun addBlack (T (R,a,x,b)) = T (B,a,x,b)
      | addBlack t = BB t

    fun addColor (R,a) = a
      | addColor (B,a) = addBlack a

    fun fixRoot E = E
      | fixRoot (T (_, a, y, b)) = T(B, a, y, b)
      | fixRoot (BB a) = a

    fun bal (color,BB a,x,T (R,b,y,c)) = (* color=black *)
	T (color,bal (R,BB a,x,b),y,c)
      | bal (color,BB a,x,T (B,b,y,T (R,c,z,d))) =
	T (color,T (B,a,x,b),y,T (B,c,z,d))
      | bal (color,BB a,x,T (B,T (R,b,y,c),z,d)) = (* d is black *)
	T (color,T (B,a,x,b),y,T (B,c,z,d))
      | bal (color,BB a,x,T (B,b,y,c)) = (* b and c are black *)
	addBlack (T (color,a,x,T (R,b,y,c)))
      | bal (color,T (R,a,x,b),y,BB c) =
	T (color,a,x,bal (R,b,y,BB c))
      | bal (color,T (B,T (R,a,x,b),y,c),z,BB d) =
	T (color,T (B,a,x,b),y,T (B,c,z,d))
      | bal (color,T (B,a,x,T (R,b,y,c)),z,BB d) =
	T (color,T (B,a,x,b),y,T (B,c,z,d))
      | bal (color,T (B,a,x,b),y,BB c) =
	addBlack (T (color,T (R,a,x,b),y,c))
      | bal ((B,T (R,T (R,a,x,b),y,c),z,d)
             | (B,T (R,a,x,T (R,b,y,c)),z,d)
             | (B,a,x,T (R,T (R,b,y,c),z,d))
             | (B,a,x,T (R,b,y,T (R,c,z,d))) ) =
        T (R,T (B,a,x,b),y,T (B,c,z,d))
      | bal body = T body

    fun delMin (T (color,E,y,b)) = (y,addColor (color,b))
      | delMin (T (color,a,y,b)) = 
	let 
	    val (x,a') = delMin a
	in 
	    (x, bal (color,a',y,b)) 
	end

    fun joinColor (color,a,E) = addColor (color,a)
      | joinColor (color,a,b) = 
	let 
	    val (x,b') = delMin b
	in 
	    bal (color,a,x,b') 
	end

    fun updateReturn f (x, t) =
        let 
	    fun update E = 
		(case f(NONE) of
		     (NONE, w) => (E, w)
		   | (SOME(v), w) => (T(R, E, (x, v), E),w))
	      | update (s as T (color, a, (y,v), b)) =
		(case Index.compare(x,y) of
		     LESS => (let val (a', w) = update a 
			      in (bal(color, a', (y,v) ,b), w)
			      end)
		   | GREATER => (let val (b', w) = update b 
				 in (bal(color, a, (y,v) ,b'), w)
				 end)
		   | EQUAL => 
			 case f(SOME(v)) of
			     (NONE, w) => (joinColor (color, a, b), w)
			   | (SOME(v'), w) => (T(color, a, (y,v'), b), w))
            val (t, w) = update t
	in 
	    (fixRoot t, w)
	end

    fun foldl f accum E = accum
      | foldl f accum (T(color, l, (k,v), r)) =
	foldl f (f(v, (foldl f accum l))) r

    fun foldlIdx f accum E = accum
      | foldlIdx f accum (T(color, l, kv, r)) =
	foldlIdx f (f(kv, (foldlIdx f accum l))) r

    fun foldr f accum E = accum
      | foldr f accum (T(color, l, (k,v), r)) =
	foldr f (f(v, (foldr f accum r))) l

    fun foldrIdx f accum E = accum
      | foldrIdx f accum (T(color, l, kv, r)) =
	foldrIdx f (f(kv, (foldrIdx f accum r))) l

    fun map f E = E
      | map f (T(p, l, (y,v), r)) =
	T(p, map f l, (y, f(v)), map f r)

    fun mapIdx f E = E
      | mapIdx f (T(p, l, (k,v), r)) =
	T(p, mapIdx f l, (k, f(k,v)), mapIdx f r)

    fun filter f _ = raise notImplemented

    fun filterIdx f _ = raise notImplemented

    fun size E = 0
      | size (T(color, l, kv, r)) = size(l) + size(r) + 1

    fun merge (f,fa,fb) (a,b) = raise notImplemented

    (* THE FOLLOWING ARE DERIVED FROM THE ABOVE *)

    fun union f (a,b) =
	merge ((fn (k,SOME(v1),SOME(v2)) => f(v1,v2)
                 | (_,SOME(v1),_) => SOME(v1)
                 | (_,_,SOME(v2)) => SOME(v2)
                 | _ => NONE),
	       (fn (t) => t),
	       (fn (t) => t))
  	      (a,b)

    fun unionIdx f (a,b) =
	merge ((fn (k,SOME(v1),SOME(v2)) => f(k,v1,v2)
                 | (_,SOME(v1),_) => SOME(v1)
                 | (_,_,SOME(v2)) => SOME(v2)
                 | _ => NONE),
	       (fn (t) => t),
	       (fn (t) => t))
  	      (a,b)

    fun intersect f (a,b) =
	merge ((fn (k,SOME(v1),SOME(v2)) => f(v1,v2)
                 | _ => NONE),
	       (fn (t) => empty),
	       (fn (t) => empty))
  	      (a,b)

    fun intersectIdx f (a,b) =
	merge ((fn (k,SOME(v1),SOME(v2)) => f(k,v1,v2)
                 | _ => NONE),
	       (fn (t) => empty),
	       (fn (t) => empty))
  	      (a,b)

    fun diff f (a,b) =
	merge ((fn (k,SOME(v1), SOME(v2)) => f(v1,v2)
                 | (_,SOME(v), NONE) => SOME(v)
                 | _ => NONE),
	       (fn (t) => t),
	       (fn (t) => empty))
  	      (a,b)

    fun diffIdx f (a,b) =
	merge ((fn (k,SOME(v1), SOME(v2)) => f(k,v1,v2)
                 | (_,SOME(v), NONE) => SOME(v)
                 | _ => NONE),
	       (fn (t) => t),
	       (fn (t) => empty))
  	      (a,b)

    fun update f (k,t) =
	let 
	    val (t',_) = updateReturn (fn v => (f(v),())) (k,t)
	in
	    t'
	end

    fun elt t k = 
	let
	    val (_,ret) = updateReturn (fn NONE => (NONE,NONE) | a => (a,a)) (k,t)
	in
	    ret
	end

    fun find(k,t) = elt t k

    fun insert((k,v),t) = update (fn _ => SOME(v)) (k,t)

    fun delete (k,t) = 
	update (fn _ => NONE) (k,t)

    fun deleteReturn (k,t) = 
	updateReturn (fn NONE => (NONE,NONE) | a => (NONE,a)) (k,t)
 
    fun toList a = foldrIdx (op ::) nil a

    fun fromList a = List.foldr insert empty a

end