-- |
-- Module      : Data.RedBlackTree.Internal
-- Copyright   : (c) 2017 Gabriel Aumala
--
-- License     : BSD3
-- Maintainer  : gabriel@criptext.com
-- Stability   : experimental
-- Portability : GHC
--
-- Data types and functions used internally by "Data.RedBlackTree". You don't need
-- to know anything about this if you only want to consume the "RedBlackTree"
-- library.
module Data.RedBlackTree.Internal (
  branchIsColor,
  getBlackHeight,
  isColor,
  emptyRedBlackTree,
  find,
  paintItBlack,
  removeBranchColor,
  whiteBranch2Tree,

  RedBlack (Red, Black),
  RedBlackNode (RedBlackNode),
  RedBlackBranch,
  RedBlackTree,
  RedBlackDirection,
  RedBlackDirections,
  WhiteBranch (WhiteBranch)
) where

import Data.RedBlackTree.BinaryTree

-- | Red black trees can only have two types of nodes: Red and Black.
data RedBlack = Red | Black deriving (Int -> RedBlack -> ShowS
[RedBlack] -> ShowS
RedBlack -> String
(Int -> RedBlack -> ShowS)
-> (RedBlack -> String) -> ([RedBlack] -> ShowS) -> Show RedBlack
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> RedBlack -> ShowS
showsPrec :: Int -> RedBlack -> ShowS
$cshow :: RedBlack -> String
show :: RedBlack -> String
$cshowList :: [RedBlack] -> ShowS
showList :: [RedBlack] -> ShowS
Show, RedBlack -> RedBlack -> Bool
(RedBlack -> RedBlack -> Bool)
-> (RedBlack -> RedBlack -> Bool) -> Eq RedBlack
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: RedBlack -> RedBlack -> Bool
== :: RedBlack -> RedBlack -> Bool
$c/= :: RedBlack -> RedBlack -> Bool
/= :: RedBlack -> RedBlack -> Bool
Eq, Eq RedBlack
Eq RedBlack =>
(RedBlack -> RedBlack -> Ordering)
-> (RedBlack -> RedBlack -> Bool)
-> (RedBlack -> RedBlack -> Bool)
-> (RedBlack -> RedBlack -> Bool)
-> (RedBlack -> RedBlack -> Bool)
-> (RedBlack -> RedBlack -> RedBlack)
-> (RedBlack -> RedBlack -> RedBlack)
-> Ord RedBlack
RedBlack -> RedBlack -> Bool
RedBlack -> RedBlack -> Ordering
RedBlack -> RedBlack -> RedBlack
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: RedBlack -> RedBlack -> Ordering
compare :: RedBlack -> RedBlack -> Ordering
$c< :: RedBlack -> RedBlack -> Bool
< :: RedBlack -> RedBlack -> Bool
$c<= :: RedBlack -> RedBlack -> Bool
<= :: RedBlack -> RedBlack -> Bool
$c> :: RedBlack -> RedBlack -> Bool
> :: RedBlack -> RedBlack -> Bool
$c>= :: RedBlack -> RedBlack -> Bool
>= :: RedBlack -> RedBlack -> Bool
$cmax :: RedBlack -> RedBlack -> RedBlack
max :: RedBlack -> RedBlack -> RedBlack
$cmin :: RedBlack -> RedBlack -> RedBlack
min :: RedBlack -> RedBlack -> RedBlack
Ord)

-- | a @RedBlackNode@ contains only two elements, the color of the node and the
-- actual content.
data RedBlackNode a = RedBlackNode {
  forall a. RedBlackNode a -> RedBlack
nodeColor :: RedBlack,
  forall a. RedBlackNode a -> a
content :: a
} deriving (Int -> RedBlackNode a -> ShowS
[RedBlackNode a] -> ShowS
RedBlackNode a -> String
(Int -> RedBlackNode a -> ShowS)
-> (RedBlackNode a -> String)
-> ([RedBlackNode a] -> ShowS)
-> Show (RedBlackNode a)
forall a. Show a => Int -> RedBlackNode a -> ShowS
forall a. Show a => [RedBlackNode a] -> ShowS
forall a. Show a => RedBlackNode a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> RedBlackNode a -> ShowS
showsPrec :: Int -> RedBlackNode a -> ShowS
$cshow :: forall a. Show a => RedBlackNode a -> String
show :: RedBlackNode a -> String
$cshowList :: forall a. Show a => [RedBlackNode a] -> ShowS
showList :: [RedBlackNode a] -> ShowS
Show)

