-- 
-- 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 (Color -> Color -> Bool
(Color -> Color -> Bool) -> (Color -> Color -> Bool) -> Eq Color
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Color -> Color -> Bool
== :: Color -> Color -> Bool
$c/= :: Color -> Color -> Bool
/= :: Color -> Color -> Bool
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 (Direction -> Direction -> Bool
(Direction -> Direction -> Bool)
-> (Direction -> Direction -> Bool) -> Eq Direction
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Direction -> Direction -> Bool
== :: Direction -> Direction -> Bool
$c/= :: Direction -> Direction -> Bool
/= :: Direction -> Direction -> Bool
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 (Int -> Step a -> ShowS
[Step a] -> ShowS
Step a -> String
(Int -> Step a -> ShowS)
-> (Step a -> String) -> ([Step a] -> ShowS) -> Show (Step a)
forall a. Show a => Int -> Step a -> ShowS
forall a. Show a => [Step a] -> ShowS
forall a. Show a => Step a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> Step a -> ShowS
showsPrec :: Int -> Step a -> ShowS
$cshow :: forall a. Show a => Step a -> String
show :: Step a -> String
$cshowList :: forall a. Show a => [Step a] -> ShowS
showList :: [Step a] -> ShowS
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 (Int -> RBZip a -> ShowS
[RBZip a] -> ShowS
RBZip a -> String
(Int -> RBZip a -> ShowS)
-> (RBZip a -> String) -> ([RBZip a] -> ShowS) -> Show (RBZip a)
forall a. Show a => Int -> RBZip a -> ShowS
forall a. Show a => [RBZip a] -> ShowS
forall a. Show a => RBZip a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> RBZip a -> ShowS
showsPrec :: Int -> RBZip a -> ShowS
$cshow :: forall a. Show a => RBZip a -> String
show :: RBZip a -> String
$cshowList :: forall a. Show a => [RBZip a] -> ShowS
showList :: [RBZip a] -> ShowS
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 :: RBTree a -> String
show (Node Color
c a
v RBTree a
l RBTree a
r) = String
"(" String -> ShowS
forall a. [a] -> [a] -> [a]
++ RBTree a -> String
forall a. Show a => a -> String
show RBTree a
l String -> ShowS
forall a. [a] -> [a] -> [a]
++ a -> String
forall a. Show a => a -> String
show a
v String -> ShowS
forall a. [a] -> [a] -> [a]
++ Color -> String
forall a. Show a => a -> String
show Color
c String -> ShowS
forall a. [a] -> [a] -> [a]
++ RBTree a -> String
forall a. Show a => a -> String
show RBTree a
r String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"
    show RBTree a
Leaf = String
"."

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

instance Show Direction where
    show :: Direction -> String
show Direction
ToLeft = String
"L"
    show Direction
ToRight = String
"R"

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

instance Show a => Show (Interval a) where
    show :: Interval a -> String
show (Interval (RealOrd a
l, RealOrd a
r)) = String
"[" String -> ShowS
forall a. [a] -> [a] -> [a]
++ RealOrd a -> String
forall a. Show a => a -> String
show RealOrd a
l String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
", " String -> ShowS
forall a. [a] -> [a] -> [a]
++ RealOrd a -> String
forall a. Show a => a -> String
show RealOrd a
r String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"]"

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

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

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

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

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

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

-- |Zip up.
topMostZip :: RBZip a -> RBZip a
topMostZip :: forall a. RBZip a -> RBZip a
topMostZip (RBZip RBTree a
s ((Step Color
c a
v Direction
d RBTree a
s1):[Step a]
path)) = case Direction
d of 
        Direction
ToLeft -> RBZip a -> RBZip a
forall a. RBZip a -> RBZip a
topMostZip (RBTree a -> [Step a] -> RBZip a
forall a. RBTree a -> Path a -> RBZip a
RBZip (Color -> a -> RBTree a -> RBTree a -> RBTree a
forall a. Color -> a -> RBTree a -> RBTree a -> RBTree a
Node Color
c a
v RBTree a
s RBTree a
s1) [Step a]
path)
        Direction
ToRight -> RBZip a -> RBZip a
forall a. RBZip a -> RBZip a
topMostZip (RBTree a -> [Step a] -> RBZip a
forall a. RBTree a -> Path a -> RBZip a
RBZip (Color -> a -> RBTree a -> RBTree a -> RBTree a
forall a. Color -> a -> RBTree a -> RBTree a -> RBTree a
Node Color
c a
v RBTree a
s1 RBTree a
s) [Step a]
path)
topMostZip RBZip a
z = RBZip a
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 :: forall a. RBZip a -> RBZip a
leftMostZip this :: RBZip a
this@(RBZip (Node Color
_ a
_ RBTree a
Leaf RBTree a
_) Path a
_) = RBZip a
this
leftMostZip (RBZip (Node Color
c a
v RBTree a
l RBTree a
r) Path a
path) = RBZip a -> RBZip a
forall a. RBZip a -> RBZip a
leftMostZip (RBTree a -> Path a -> RBZip a
forall a. RBTree a -> Path a -> RBZip a
RBZip RBTree a
l ((Color -> a -> Direction -> RBTree a -> Step a
forall a. Color -> a -> Direction -> RBTree a -> Step a
Step Color
c a
v Direction
ToLeft RBTree a
r)Step a -> Path a -> Path a
forall a. a -> [a] -> [a]
:Path a
path))
leftMostZip RBZip a
z = RBZip a
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 :: forall a. RBZip a -> RBZip a
rightMostZip this :: RBZip a
this@(RBZip (Node Color
_ a
_ RBTree a
_ RBTree a
Leaf) Path a
_) = RBZip a
this
rightMostZip (RBZip (Node Color
c a
v RBTree a
l RBTree a
r) Path a
path) = RBZip a -> RBZip a
forall a. RBZip a -> RBZip a
rightMostZip (RBTree a -> Path a -> RBZip a
forall a. RBTree a -> Path a -> RBZip a
RBZip RBTree a
r ((Color -> a -> Direction -> RBTree a -> Step a
forall a. Color -> a -> Direction -> RBTree a -> Step a
Step Color
c a
v Direction
ToRight RBTree a
l)Step a -> Path a -> Path a
forall a. a -> [a] -> [a]
:Path a
path))
rightMostZip RBZip a
z = RBZip a
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 :: forall a. RBZip a -> RBZip a
leftParentZip (RBZip RBTree a
l ((Step Color
c a
v Direction
ToLeft RBTree a
r):[Step a]
path)) = RBZip a -> RBZip a
forall a. RBZip a -> RBZip a
leftParentZip (RBTree a -> [Step a] -> RBZip a
forall a. RBTree a -> Path a -> RBZip a
RBZip (Color -> a -> RBTree a -> RBTree a -> RBTree a
forall a. Color -> a -> RBTree a -> RBTree a -> RBTree a
Node Color
c a
v RBTree a
l RBTree a
r) [Step a]
path)
leftParentZip (RBZip RBTree a
r ((Step Color
c a
v Direction
ToRight RBTree a
l):[Step a]
path)) = RBTree a -> [Step a] -> RBZip a
forall a. RBTree a -> Path a -> RBZip a
RBZip (Color -> a -> RBTree a -> RBTree a -> RBTree a
forall a. Color -> a -> RBTree a -> RBTree a -> RBTree a
Node Color
c a
v RBTree a
l RBTree a
r) [Step a]
path
leftParentZip (RBZip RBTree a
_ []) = RBTree a -> [Step a] -> RBZip a
forall a. RBTree a -> Path a -> RBZip a
RBZip RBTree a
forall a. RBTree a
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 :: forall a. RBZip a -> RBZip a
rightParentZip (RBZip RBTree a
r ((Step Color
c a
v Direction
ToRight RBTree a
l):[Step a]
path)) = RBZip a -> RBZip a
forall a. RBZip a -> RBZip a
rightParentZip (RBTree a -> [Step a] -> RBZip a
forall a. RBTree a -> Path a -> RBZip a
RBZip (Color -> a -> RBTree a -> RBTree a -> RBTree a
forall a. Color -> a -> RBTree a -> RBTree a -> RBTree a
Node Color
c a
v RBTree a
l RBTree a
r) [Step a]
path)
rightParentZip (RBZip RBTree a
l ((Step Color
c a
v Direction
ToLeft RBTree a
r):[Step a]
path)) = RBTree a -> [Step a] -> RBZip a
forall a. RBTree a -> Path a -> RBZip a
RBZip (Color -> a -> RBTree a -> RBTree a -> RBTree a
forall a. Color -> a -> RBTree a -> RBTree a -> RBTree a
Node Color
c a
v RBTree a
l RBTree a
r) [Step a]
path
rightParentZip (RBZip RBTree a
_ []) = RBTree a -> [Step a] -> RBZip a
forall a. RBTree a -> Path a -> RBZip a
RBZip RBTree a
forall a. RBTree a
Leaf [] -- no such parent, return a empty zip

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

-- |find predecessor of a node/leaf.
succZip :: RBZip a -> RBZip a
succZip :: forall a. RBZip a -> RBZip a
succZip (RBZip (Node Color
c a
v RBTree a
l r :: RBTree a
r@(Node Color
_ a
_ RBTree a
_ RBTree a
_)) Path a
path) = RBZip a -> RBZip a
forall a. RBZip a -> RBZip a
leftMostZip (RBTree a -> Path a -> RBZip a
forall a. RBTree a -> Path a -> RBZip a
RBZip RBTree a
r ((Color -> a -> Direction -> RBTree a -> Step a
forall a. Color -> a -> Direction -> RBTree a -> Step a
Step Color
c a
v Direction
ToRight RBTree a
l)Step a -> Path a -> Path a
forall a. a -> [a] -> [a]
:Path a
path))
succZip z :: RBZip a
z@(RBZip RBTree a
Leaf Path a
_) = case RBZip a
lp of
  RBZip RBTree a
Leaf [] -> RBZip a
z -- itself
  RBZip a
_ -> RBZip a
lp
  where lp :: RBZip a
lp = RBZip a -> RBZip a
forall a. RBZip a -> RBZip a
rightParentZip RBZip a
z
succZip z :: RBZip a
z@(RBZip (Node Color
c a
v RBTree a
l RBTree a
r) Path a
path) = case RBZip a
lp of
  RBZip RBTree a
Leaf [] -> RBTree a -> Path a -> RBZip a
forall a. RBTree a -> Path a -> RBZip a
RBZip RBTree a
r ((Color -> a -> Direction -> RBTree a -> Step a
forall a. Color -> a -> Direction -> RBTree a -> Step a
Step Color
c a
v Direction
ToRight RBTree a
l)Step a -> Path a -> Path a
forall a. a -> [a] -> [a]
:Path a
path)
  RBZip a
_ -> RBZip a
lp
  where lp :: RBZip a
lp = RBZip a -> RBZip a
forall a. RBZip a -> RBZip a
rightParentZip RBZip a
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 :: forall a. a -> RBTree a -> a
leftmostV a
v RBTree a
Leaf = a
v
leftmostV a
_ (Node Color
_ a
vl RBTree a
l RBTree a
_) = a -> RBTree a -> a
forall a. a -> RBTree a -> a
leftmostV a
vl RBTree a
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 :: forall a. Ord a => RBTree a -> a -> RBTree a
insertOrd = (a -> a -> Ordering) -> RBTree a -> a -> RBTree a
forall a. (a -> a -> Ordering) -> RBTree a -> a -> RBTree a
insert a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare

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

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

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

insertRedZip :: (a -> a -> Ordering) -> RBZip a -> a -> RBZip a
insertRedZip :: forall a. (a -> a -> Ordering) -> RBZip a -> a -> RBZip a
insertRedZip a -> a -> Ordering
_ (RBZip RBTree a
Leaf Path a
path) a
v = RBTree a -> Path a -> RBZip a
forall a. RBTree a -> Path a -> RBZip a
RBZip (Color -> a -> RBTree a -> RBTree a -> RBTree a
forall a. Color -> a -> RBTree a -> RBTree a -> RBTree a
Node Color
Red a
v RBTree a
forall a. RBTree a
Leaf RBTree a
forall a. RBTree a
Leaf) Path a
path
insertRedZip a -> a -> Ordering
f (RBZip (Node Color
c a
v RBTree a
l RBTree a
r) Path a
path) a
new
    | a -> a -> Ordering
f a
new a
v Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== Ordering
GT = (a -> a -> Ordering) -> RBZip a -> a -> RBZip a
forall a. (a -> a -> Ordering) -> RBZip a -> a -> RBZip a
insertRedZip a -> a -> Ordering
f (RBTree a -> Path a -> RBZip a
forall a. RBTree a -> Path a -> RBZip a
RBZip RBTree a
r ((Color -> a -> Direction -> RBTree a -> Step a
forall a. Color -> a -> Direction -> RBTree a -> Step a
Step Color
c a
v Direction
ToRight RBTree a
l)Step a -> Path a -> Path a
forall a. a -> [a] -> [a]
:Path a
path)) a
new
    | Bool
otherwise     = (a -> a -> Ordering) -> RBZip a -> a -> RBZip a
forall a. (a -> a -> Ordering) -> RBZip a -> a -> RBZip a
insertRedZip a -> a -> Ordering
f (RBTree a -> Path a -> RBZip a
forall a. RBTree a -> Path a -> RBZip a
RBZip RBTree a
l ((Color -> a -> Direction -> RBTree a -> Step a
forall a. Color -> a -> Direction -> RBTree a -> Step a
Step Color
c a
v Direction
ToLeft RBTree a
r)Step a -> Path a -> Path a
forall a. a -> [a] -> [a]
:Path a
path)) a
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 :: forall a. RBZip a -> RBZip a
insertFixup (RBZip a :: RBTree a
a@(Node Color
Red a
_ RBTree a
_ RBTree a
_) ((Step Color
Red a
vb Direction
db RBTree a
sb):(Step Color
Black a
vc Direction
dc d :: RBTree a
d@(Node Color
Red a
_ RBTree a
_ RBTree a
_)):[Step a]
path)) =
    RBZip a -> RBZip a
forall a. RBZip a -> RBZip a
insertFixup (RBTree a -> [Step a] -> RBZip a
forall a. RBTree a -> Path a -> RBZip a
RBZip RBTree a
newC [Step a]
path)
    where newC :: RBTree a
newC = Color -> a -> RBTree a -> RBTree a -> RBTree a
forall a. Color -> a -> RBTree a -> RBTree a -> RBTree a
Node Color
Red a
vc RBTree a
newCL RBTree a
newCR
          (RBTree a
newCL,RBTree a
newCR) = case Direction
dc of
              Direction
ToLeft -> (RBTree a
newB,RBTree a
newD)
              Direction
ToRight -> (RBTree a
newD,RBTree a
newB)
          newB :: RBTree a
newB = Color -> a -> RBTree a -> RBTree a -> RBTree a
forall a. Color -> a -> RBTree a -> RBTree a -> RBTree a
Node Color
Black a
vb RBTree a
newBL RBTree a
newBR
          (RBTree a
newBL,RBTree a
newBR) = case Direction
db of
              Direction
ToLeft -> (RBTree a
a,RBTree a
sb)
              Direction
ToRight -> (RBTree a
sb,RBTree a
a)
          !newD :: RBTree a
newD = RBTree a -> RBTree a
forall a. RBTree a -> RBTree a
setBlack RBTree a
d
insertFixup (RBZip a :: RBTree a
a@(Node Color
Red a
va RBTree a
sal RBTree a
sar) ((Step Color
Red a
vb Direction
db RBTree a
sb):(Step Color
Black a
vc Direction
dc RBTree a
d):[Step a]
path)) =
    RBTree a -> [Step a] -> RBZip a
forall a. RBTree a -> Path a -> RBZip a
RBZip RBTree a
newZ (Step a
newPStep a -> [Step a] -> [Step a]
forall a. a -> [a] -> [a]
:[Step a]
path)
    where (RBTree a
newZ, Step a
newP) = case (Direction
dc,Direction
db) of 
              (Direction
ToLeft,Direction
ToLeft) -> (RBTree a
a,Color -> a -> Direction -> RBTree a -> Step a
forall a. Color -> a -> Direction -> RBTree a -> Step a
Step Color
Black a
vb Direction
dc (Color -> a -> RBTree a -> RBTree a -> RBTree a
forall a. Color -> a -> RBTree a -> RBTree a -> RBTree a
Node Color
Red a
vc RBTree a
sb RBTree a
d))
              (Direction
ToLeft,Direction
ToRight) -> (Color -> a -> RBTree a -> RBTree a -> RBTree a
forall a. Color -> a -> RBTree a -> RBTree a -> RBTree a
Node Color
Red a
vb RBTree a
sb RBTree a
sal, Color -> a -> Direction -> RBTree a -> Step a
forall a. Color -> a -> Direction -> RBTree a -> Step a
Step Color
Black a
va Direction
dc (Color -> a -> RBTree a -> RBTree a -> RBTree a
forall a. Color -> a -> RBTree a -> RBTree a -> RBTree a
Node Color
Red a
vc RBTree a
sar RBTree a
d))
              (Direction
ToRight,Direction
ToLeft) -> (Color -> a -> RBTree a -> RBTree a -> RBTree a
forall a. Color -> a -> RBTree a -> RBTree a -> RBTree a
Node Color
Red a
vb RBTree a
sar RBTree a
sb, Color -> a -> Direction -> RBTree a -> Step a
forall a. Color -> a -> Direction -> RBTree a -> Step a
Step Color
Black a
va Direction
dc (Color -> a -> RBTree a -> RBTree a -> RBTree a
forall a. Color -> a -> RBTree a -> RBTree a -> RBTree a
Node Color
Red a
vc RBTree a
d RBTree a
sal))
              (Direction
ToRight,Direction
ToRight) -> (RBTree a
a,Color -> a -> Direction -> RBTree a -> Step a
forall a. Color -> a -> Direction -> RBTree a -> Step a
Step Color
Black a
vb Direction
dc (Color -> a -> RBTree a -> RBTree a -> RBTree a
forall a. Color -> a -> RBTree a -> RBTree a -> RBTree a
Node Color
Red a
vc RBTree a
d RBTree a
sb))
insertFixup RBZip a
t = RBZip a
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 :: forall a. Ord a => RBTree a -> a -> Maybe a
searchOrd = (a -> a -> Ordering) -> RBTree a -> a -> Maybe a
forall b a. (b -> a -> Ordering) -> RBTree a -> b -> Maybe a
search a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare

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

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

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

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

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


searchZip :: (b -> a -> Ordering) -> RBZip a -> b -> Maybe (RBZip a)
searchZip :: forall b a. (b -> a -> Ordering) -> RBZip a -> b -> Maybe (RBZip a)
searchZip b -> a -> Ordering
_ (RBZip RBTree a
Leaf Path a
_) b
_ = Maybe (RBZip a)
forall a. Maybe a
Nothing
searchZip b -> a -> Ordering
f this :: RBZip a
this@(RBZip (Node Color
c a
v RBTree a
l RBTree a
r) Path a
path) b
vs = case b -> a -> Ordering
f b
vs a
v of
    Ordering
LT -> (b -> a -> Ordering) -> RBZip a -> b -> Maybe (RBZip a)
forall b a. (b -> a -> Ordering) -> RBZip a -> b -> Maybe (RBZip a)
searchZip b -> a -> Ordering
f (RBTree a -> Path a -> RBZip a
forall a. RBTree a -> Path a -> RBZip a
RBZip RBTree a
l ((Color -> a -> Direction -> RBTree a -> Step a
forall a. Color -> a -> Direction -> RBTree a -> Step a
Step Color
c a
v Direction
ToLeft RBTree a
r)Step a -> Path a -> Path a
forall a. a -> [a] -> [a]
:Path a
path)) b
vs
    Ordering
GT -> (b -> a -> Ordering) -> RBZip a -> b -> Maybe (RBZip a)
forall b a. (b -> a -> Ordering) -> RBZip a -> b -> Maybe (RBZip a)
searchZip b -> a -> Ordering
f (RBTree a -> Path a -> RBZip a
forall a. RBTree a -> Path a -> RBZip a
RBZip RBTree a
r ((Color -> a -> Direction -> RBTree a -> Step a
forall a. Color -> a -> Direction -> RBTree a -> Step a
Step Color
c a
v Direction
ToRight RBTree a
l)Step a -> Path a -> Path a
forall a. a -> [a] -> [a]
:Path a
path)) b
vs
    Ordering
EQ -> RBZip a -> Maybe (RBZip a)
forall a. a -> Maybe a
Just RBZip a
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 :: forall b a. (b -> a -> Ordering) -> RBZip a -> b -> RBZip a
searchZipTrace b -> a -> Ordering
_ z :: RBZip a
z@(RBZip RBTree a
Leaf Path a
_) b
_ = RBZip a
z
searchZipTrace b -> a -> Ordering
f this :: RBZip a
this@(RBZip (Node Color
c a
v RBTree a
l RBTree a
r) Path a
path) b
vs = case b -> a -> Ordering
f b
vs a
v of
    Ordering
LT -> (b -> a -> Ordering) -> RBZip a -> b -> RBZip a
forall b a. (b -> a -> Ordering) -> RBZip a -> b -> RBZip a
searchZipTrace b -> a -> Ordering
f (RBTree a -> Path a -> RBZip a
forall a. RBTree a -> Path a -> RBZip a
RBZip RBTree a
l ((Color -> a -> Direction -> RBTree a -> Step a
forall a. Color -> a -> Direction -> RBTree a -> Step a
Step Color
c a
v Direction
ToLeft RBTree a
r)Step a -> Path a -> Path a
forall a. a -> [a] -> [a]
:Path a
path)) b
vs
    Ordering
GT -> (b -> a -> Ordering) -> RBZip a -> b -> RBZip a
forall b a. (b -> a -> Ordering) -> RBZip a -> b -> RBZip a
searchZipTrace b -> a -> Ordering
f (RBTree a -> Path a -> RBZip a
forall a. RBTree a -> Path a -> RBZip a
RBZip RBTree a
r ((Color -> a -> Direction -> RBTree a -> Step a
forall a. Color -> a -> Direction -> RBTree a -> Step a
Step Color
c a
v Direction
ToRight RBTree a
l)Step a -> Path a -> Path a
forall a. a -> [a] -> [a]
:Path a
path)) b
vs
    Ordering
EQ -> RBZip a
this

-- |Search \'Ord\' things, see 'searchInterval'
searchIntervalOrd :: (Ord a) => RBTree a -> a -> Interval a
searchIntervalOrd :: forall a. Ord a => RBTree a -> a -> Interval a
searchIntervalOrd RBTree a
t a
a = (a -> a -> Ordering) -> RBTree a -> a -> Interval a
forall b a. (b -> a -> Ordering) -> RBTree a -> b -> Interval a
searchInterval a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare RBTree a
t a
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 :: forall b a. (b -> a -> Ordering) -> RBTree a -> b -> Interval a
searchInterval b -> a -> Ordering
f RBTree a
t b
a = case RBZip a
r of
    RBZip RBTree a
Leaf Path a
_ -> (RealOrd a, RealOrd a) -> Interval a
forall a. (RealOrd a, RealOrd a) -> Interval a
Interval (RBZip a -> RealOrd a
forall {a}. RBZip a -> RealOrd a
toNRealOrd (RBZip a -> RBZip a
forall a. RBZip a -> RBZip a
predZip RBZip a
r), RBZip a -> RealOrd a
forall {a}. RBZip a -> RealOrd a
toPRealOrd (RBZip a -> RBZip a
forall a. RBZip a -> RBZip a
succZip RBZip a
r))
    RBZip a
_ -> (RealOrd a, RealOrd a) -> Interval a
forall a. (RealOrd a, RealOrd a) -> Interval a
Interval (RBZip a -> RealOrd a
forall {a}. RBZip a -> RealOrd a
toNRealOrd RBZip a
r, RBZip a -> RealOrd a
forall {a}. RBZip a -> RealOrd a
toPRealOrd RBZip a
r)
    where r :: RBZip a
r = (b -> a -> Ordering) -> RBZip a -> b -> RBZip a
forall b a. (b -> a -> Ordering) -> RBZip a -> b -> RBZip a
searchZipTrace b -> a -> Ordering
f (RBTree a -> RBZip a
forall a. RBTree a -> RBZip a
toZip RBTree a
t) b
a
          toNRealOrd :: RBZip a -> RealOrd a
toNRealOrd (RBZip RBTree a
Leaf Path a
_) = RealOrd a
forall a. RealOrd a
NInfinity
          toNRealOrd (RBZip (Node Color
_ a
v RBTree a
_ RBTree a
_) Path a
_) = a -> RealOrd a
forall a. a -> RealOrd a
RealValue a
v
          toPRealOrd :: RBZip a -> RealOrd a
toPRealOrd (RBZip RBTree a
Leaf Path a
_) = RealOrd a
forall a. RealOrd a
PInfinity
          toPRealOrd (RBZip (Node Color
_ a
v RBTree a
_ RBTree a
_) Path a
_) = a -> RealOrd a
forall a. a -> RealOrd a
RealValue a
v

-- Delete functions.

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

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

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

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

deleteZip :: RBZip a -> RBZip a
deleteZip :: forall a. RBZip a -> RBZip a
deleteZip z :: RBZip a
z@(RBZip RBTree a
Leaf Path a
_) = RBZip a
z

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

-- case 2: right null
deleteZip (RBZip (Node Color
c a
_ RBTree a
l RBTree a
Leaf) Path a
path) = case Color
c of
    Color
Red -> RBTree a -> Path a -> RBZip a
forall a. RBTree a -> Path a -> RBZip a
RBZip RBTree a
l Path a
path
    Color
Black -> RBZip a -> RBZip a
forall a. RBZip a -> RBZip a
deleteFixup (RBTree a -> Path a -> RBZip a
forall a. RBTree a -> Path a -> RBZip a
RBZip RBTree a
l Path a
path)

-- case 3: both not null
deleteZip (RBZip (Node Color
c a
_ RBTree a
l r :: RBTree a
r@(Node Color
_ a
vr RBTree a
srl RBTree a
_)) Path a
path) = RBZip a -> RBZip a
forall a. RBZip a -> RBZip a
deleteZip RBZip a
newX
    where !newX :: RBZip a
newX = RBZip a -> RBZip a
forall a. RBZip a -> RBZip a
leftMostZip (RBTree a -> Path a -> RBZip a
forall a. RBTree a -> Path a -> RBZip a
RBZip RBTree a
r ((Color -> a -> Direction -> RBTree a -> Step a
forall a. Color -> a -> Direction -> RBTree a -> Step a
Step Color
c a
newV Direction
ToRight RBTree a
l)Step a -> Path a -> Path a
forall a. a -> [a] -> [a]
:Path a
path))
          !newV :: a
newV = a -> RBTree a -> a
forall a. a -> RBTree a -> a
leftmostV a
vr RBTree a
srl

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

-- endcase : 'a' may be Leaf!
deleteFixup :: forall a. RBZip a -> RBZip a
deleteFixup (RBZip a :: RBTree a
a@(Node Color
Red a
_ RBTree a
_ RBTree a
_) Path a
path) = RBTree a -> Path a -> RBZip a
forall a. RBTree a -> Path a -> RBZip a
RBZip (RBTree a -> RBTree a
forall a. RBTree a -> RBTree a
setBlack RBTree a
a) Path a
path

-- case 1: brother of x is Red
deleteFixup (RBZip RBTree a
a ((Step Color
_ a
vb Direction
db (Node Color
Red a
vd RBTree a
l RBTree a
r)):Path a
path)) =
    RBZip a -> RBZip a
forall a. RBZip a -> RBZip a
deleteFixup (RBZip a -> RBZip a) -> RBZip a -> RBZip a
forall a b. (a -> b) -> a -> b
$ RBTree a -> Path a -> RBZip a
forall a. RBTree a -> Path a -> RBZip a
RBZip RBTree a
a ((Color -> a -> Direction -> RBTree a -> Step a
forall a. Color -> a -> Direction -> RBTree a -> Step a
Step Color
Red a
vb Direction
db RBTree a
newW)Step a -> Path a -> Path a
forall a. a -> [a] -> [a]
:(Color -> a -> Direction -> RBTree a -> Step a
forall a. Color -> a -> Direction -> RBTree a -> Step a
Step Color
Black a
vd Direction
db RBTree a
newS)Step a -> Path a -> Path a
forall a. a -> [a] -> [a]
:Path a
path)
    where (!RBTree a
newW, !RBTree a
newS) = case Direction
db of
              Direction
ToLeft -> (RBTree a
l,RBTree a
r)
              Direction
ToRight -> (RBTree a
r,RBTree a
l)

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

-- case 3: x's brother s is black, but s's inner child is Red
deleteFixup (RBZip RBTree a
a ((Step Color
cb a
vb Direction
ToLeft (Node Color
Black a
vd (Node Color
Red a
vc RBTree a
scl RBTree a
scr) RBTree a
e)):Path a
path)) = 
    RBZip a -> RBZip a
forall a. RBZip a -> RBZip a
deleteFixup (RBZip a -> RBZip a) -> RBZip a -> RBZip a
forall a b. (a -> b) -> a -> b
$ RBTree a -> Path a -> RBZip a
forall a. RBTree a -> Path a -> RBZip a
RBZip RBTree a
a ((Color -> a -> Direction -> RBTree a -> Step a
forall a. Color -> a -> Direction -> RBTree a -> Step a
Step Color
cb a
vb Direction
ToLeft (Color -> a -> RBTree a -> RBTree a -> RBTree a
forall a. Color -> a -> RBTree a -> RBTree a -> RBTree a
Node Color
Black a
vc RBTree a
scl (Color -> a -> RBTree a -> RBTree a -> RBTree a
forall a. Color -> a -> RBTree a -> RBTree a -> RBTree a
Node Color
Red a
vd RBTree a
scr RBTree a
e)))Step a -> Path a -> Path a
forall a. a -> [a] -> [a]
:Path a
path)
deleteFixup (RBZip RBTree a
a ((Step Color
cb a
vb Direction
ToRight (Node Color
Black a
vd RBTree a
e (Node Color
Red a
vc RBTree a
scl RBTree a
scr))):Path a
path)) = 
    RBZip a -> RBZip a
forall a. RBZip a -> RBZip a
deleteFixup (RBZip a -> RBZip a) -> RBZip a -> RBZip a
forall a b. (a -> b) -> a -> b
$ RBTree a -> Path a -> RBZip a
forall a. RBTree a -> Path a -> RBZip a
RBZip RBTree a
a ((Color -> a -> Direction -> RBTree a -> Step a
forall a. Color -> a -> Direction -> RBTree a -> Step a
Step Color
cb a
vb Direction
ToRight (Color -> a -> RBTree a -> RBTree a -> RBTree a
forall a. Color -> a -> RBTree a -> RBTree a -> RBTree a
Node Color
Black a
vc (Color -> a -> RBTree a -> RBTree a -> RBTree a
forall a. Color -> a -> RBTree a -> RBTree a -> RBTree a
Node Color
Red a
vd RBTree a
e RBTree a
scl) RBTree a
scr))Step a -> Path a -> Path a
forall a. a -> [a] -> [a]
:Path a
path)

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

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

-- Verification functions

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

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