module RedBlack where

----------------------------------------------------------------------------------- 
-- From: http://www.mew.org/~kazu/proj/red-black-tree/ ----------------------------
----------------------------------------------------------------------------------- 

data RBTree a = Leaf | Fork Color (RBTree a) a (RBTree a)
data Color    = Red  | Black

{-@
inline measure cheight :: Color -> Int
cheight (x) = if (x == Black) then 1 else 0
@-}

{-@ 
measure height        :: RBTree a -> Int   
height (Leaf)         = 0
height (Fork c l x r) = (cheight c) + (height l)
@-}

{-@ 
measure color        :: RBTree a -> Color 
color (Leaf)         = Black
color (Fork c l x r) = c
@-}

{-@ predicate isBlack X     = (X == Black)                                                 @-}
{-@ predicate isRed X       = (X == Red)                                                   @-}
{-@ predicate okRoot1 C L R = ((isRed C)  => ((isBlack (color L)) || (isBlack (color R)))) @-}
{-@ predicate okRoot2 C L R = ((isRed C)  => ((isBlack (color L)) && (isBlack (color R)))) @-}
{-@ predicate eqHeight  L R = ((height L) == (height R))                                   @-}

{-@ 
measure invKids2        :: RBTree a -> Prop 
invKids2 (Leaf)         =  true
invKids2 (Fork c l x r) =  (OkRoot2 c l r) && (invKids2 l) && (invKids2 r) 
@-}

{-@ 
measure invKids1        :: RBTree a -> Prop 
invKids1 (Leaf)         = true
invKids1 (Fork c l x r) = (OkRoot1 c l r) && (invKids2 l) && (invKids2 r)
@-}

{-@ 
measure invHeight        :: RBTree a -> Prop 
invHeight (Leaf)         =  true
invHeight (Fork c l x r) =  (eqHeight l r) && (invHeight l) && (invHeight r)
@-}

{-@ type RB1Tree a = {v: (RBTree a) | (invKids1 v) && (invHeight v)} @-}
{-@ type RB2Tree a = {v: (RBTree a) | (invKids2 v) && (invHeight v)} @-}

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

{-@ empty :: IRBTree a @-}
empty     = Leaf

{-@ insert       :: Ord a => a -> RB2Tree a -> RB2Tree a @-}
insert x t       = Fork Black d e f
  where
    Fork _ d e f = ins x t

ins x Leaf             = Fork R Leaf x Leaf
ins x t@(Fork c l y r) = case compare x y of
                           LT -> balanceL c (ins x l) y r
                           GT -> balanceR c l y (ins x r)
                           EQ -> t

{-@ balanceL :: c:Color 
             -> l:(RB1Tree a) 
             -> x:a 
             -> r:{v : (RB2Tree a) | (eqHeight l v)} 
             -> {v:(RB2Tree a) | (height v) = (cheight c) + (height l) } @-}
balanceL :: Color -> RBTree a -> a -> RBTree a -> RBTree a
balanceL B (Fork R (Fork R a x b) y c) z d = Fork R (Fork B a x b) y (Fork B c z d)
balanceL B (Fork R a x (Fork R b y c)) z d = Fork R (Fork B a x b) y (Fork B c z d)
balanceL k a x b                           = Fork k a x b

{-@ balanceR :: c:Color 
             -> l:(RB2Tree a) 
             -> x:a 
             -> r:{v : (RB1Tree a) | (eqHeight l v)} 
             -> {v:(RB2Tree a) | (height v) = (cheight c) + (height l) } @-}
balanceR :: Color -> RBTree a -> a -> RBTree a -> RBTree a
balanceR B a x (Fork R b y (Fork R c z d)) = Fork R (Fork B a x b) y (Fork B c z d)
balanceR B a x (Fork R (Fork R b y c) z d) = Fork R (Fork B a x b) y (Fork B c z d)
balanceR k a x b                           = Fork k a x b