instance (BinaryTreeNode a) => BinaryTreeNode (RedBlackNode a)  where
  mergeNodes :: RedBlackNode a -> RedBlackNode a -> RedBlackNode a
mergeNodes RedBlackNode a
leftNode RedBlackNode a
rightNode = RedBlack -> a -> RedBlackNode a
forall a. RedBlack -> a -> RedBlackNode a
RedBlackNode RedBlack
color a
mergedContent
    where RedBlackNode RedBlack
color a
leftContent = RedBlackNode a
leftNode
          RedBlackNode RedBlack
_ a
rightContent = RedBlackNode a
rightNode
          mergedContent :: a
mergedContent = a
leftContent a -> a -> a
forall a. BinaryTreeNode a => a -> a -> a
`mergeNodes` a
rightContent

instance (BinaryTreeNode a) => Ord (RedBlackNode a) where
  (RedBlackNode RedBlack
_ a
lcontent) <= :: RedBlackNode a -> RedBlackNode a -> Bool
<= (RedBlackNode RedBlack
_ a
rcontent) =
    a
lcontent a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
rcontent

instance (BinaryTreeNode a) => Eq (RedBlackNode a) where
  (RedBlackNode RedBlack
_ a
lcontent) == :: RedBlackNode a -> RedBlackNode a -> Bool
== (RedBlackNode RedBlack
_ a
rcontent) =
    a
lcontent a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
rcontent



-- | A "BinaryTree" with only nodes of type @RedBlackNode@. is either a @Leaf@
-- (empty) or a @RedBlackNode@ with 2 @RedBlackTree@ children: left and right
type RedBlackTree a = BinaryTree (RedBlackNode a)

-- A @TreeBranch@ with only nodes of type @RedBlackNode. Holds the data of a
-- @RedBlackTree@ created with the @Branch@ constructor. Useful
-- type when you want to guarantee that the element is not a @Leaf@
type RedBlackBranch a = TreeBranch (RedBlackNode a)

-- @TreeDirection@ for trees of type @RedBlackTree@. Minimum necessary to
-- reconstruct the parent of any focused node. First argument is the @BranchType@
-- of the focused node relative to the parent. Second argument is the parent's
-- node. The third argument is the sibling tree of the focused node.
type RedBlackDirection a = TreeDirection (RedBlackNode a)

-- List of @RedBlackDirection@
type RedBlackDirections a = [ RedBlackDirection a ]

