
  0 | module RedBlackSet
  1 | 
  2 | import Set
  3 | 
  4 | %default total
  5 | 
  6 | data Color = R | B
  7 | 
  8 | export
  9 | data RedBlackSet a = E | T Color (RedBlackSet a) a (RedBlackSet a)
 10 | 
 11 | balance : Color -> RedBlackSet a -> a -> RedBlackSet a -> RedBlackSet a
 12 | balance B (T R (T R e x b) y c) z d = T R (T B e x b) y (T B c z d)
 13 | balance B (T R e x (T R b y c)) z d = T R (T B e x b) y (T B c z d)
 14 | balance B e x (T R (T R b y c) z d) = T R (T B e x b) y (T B c z d)
 15 | balance B e x (T R b y (T R c z d)) = T R (T B e x b) y (T B c z d)
 16 | balance color e x b = T color e x b
 17 | 
 18 | export
 19 | Ord a => Set RedBlackSet a where
 20 |   empty = E
 21 | 
 22 |   member x E = False
 23 |   member x (T _ l y r) = case compare x y of
 24 |     LT => member x l
 25 |     GT => member x r
 26 |     EQ => True
 27 | 
 28 |   insert x s = assert_total $ let T _ l y r = ins s in T B l y r
 29 |    where
 30 |     -- ins always returns a T (never an E)
 31 |     ins : RedBlackSet a -> RedBlackSet a
 32 |     ins E = T R E x E
 33 |     ins s'@(T color l y r) = case compare x y of
 34 |       LT => balance color (ins l) y r
 35 |       GT => balance color l y (ins r)
 36 |       EQ => s'
 37 |