compiler-ml/chapter1/redblack.sml

exception LookupError

type key = string

datatype Color = R | B
datatype 'a Tree = E | T of Color * 'a Tree * (key * 'a) * 'a Tree

val empty = E


fun member key E = false
  | member key (T (_, l, (k, _), r)) =
    if key < k then member key l
    else if key > k then member key r
    else true



fun lookup key (T (_, l, (k, v), r)) =
    if key < k then lookup key l
    else if key > k then lookup key r
    else v
  | lookup key E = raise LookupError


fun balance (B,T(R,T(R,lll,llb,llr),lb,lr),b,r) = T(R,
                                                    T(B,lll,llb,llr),
                                                    lb,
                                                    T(B,lr,b,r))
  | balance (B,T(R,ll,lb,T(R,lrl,lrb,lrr)),b,r) = T(R,
                                                    T(B,ll,lb,lrl),
                                                    lrb,
                                                    T(B,lrr,b,r))
  | balance (B,l,b,T(R,T(R,rll,rlb,rlr),rb,rl)) = T(R,
                                                    T(B,l,b,rll),
                                                    rlb,
                                                    T(B,rlr,rb,rl))
  | balance (B,l,b,T(R,rl,rb,T(R,rrl,rrb,rrr))) = T(R,
                                                    T(B,l,b,rl),
                                                    rb,
                                                    T(B,rrl,rrb,rrr))
  | balance body = T body


fun insert (key,newval) t =
  let 
      fun ins E = T (R, E, (key,newval), E)
        | ins (t as T (color, l, (k,v), r)) =
          if key < k then balance (color, ins l, (k,v), r)
          else if key > k then balance (color, l, (k,v), ins r)
          else T(color, l, (key, newval), r)
      val T (_, l, b, r) = ins t
  in
      T (B, l, b, r)
  end