-- 
-- Module : RBTree
-- Author : Wu Xingbo
-- Copyright (c) 2010, 2011 Wu Xingbo (wuxb45@gmail.com)
-- New BSD License

{-# LANGUAGE BangPatterns #-}
-- |
--  Pure Haskell Red-Black tree implementation
--
module Data.Tree.RBTree (
  -- * Tree Types
  Color (Red, Black), RBTree (Node, Leaf), emptyRB,
  -- * Interval Types
  Interval (Interval), RealOrd (PInfinity, NInfinity, RealValue),
  -- * Insertion
  (<</), insert, insertOrd, insertOrdList,
  -- * Delete
  (<<\), delete, deleteOrd, deleteOrdList,
  -- * Search
  (<<?), search, searchOrd, searchFast, searchMax, searchMin,
  searchInterval, searchIntervalOrd,
  -- * Verification
  vD, vR
)
where

import Control.Monad(liftM2)

-- |Color of a 'Node'.
--  Leaf is assumed to be Black.
data Color = 
    Red
  | Black 
  deriving (Eq)

-- |Basic RBTree Structure.
data RBTree a = Node Color a !(RBTree a) !(RBTree a) -- ^A Node that holds an element and has two leaves.
              | Leaf  -- ^A Black leaf.

-- |Represents the direction of one step.
data Direction = ToLeft | ToRight deriving (Eq)

-- |Records the one step from a parent node to one of its children nodes.
data Step a = Step Color a Direction !(RBTree a) deriving (Show)

-- |A Path is a series of Steps.
type Path a = [Step a]

-- |RBTree in a 'Zip' mode.
--
--  Current Node can start from any node inside the tree, with a Path back to Root node.
--  RBZip is equivalent to RBTree in Logic.
--  All RBZip can be convert to a RBTree by Trace back to Root point.
data RBZip a = 
  RBZip !(RBTree a) !(Path a)  -- ^ RBZip sub-tree path
  deriving (Show)

-- |used for range query.
data Interval a = Interval (RealOrd a, RealOrd a)

-- |Interval value from -INF to +INF.
data RealOrd a = 
    PInfinity  -- ^positive infinity
  | NInfinity  -- ^positive infinity
  | RealValue a  -- ^Normal value, not need to be Ord.

-- |Simply show tree in (), hard to read but easy to parse.
instance Show a => Show (RBTree a) where
    show (Node c v l r) = "(" ++ show l ++ show v ++ show c ++ show r ++ ")"
    show Leaf = "."

-- |for distinguish Red/Black, show \'*\' for Red and nothing for Black.
instance Show Color where
    show Red = "*"
    show Black = ""

instance Show Direction where
    show ToLeft = "L"
    show ToRight = "R"

instance Show a => Show (RealOrd a) where
    show PInfinity = "+INF"
    show NInfinity = "-INF"
    show (RealValue a) = show a

instance Show a => Show (Interval a) where
    show (Interval (l, r)) = "[" ++ show l ++ ", " ++ show r ++ "]"

-- |Gen an empty Tree.
emptyRB :: RBTree a
emptyRB = Leaf

-- |Get the root node Color of current sub-tree, Leaf is also Black.
getColor :: RBTree a -> Color
getColor (Node c _ _ _) = c
getColor Leaf = Black

-- |Set current Root to Black.
setBlack :: RBTree a -> RBTree a
setBlack (Node _ v l r) = Node Black v l r
setBlack Leaf = Leaf

-- |Set current Root to Red.
setRed :: RBTree a -> RBTree a
setRed (Node _ v l r) = Node Red v l r
setRed Leaf = Leaf -- never happen

-- |Conversion : RBTree \<==> RBZip.
toZip :: RBTree a -> RBZip a
toZip t = RBZip t []

-- |convert a zip to tree.
toTree :: RBZip a -> RBTree a
toTree z = tree
    where (RBZip tree _) = topMostZip z

-- |Zip up.
topMostZip :: RBZip a -> RBZip a
topMostZip (RBZip s ((Step c v d s1):path)) = case d of 
        ToLeft -> topMostZip (RBZip (Node c v s s1) path)
        ToRight -> topMostZip (RBZip (Node c v s1 s) path)
topMostZip z = z

-- |Get the Left-most non-leaf node from a Zip, or get Leaf if it is a Leaf.
leftMostZip :: RBZip a -> RBZip a
leftMostZip this@(RBZip (Node _ _ Leaf _) _) = this
leftMostZip (RBZip (Node c v l r) path) = leftMostZip (RBZip l ((Step c v ToLeft r):path))
leftMostZip z = z --only when leaf itself from start over

-- |Get the Right-most non-leaf node from a Zip, or get Leaf if it is a Leaf.
rightMostZip :: RBZip a -> RBZip a
rightMostZip this@(RBZip (Node _ _ _ Leaf) _) = this
rightMostZip (RBZip (Node c v l r) path) = rightMostZip (RBZip r ((Step c v ToRight l):path))
rightMostZip z = z --leaf itself

-- |Zip up until the sub-tree has a left-parent, used to find the biggest lower-order element of the current node.
leftParentZip :: RBZip a -> RBZip a
leftParentZip (RBZip l ((Step c v ToLeft r):path)) = leftParentZip (RBZip (Node c v l r) path)
leftParentZip (RBZip r ((Step c v ToRight l):path)) = RBZip (Node c v l r) path
leftParentZip (RBZip _ []) = RBZip Leaf [] -- no such parent, return a empty zip

-- |Zip up until the sub-tree has a right-parent, used to find the smallest upper-order element of the current node.
rightParentZip :: RBZip a -> RBZip a
rightParentZip (RBZip r ((Step c v ToRight l):path)) = rightParentZip (RBZip (Node c v l r) path)
rightParentZip (RBZip l ((Step c v ToLeft r):path)) = RBZip (Node c v l r) path
rightParentZip (RBZip _ []) = RBZip Leaf [] -- no such parent, return a empty zip

-- |find predecessor of a node/leaf.
predZip :: RBZip a -> RBZip a
predZip (RBZip (Node c v l@(Node _ _ _ _) r) path) = rightMostZip (RBZip l ((Step c v ToLeft r):path))
predZip z@(RBZip Leaf _) = case lp of
  RBZip Leaf [] -> z -- itself
  _ -> lp
  where lp = leftParentZip z
predZip z@(RBZip (Node c v l r) path) = case lp of
  RBZip Leaf [] -> RBZip l ((Step c v ToLeft r):path)
  _ -> lp
  where lp = leftParentZip z

-- |find predecessor of a node/leaf.
succZip :: RBZip a -> RBZip a
succZip (RBZip (Node c v l r@(Node _ _ _ _)) path) = leftMostZip (RBZip r ((Step c v ToRight l):path))
succZip z@(RBZip Leaf _) = case lp of
  RBZip Leaf [] -> z -- itself
  _ -> lp
  where lp = rightParentZip z
succZip z@(RBZip (Node c v l r) path) = case lp of
  RBZip Leaf [] -> RBZip r ((Step c v ToRight l):path)
  _ -> lp
  where lp = rightParentZip z

-- |Get the Leftmost non-leaf node's value from a Zip.
-- [@param 1@] current node's value.
-- [@param 2@] current node's left child.
leftmostV :: a -> RBTree a -> a
leftmostV v Leaf = v
leftmostV _ (Node _ vl l _) = leftmostV vl l

-- Insertion functions. x will be in left of y if x equals to y and y has already in the tree.

-- |Insert \'Ord\' things.
insertOrd :: (Ord a) => RBTree a -> a -> RBTree a
insertOrd = insert compare

-- |Insert a bunch of \'Ord\' things.
insertOrdList :: (Ord a) => RBTree a -> [a] -> RBTree a
insertOrdList = foldl insertOrd

-- |Insert anything.
-- |you have to provide a compare function.
insert :: (a -> a -> Ordering) -> RBTree a -> a ->RBTree a
insert f t v = setBlack . toTree . insertFixup . (insertRedZip f (toZip t)) $ v

-- |Insert Operator for insertOrd
(<</) :: (Ord a) => RBTree a -> a -> RBTree a
t <</ e = insertOrd t e

insertRedZip :: (a -> a -> Ordering) -> RBZip a -> a -> RBZip a
insertRedZip _ (RBZip Leaf path) v = RBZip (Node Red v Leaf Leaf) path
insertRedZip f (RBZip (Node c v l r) path) new
    | f new v == GT = insertRedZip f (RBZip r ((Step c v ToRight l):path)) new
    | otherwise     = insertRedZip f (RBZip l ((Step c v ToLeft r):path)) new

-- insertFixup
--
-- a : current node
-- b : parent of a
-- c : parent of b
-- d : brother of b
-- vx : value of x
-- dx : direction of x
-- sx : sub-tree of x in the path
-- sxy : sub-tree of x in y side
insertFixup :: RBZip a -> RBZip a
insertFixup (RBZip a@(Node Red _ _ _) ((Step Red vb db sb):(Step Black vc dc d@(Node Red _ _ _)):path)) =
    insertFixup (RBZip newC path)
    where newC = Node Red vc newCL newCR
          (newCL,newCR) = case dc of
              ToLeft -> (newB,newD)
              ToRight -> (newD,newB)
          newB = Node Black vb newBL newBR
          (newBL,newBR) = case db of
              ToLeft -> (a,sb)
              ToRight -> (sb,a)
          !newD = setBlack d
insertFixup (RBZip a@(Node Red va sal sar) ((Step Red vb db sb):(Step Black vc dc d):path)) =
    RBZip newZ (newP:path)
    where (newZ, newP) = case (dc,db) of 
              (ToLeft,ToLeft) -> (a,Step Black vb dc (Node Red vc sb d))
              (ToLeft,ToRight) -> (Node Red vb sb sal, Step Black va dc (Node Red vc sar d))
              (ToRight,ToLeft) -> (Node Red vb sar sb, Step Black va dc (Node Red vc d sal))
              (ToRight,ToRight) -> (a,Step Black vb dc (Node Red vc d sb))
insertFixup t = t

-- Search functions. return \'Just result\' on success, otherwise Nothing.

-- |Search for \'Ord\' things. see 'search'
searchOrd :: (Ord a) => RBTree a -> a -> Maybe a
searchOrd = search compare

-- |search for any thing, you should provide proper compare function.
search :: (b -> a -> Ordering) -> RBTree a -> b -> Maybe a
search f t v = case rZip of
    Just (RBZip (Node _ v' _ _) _) -> Just v'
    _ -> Nothing
    where rZip = searchZip f (toZip t) v

-- |Search operator for searchOrd
(<<?) :: (Ord a) => RBTree a -> a -> Maybe a
t <<? e = searchOrd t e

-- |a faster 'search' function implemetation. strongly recommanded.
searchFast :: (b -> a -> Ordering) -> RBTree a -> b -> Maybe a
searchFast f (Node _ v l r) vs = case f vs v of
    LT -> searchFast f l vs
    GT -> searchFast f r vs
    EQ -> Just v
searchFast _ Leaf _ = Nothing

-- |Search the Maximum value in the tree, equals to get the right-most element.
searchMax :: (Ord a) => RBTree a -> Maybe a
searchMax t = case r of
    RBZip (Node _ v _ _) _ -> Just v
    _ -> Nothing
    where r = rightMostZip . toZip $ t

-- |Search the Minimum value in the tree, equals to get the left-most element.
searchMin :: (Ord a) => RBTree a -> Maybe a
searchMin t = case r of
    RBZip (Node _ v _ _) _ -> Just v
    _ -> Nothing
    where r = leftMostZip . toZip $ t


searchZip :: (b -> a -> Ordering) -> RBZip a -> b -> Maybe (RBZip a)
searchZip _ (RBZip Leaf _) _ = Nothing
searchZip f this@(RBZip (Node c v l r) path) vs = case f vs v of
    LT -> searchZip f (RBZip l ((Step c v ToLeft r):path)) vs
    GT -> searchZip f (RBZip r ((Step c v ToRight l):path)) vs
    EQ -> Just this

-- searchZipTrace : always returns the current point that the search stops.
-- returns a Zip-Node on equal, otherwise a Zip-Leaf
searchZipTrace :: (b -> a -> Ordering) -> RBZip a -> b -> RBZip a
searchZipTrace _ z@(RBZip Leaf _) _ = z
searchZipTrace f this@(RBZip (Node c v l r) path) vs = case f vs v of
    LT -> searchZipTrace f (RBZip l ((Step c v ToLeft r):path)) vs
    GT -> searchZipTrace f (RBZip r ((Step c v ToRight l):path)) vs
    EQ -> this

-- |Search \'Ord\' things, see 'searchInterval'
searchIntervalOrd :: (Ord a) => RBTree a -> a -> Interval a
searchIntervalOrd t a = searchInterval compare t a

-- |Search for a Interval.
--
--  For example: tree has 1,3,5,7. search for 3 returns [3,3] that indicates itself
--      search for 4 returns [3,5] indicates that 4 is between the element 3 and 5
--
--  The given value be or not be an element of the tree.
searchInterval :: (b -> a -> Ordering) -> RBTree a -> b -> Interval a
searchInterval f t a = case r of
    RBZip Leaf _ -> Interval (toNRealOrd (predZip r), toPRealOrd (succZip r))
    _ -> Interval (toNRealOrd r, toPRealOrd r)
    where r = searchZipTrace f (toZip t) a
          toNRealOrd (RBZip Leaf _) = NInfinity
          toNRealOrd (RBZip (Node _ v _ _) _) = RealValue v
          toPRealOrd (RBZip Leaf _) = PInfinity
          toPRealOrd (RBZip (Node _ v _ _) _) = RealValue v

-- Delete functions.

-- |Delete an \'Ord\' thing. see 'delete'.
deleteOrd :: (Ord a) => RBTree a -> a -> RBTree a
deleteOrd = delete compare

-- |Delete a sequence of elements.
deleteOrdList :: (Ord a) => RBTree a -> [a] -> RBTree a
deleteOrdList = foldl deleteOrd 

-- |If there is no relevant element in tree, tree will be returned unmodified.
delete :: (a -> a -> Ordering) -> RBTree a -> a -> RBTree a
delete f t a = 
    case searchZip f (toZip t) a of
        Just z -> toTree . deleteZip $ z
        Nothing -> t

-- |Delete Operator for deleteOrd
(<<\) :: (Ord a) => RBTree a -> a -> RBTree a
t <<\ e = deleteOrd t e

deleteZip :: RBZip a -> RBZip a
deleteZip z@(RBZip Leaf _) = z

-- case 1: left null
deleteZip (RBZip (Node c _ Leaf r) path) = case c of --r may be Leaf
    Red -> RBZip r path
    Black -> deleteFixup (RBZip r path)

-- case 2: right null
deleteZip (RBZip (Node c _ l Leaf) path) = case c of
    Red -> RBZip l path
    Black -> deleteFixup (RBZip l path)

-- case 3: both not null
deleteZip (RBZip (Node c _ l r@(Node _ vr srl _)) path) = deleteZip newX
    where !newX = leftMostZip (RBZip r ((Step c newV ToRight l):path))
          !newV = leftmostV vr srl

-- |fixup.
deleteFixup :: RBZip a -> RBZip a

-- endcase : 'a' may be Leaf!
deleteFixup (RBZip a@(Node Red _ _ _) path) = RBZip (setBlack a) path

-- case 1: brother of x is Red
deleteFixup (RBZip a ((Step _ vb db (Node Red vd l r)):path)) =
    deleteFixup $ RBZip a ((Step Red vb db newW):(Step Black vd db newS):path)
    where (!newW, !newS) = case db of
              ToLeft -> (l,r)
              ToRight -> (r,l)

-- case 4: x's brother s is black, but s's outter child is Red
-- c may be leaf
deleteFixup (RBZip a ((Step cb vb ToLeft (Node Black vd c e@(Node Red _ _ _))):path)) = 
    deleteFixup . topMostZip $ RBZip (Node cb vd (Node Black vb a c) (setBlack e)) path
deleteFixup (RBZip a ((Step cb vb ToRight (Node Black vd e@(Node Red _ _ _) c)):path)) = 
    deleteFixup . topMostZip $ RBZip (Node cb vd (setBlack e) (Node Black vb c a)) path

-- case 3: x's brother s is black, but s's inner child is Red
deleteFixup (RBZip a ((Step cb vb ToLeft (Node Black vd (Node Red vc scl scr) e)):path)) = 
    deleteFixup $ RBZip a ((Step cb vb ToLeft (Node Black vc scl (Node Red vd scr e))):path)
deleteFixup (RBZip a ((Step cb vb ToRight (Node Black vd e (Node Red vc scl scr))):path)) = 
    deleteFixup $ RBZip a ((Step cb vb ToRight (Node Black vc (Node Red vd e scl) scr)):path)

-- case 2: s's both children are not Red (Black or Leaf).
deleteFixup (RBZip a ((Step cb vb db d@(Node Black _ _ _)):path)) = 
    deleteFixup $ (RBZip (Node cb vb newL newR) path)
    where (!newL, !newR) = case db of
              ToLeft -> (a,d')
              ToRight -> (d',a)
          !d' = setRed d

-- any other case: set current node to black and return.
deleteFixup (RBZip a path) = RBZip (setBlack a) path

-- Verification functions

-- |Verify black-depth are all the same. 
--  Return Just \'depth\' on success, otherwise Nothing.
vD :: RBTree a -> Maybe Int
vD Leaf = Just 1
vD (Node c _ l r) = 
    case dl == dr of 
        True -> liftM2 (+) inc dl
        False -> Nothing
    where !dl = vD l
          !dr = vD r
          !inc = case c of
              Red -> Just 0
              Black -> Just 1

-- |vR : verify no \'red-red\' pattern in x and x\'s parent
vR :: RBTree a -> Bool
vR Leaf = True
vR (Node Black _ l r) = (vR l) && (vR r)
vR (Node Red _ l r) = 
    (cl /= Red) && (cr /= Red) && (vR l) && (vR r)
    where !cl = getColor l
          !cr = getColor r