module LLRB where
    
import Test.QuickCheck
import Data.List (nub, sort)

------------------------------------------------------------------------

data Tree a
  = Lf
  | Nd Color a (Tree a) (Tree a)
 deriving (Show)

data Color
  = Red
  | Black
 deriving (Show)

------------------------------------------------------------------------

insert :: Ord a => a -> Tree a -> Tree a
insert x t = makeBlack $ insert' x t
  where
    insert' :: Ord a => a -> Tree a -> Tree a
    insert' x Lf = Nd Red x Lf Lf
    insert' x t@(Nd c y l r)
      = colorFlip $ rotateR $ rotateL $ case compare x y of
        LT -> Nd c y (insert' x l) r
        EQ -> t
        GT -> Nd c y l (insert' x r)
	
    makeBlack, colorFlip, rotateL, rotateR :: Tree a -> Tree a
    makeBlack Lf           = Lf
    makeBlack (Nd _ x l r) = Nd Black x l r
                  
    rotateL (Nd c x l (Nd Red x' l' r'))
      = Nd c x' (Nd Red x l l') r'
    rotateL t = t

    rotateR (Nd c x (Nd Red x' l'@(Nd Red _ _ _) r') r)
      = Nd c x' l' (Nd Red x r' r)
    rotateR t = t

    colorFlip (Nd Black x l@(Nd Red _ _ _) r@(Nd Red _ _ _))
      = Nd Red x (makeBlack l) (makeBlack r)
    colorFlip t = t

------------------------------------------------------------------------

fromList :: Ord a => [a] -> Tree a
fromList = foldr insert Lf
           
toList :: Tree a -> [a]
toList Lf           = []
toList (Nd _ x l r) = toList l ++ [x] ++ toList r

------------------------------------------------------------------------

prop_redBlackTree :: [Int] -> Bool
prop_redBlackTree xs =  binarySearchTree xs
                     && redInvariant         t
                     && blackInvariant       t
                     && leftLeaningInvariant t
  where
    t = fromList xs

    binarySearchTree :: Ord a => [a] -> Bool
    binarySearchTree xs = sort (nub xs) == (toList (fromList xs))

    redInvariant :: Tree a -> Bool
    redInvariant Lf = True
    redInvariant (Nd Black _ l r) =  redInvariant  l && redInvariant r
    redInvariant (Nd Red   _ l r) =  isNotRed      l && isNotRed     r
                                  && redInvariant  l && redInvariant r

    isNotRed :: Tree a -> Bool
    isNotRed Lf               = True
    isNotRed (Nd Black _ _ _) = True
    isNotRed (Nd Red   _ _ _) = False 

    blackInvariant :: Tree a -> Bool
    blackInvariant Lf           = True
    blackInvariant (Nd _ _ l r) =  blackDepth     l == blackDepth     r
                                && blackInvariant l && blackInvariant r
      where
        blackDepth :: Tree a -> Int
        blackDepth Lf               = 0
        blackDepth (Nd Black _ l r) = 1 + blackDepth l
        blackDepth (Nd Red   _ l r) =     blackDepth l

    leftLeaningInvariant :: Tree a -> Bool
    leftLeaningInvariant Lf               = True
    leftLeaningInvariant (Nd Black _ l r) =  isNotRed r
                                          && leftLeaningInvariant l
                                          && leftLeaningInvariant r
    leftLeaningInvariant (Nd Red   _ l r) =  leftLeaningInvariant l
                                          && leftLeaningInvariant r