-- Holds all the data of a @RedBlackBranch@ except for the color of the node
-- at the top of the branch
data WhiteBranch a = WhiteBranch (RedBlackTree a) a (RedBlackTree a)
  deriving (WhiteBranch a -> WhiteBranch a -> Bool
(WhiteBranch a -> WhiteBranch a -> Bool)
-> (WhiteBranch a -> WhiteBranch a -> Bool) -> Eq (WhiteBranch a)
forall a.
BinaryTreeNode a =>
WhiteBranch a -> WhiteBranch a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a.
BinaryTreeNode a =>
WhiteBranch a -> WhiteBranch a -> Bool
== :: WhiteBranch a -> WhiteBranch a -> Bool
$c/= :: forall a.
BinaryTreeNode a =>
WhiteBranch a -> WhiteBranch a -> Bool
/= :: WhiteBranch a -> WhiteBranch a -> Bool
Eq, Eq (WhiteBranch a)
Eq (WhiteBranch a) =>
(WhiteBranch a -> WhiteBranch a -> Ordering)
-> (WhiteBranch a -> WhiteBranch a -> Bool)
-> (WhiteBranch a -> WhiteBranch a -> Bool)
-> (WhiteBranch a -> WhiteBranch a -> Bool)
-> (WhiteBranch a -> WhiteBranch a -> Bool)
-> (WhiteBranch a -> WhiteBranch a -> WhiteBranch a)
-> (WhiteBranch a -> WhiteBranch a -> WhiteBranch a)
-> Ord (WhiteBranch a)
WhiteBranch a -> WhiteBranch a -> Bool
WhiteBranch a -> WhiteBranch a -> Ordering
WhiteBranch a -> WhiteBranch a -> WhiteBranch a
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall a. BinaryTreeNode a => Eq (WhiteBranch a)
forall a.
BinaryTreeNode a =>
WhiteBranch a -> WhiteBranch a -> Bool
forall a.
BinaryTreeNode a =>
WhiteBranch a -> WhiteBranch a -> Ordering
forall a.
BinaryTreeNode a =>
WhiteBranch a -> WhiteBranch a -> WhiteBranch a
$ccompare :: forall a.
BinaryTreeNode a =>
WhiteBranch a -> WhiteBranch a -> Ordering
compare :: WhiteBranch a -> WhiteBranch a -> Ordering
$c< :: forall a.
BinaryTreeNode a =>
WhiteBranch a -> WhiteBranch a -> Bool
< :: WhiteBranch a -> WhiteBranch a -> Bool
$c<= :: forall a.
BinaryTreeNode a =>
WhiteBranch a -> WhiteBranch a -> Bool
<= :: WhiteBranch a -> WhiteBranch a -> Bool
$c> :: forall a.
BinaryTreeNode a =>
WhiteBranch a -> WhiteBranch a -> Bool
> :: WhiteBranch a -> WhiteBranch a -> Bool
$c>= :: forall a.
BinaryTreeNode a =>
WhiteBranch a -> WhiteBranch a -> Bool
>= :: WhiteBranch a -> WhiteBranch a -> Bool
$cmax :: forall a.
BinaryTreeNode a =>
WhiteBranch a -> WhiteBranch a -> WhiteBranch a
max :: WhiteBranch a -> WhiteBranch a -> WhiteBranch a
$cmin :: forall a.
BinaryTreeNode a =>
WhiteBranch a -> WhiteBranch a -> WhiteBranch a
min :: WhiteBranch a -> WhiteBranch a -> WhiteBranch a
Ord, Int -> WhiteBranch a -> ShowS
[WhiteBranch a] -> ShowS
WhiteBranch a -> String
(Int -> WhiteBranch a -> ShowS)
-> (WhiteBranch a -> String)
-> ([WhiteBranch a] -> ShowS)
-> Show (WhiteBranch a)
forall a.
(BinaryTreeNode a, Show a) =>
Int -> WhiteBranch a -> ShowS
forall a. (BinaryTreeNode a, Show a) => [WhiteBranch a] -> ShowS
forall a. (BinaryTreeNode a, Show a) => WhiteBranch a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a.
(BinaryTreeNode a, Show a) =>
Int -> WhiteBranch a -> ShowS
showsPrec :: Int -> WhiteBranch a -> ShowS
$cshow :: forall a. (BinaryTreeNode a, Show a) => WhiteBranch a -> String
show :: WhiteBranch a -> String
$cshowList :: forall a. (BinaryTreeNode a, Show a) => [WhiteBranch a] -> ShowS
showList :: [WhiteBranch a] -> ShowS
Show)


isColor :: (BinaryTreeNode a) => RedBlackNode a -> RedBlack -> Bool
isColor :: forall a. BinaryTreeNode a => RedBlackNode a -> RedBlack -> Bool
isColor (RedBlackNode RedBlack
color a
_) RedBlack
expectedColor = RedBlack
color RedBlack -> RedBlack -> Bool
forall a. Eq a => a -> a -> Bool
== RedBlack
expectedColor

branchIsColor :: (BinaryTreeNode a) => TreeBranch (RedBlackNode a) -> RedBlack
  -> Bool
branchIsColor :: forall a.
BinaryTreeNode a =>
TreeBranch (RedBlackNode a) -> RedBlack -> Bool
branchIsColor (TreeBranch BinaryTree (RedBlackNode a)
leftChild RedBlackNode a
node BinaryTree (RedBlackNode a)
rightChild) = RedBlackNode a -> RedBlack -> Bool
forall a. BinaryTreeNode a => RedBlackNode a -> RedBlack -> Bool
isColor RedBlackNode a
node

treeIsColor :: (BinaryTreeNode a) => RedBlackTree a -> RedBlack -> Bool
treeIsColor :: forall a. BinaryTreeNode a => RedBlackTree a -> RedBlack -> Bool
treeIsColor BinaryTree (RedBlackNode a)
Leaf RedBlack
expectedColor = RedBlack
expectedColor RedBlack -> RedBlack -> Bool
forall a. Eq a => a -> a -> Bool
== RedBlack
Black
treeIsColor (Branch BinaryTree (RedBlackNode a)
leftChild RedBlackNode a
node BinaryTree (RedBlackNode a)
rightChild) RedBlack
expectedColor =
  RedBlackNode a -> RedBlack -> Bool
forall a. BinaryTreeNode a => RedBlackNode a -> RedBlack -> Bool
isColor RedBlackNode a
node RedBlack
expectedColor

paintItBlack :: (BinaryTreeNode a) => RedBlackNode a -> RedBlackNode a
paintItBlack :: forall a. BinaryTreeNode a => RedBlackNode a -> RedBlackNode a
paintItBlack (RedBlackNode RedBlack
_ a
item) = RedBlack -> a -> RedBlackNode a
forall a. RedBlack -> a -> RedBlackNode a
RedBlackNode RedBlack
Black a
item

