(******************************************************************************
 ** SplayTable.sml
 ** sml
 **
 ** Guy Blelloch 
 **
 ** Implementation of TABLE using Splay Trees.
 ** Uses the SML/NJ Lib SplayTree code.
 ******************************************************************************)

functor SplayTable (structure Index : TOTAL_ORD) 
    : TABLE where type index = Index.t =
struct

    exception NoEntry

    structure ST = SplayTree

    type index = Index.t

    type 'a table = (index*'a) ST.splay

    val empty = ST.SplayNil
	
    fun isEmpty(t) = case t of ST.SplayNil => true | _ => false

    fun split(k,tree) =
	case ST.splay((fn (k',_) => Index.compare(k',k)), tree) of
	    (_, ST.SplayNil) => (ST.SplayNil, NONE, ST.SplayNil)
	  | (EQUAL, ST.SplayObj{value=(k,v), left=l, right=r}) => (l,SOME(v),r)
	  | (LESS, ST.SplayObj{value=(k,v), left=l, right=r}) => 
		(ST.SplayObj{value=(k,v), left=l, right=ST.SplayNil}, NONE, r)
	  | (GREATER, ST.SplayObj{value=(k,v), left=l, right=r}) => 
		(l, NONE, ST.SplayObj{value=(k,v), left=ST.SplayNil, right=r})

    fun updateReturn f (k, t) =
	let
	    val (l,m,r) = split(k, t)
	in
	    case f(m) of
		(NONE, v)    => (ST.join(l,r),v)
	      | (SOME(d), v) => (ST.SplayObj {value=(k,d),left=l,right=r}, v)
	end

    fun merge (f,fa,fb) (a,b) =
	let
	    fun merge'(ST.SplayNil,b) = fb(b)
	      | merge'(a,ST.SplayNil) = fa(a)
	      | merge'(ST.SplayObj({value=(k1,v1), left=l1, right=r1}), t2) =
		let val (l2,m,r2) = split(k1,t2)
		    val ln = merge'(l1,l2)
		    val rn = merge'(r1,r2)
		in 
		    case f(k1,SOME(v1),m) of
			NONE => ST.join(ln,rn)
		      | SOME(v) => ST.SplayObj {value=(k1,v), left=ln, right=rn}
		end
	in
	    merge'(a,b)
	end

    fun foldl f accum ST.SplayNil = accum
      | foldl f accum (ST.SplayObj {value=(k,v),left=l,right=r}) =
	foldl f (f(v, (foldl f accum l))) r

    fun foldlIdx f accum ST.SplayNil = accum
      | foldlIdx f accum (ST.SplayObj {value=kv,left=l,right=r}) =
	foldlIdx f (f(kv, (foldlIdx f accum l))) r

    fun foldr f accum ST.SplayNil = accum
      | foldr f accum (ST.SplayObj {value=(k,v),left=l,right=r}) =
	foldr f (f(v, (foldr f accum r))) l

    fun foldrIdx f accum ST.SplayNil = accum
      | foldrIdx f accum (ST.SplayObj {value=kv,left=l,right=r}) =
	foldrIdx f (f(kv, (foldrIdx f accum r))) l

    fun map f ST.SplayNil = empty
      | map f (ST.SplayObj{value=(k,v), left=l, right=r}) =
	ST.SplayObj{value = (k, f(v)), 
		    left = (map f l), 
		    right = (map f r)}

    fun mapIdx f ST.SplayNil = empty
      | mapIdx f (ST.SplayObj{value=(k,v), left=l, right=r}) =
	ST.SplayObj{value = (k, f(k,v)), 
		    left = (mapIdx f l), 
		    right = (mapIdx f r)}

    fun filter f ST.SplayNil = empty
      | filter f (ST.SplayObj{value=(k,v), left=l, right=r}) =
	case f(v) of
	    NONE => ST.join(filter f l, filter f r)
	  | SOME(v') => ST.SplayObj{value = (k, v'), 
				    left = (filter f l), 
				    right = (filter f r)}

    fun filterIdx f ST.SplayNil = empty
      | filterIdx f (ST.SplayObj{value=(k,v), left=l, right=r}) =
	case f(k,v) of
	    NONE => ST.join(filterIdx f l, filterIdx f r)
	  | SOME(v') => ST.SplayObj{value = (k, v'), 
				    left = (filterIdx f l), 
				    right = (filterIdx f r)}

    fun size ST.SplayNil = 0
      | size (ST.SplayObj{value=_, left=l, right=r}) = size(l) + size(r) + 1

    (* 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