Red-black trees with deletion in Haskell

rb.hs

import System.Random
import Data.Array.IO
import Control.Monad

data Color = R | B deriving (Eq, Show)
data Tree a = E | N Color (Tree a) a (Tree a) deriving (Eq, Show)
data Result a b = D a | T b deriving (Eq, Show)

sseq (D x) f = D x
sseq (T x) f = f x

fromResult (D x) = x
fromResult (T x) = x

f <$$> (D x) = D (f x)
f <$$> (T x) = T (f x)

balance (N B (N R (N R a x b) y c) z d) = T (N R (N B a x b) y (N B c z d))
balance (N B (N R a x (N R b y c)) z d) = T (N R (N B a x b) y (N B c z d))
balance (N B a x (N R (N R b y c) z d)) = T (N R (N B a x b) y (N B c z d))
balance (N B a x (N R b y (N R c z d))) = T (N R (N B a x b) y (N B c z d))
balance (N B a x b) = D (N B a x b)
balance (N R a x b) = T (N R a x b)

blacken (N _ a y b) = N B a y b
blacken s = s

insert x s = (blacken . fromResult . ins) s where
	ins E = T (N R E x E)
	ins (N k a y b)
		| x < y = sseq ((\a -> N k a y b) <$$> ins a) balance
		| x == y = D (N k a y b)
		| x > y = sseq ((\b -> N k a y b) <$$> ins b) balance

balance1 (N k (N R (N R a x b) y c) z d) = D (N k (N B a x b) y (N B c z d))
balance1 (N k (N R a x (N R b y c)) z d) = D (N k (N B a x b) y (N B c z d))
balance1 (N k a x (N R (N R b y c) z d)) = D (N k (N B a x b) y (N B c z d))
balance1 (N k a x (N R b y (N R c z d))) = D (N k (N B a x b) y (N B c z d))
balance1 s = blacken1 s

blacken1 (N R a y b) = D (N B a y b)
blacken1 s = T s

eqL (N k a y (N B c z d)) = balance1 (N k a y (N R c z d))
eqL (N k a y (N R c z d)) = (\a -> N B a z d) <$$> eqL (N R a y c)
eqR (N k (N B a x b) y c) = balance1 (N k (N R a x b) y c)
eqR (N k (N R a x b) y c) = (\b -> N B a x b) <$$> eqR (N R b y c)

delete x s = (fromResult . del) s where
	del E = D E
	del (N k a y b)
		| x < y = sseq ((\a -> N k a y b) <$$> del a) eqL
		| x == y = delCur (N k a y b)
		| x > y = sseq((\b -> N k a y b) <$$> del b) eqR

delCur (N R a y E) = D a
delCur (N B a y E) = blacken1 a
delCur (N k a y b) = sseq ((\b -> N k a min b) <$$> b1) eqR where
	(b1, min) = delMin b

delMin (N R E y b) = (D b, y)
delMin (N B E y b) = (blacken1 b, y)
delMin (N k a y b) = (sseq ((\a -> N k a y b) <$$> a1) eqL, min) where
	(a1, min) = delMin a



shuffle :: [a] -> IO [a]
shuffle xs = do
        ar <- newArray n xs
        forM [1..n] $ \i -> do
            j <- randomRIO (i,n)
            vi <- readArray ar i
            vj <- readArray ar j
            writeArray ar j vi
            return vj
  where
    n = length xs
    newArray :: Int -> [a] -> IO (IOArray Int a)
    newArray n xs =  newListArray (1,n) xs

depth E = 0
depth (N _ a _ b) = 1 + max (depth a) (depth b)

walk E = []
walk (N _ a x b) = concat [walk a, [x], walk b]

test on = do
	let adds = foldl (flip insert) E on
	print $ depth adds
	print $ walk adds
	let rems = foldl (flip delete) adds [20..50]
	print $ depth rems
	print $ walk rems
	
main = do
	let arr = [1..100]
	test arr
	rnd <- shuffle arr
	test rnd

Code from "Faster, Simpler Red-Black Trees" by Cameron Moy.