removeBranchColor :: (BinaryTreeNode a) => RedBlackBranch a -> WhiteBranch a
removeBranchColor :: forall a. BinaryTreeNode a => RedBlackBranch a -> WhiteBranch a
removeBranchColor (TreeBranch BinaryTree (RedBlackNode a)
leftChild (RedBlackNode RedBlack
_ a
content) BinaryTree (RedBlackNode a)
rightChild) =
  BinaryTree (RedBlackNode a)
-> a -> BinaryTree (RedBlackNode a) -> WhiteBranch a
forall a. RedBlackTree a -> a -> RedBlackTree a -> WhiteBranch a
WhiteBranch BinaryTree (RedBlackNode a)
leftChild a
content BinaryTree (RedBlackNode a)
rightChild

whiteBranch2Tree :: (BinaryTreeNode a) => WhiteBranch a -> RedBlack ->
  RedBlackTree a
whiteBranch2Tree :: forall a.
BinaryTreeNode a =>
WhiteBranch a -> RedBlack -> RedBlackTree a
whiteBranch2Tree (WhiteBranch RedBlackTree a
leftChild a
content RedBlackTree a
rightChild) RedBlack
color =
  RedBlackTree a
-> RedBlackNode a -> RedBlackTree a -> RedBlackTree a
forall a. BinaryTree a -> a -> BinaryTree a -> BinaryTree a
Branch RedBlackTree a
leftChild RedBlackNode a
newNode RedBlackTree a
rightChild
  where newNode :: RedBlackNode a
newNode = RedBlack -> a -> RedBlackNode a
forall a. RedBlack -> a -> RedBlackNode a
RedBlackNode RedBlack
color a
content

-- | Returns the black-height of a "RedBlackTree", the uniform number of black
-- nodes in any path from the root to any leaf at the bottom of the tree. This
-- is a O(Log(n)) operation.
getBlackHeight :: (BinaryTreeNode a) => RedBlackTree a -> Int
getBlackHeight :: forall a. BinaryTreeNode a => RedBlackTree a -> Int
getBlackHeight BinaryTree (RedBlackNode a)
Leaf = Int
1
getBlackHeight (Branch BinaryTree (RedBlackNode a)
_ (RedBlackNode RedBlack
Black a
_) BinaryTree (RedBlackNode a)
rightSubtree) =
  Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ BinaryTree (RedBlackNode a) -> Int
forall a. BinaryTreeNode a => RedBlackTree a -> Int
getBlackHeight BinaryTree (RedBlackNode a)
rightSubtree
getBlackHeight (Branch BinaryTree (RedBlackNode a)
_ (RedBlackNode RedBlack
Red a
_) BinaryTree (RedBlackNode a)
rightSubtree) =
  BinaryTree (RedBlackNode a) -> Int
forall a. BinaryTreeNode a => RedBlackTree a -> Int
getBlackHeight BinaryTree (RedBlackNode a)
rightSubtree

getNodeContent :: (BinaryTreeNode a) => RedBlackNode a -> a
getNodeContent :: forall a. BinaryTreeNode a => RedBlackNode a -> a
getNodeContent (RedBlackNode RedBlack
_ a
content) = a
content

-- | Lookup a target node in the tree. The target value doesn't need to be the
-- exact same value that is already in the tree. It only needs to satisfy the
-- "Eq" instance
find :: (BinaryTreeNode a) => RedBlackTree a -> a -> Maybe a
find :: forall a. BinaryTreeNode a => RedBlackTree a -> a -> Maybe a
find RedBlackTree a
redBlackTree a
target = (RedBlackNode a -> a) -> Maybe (RedBlackNode a) -> Maybe a
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap RedBlackNode a -> a
forall a. BinaryTreeNode a => RedBlackNode a -> a
getNodeContent Maybe (RedBlackNode a)
maybeResult
  where maybeResult :: Maybe (RedBlackNode a)
maybeResult = RedBlackTree a -> RedBlackNode a -> Maybe (RedBlackNode a)
forall a. BinaryTreeNode a => BinaryTree a -> a -> Maybe a
binaryTreeFind RedBlackTree a
redBlackTree (RedBlack -> a -> RedBlackNode a
forall a. RedBlack -> a -> RedBlackNode a
RedBlackNode RedBlack
Black a
target)

-- | Convenient function to "create" a new empty tree.
emptyRedBlackTree :: RedBlackTree a
emptyRedBlackTree :: forall a. RedBlackTree a
emptyRedBlackTree = BinaryTree (RedBlackNode a)
forall a. BinaryTree a
Leaf