module Data.Tree.Braun
(
fromList
,replicate
,singleton
,empty
,Builder
,consB
,nilB
,runB
,
cons
,uncons
,uncons'
,tail
,
foldrBraun
,toList
,
(!)
,(!?)
,size
,UpperBound(..)
,ub)
where
import Data.Tree.Binary (Tree (..))
import qualified Data.Tree.Binary as Binary
import GHC.Base (build)
import Prelude hiding (tail, replicate)
import Data.Tree.Braun.Internal (zipLevels)
import GHC.Stack
singleton :: a -> Tree a
singleton = Binary.singleton
empty :: Tree a
empty = Leaf
fromList :: [a] -> Tree a
fromList xs = runB (foldr consB nilB xs)
type Builder a b = (Int -> Int -> (([Tree a] -> [Tree a] -> [Tree a]) -> [Tree a] -> b) -> b)
consB :: a -> Builder a b -> Builder a b
consB e a !k 1 p = a (k*2) k (\ys zs -> p (\_ _ -> []) (zipLevels e ys zs (drop k zs)))
consB e a !k !m p = a k (m1) (p . zipLevels e)
nilB :: Builder a b
nilB _ _ p = p (\_ _ -> []) [Leaf]
runB :: Builder a (Tree a) -> Tree a
runB b = b 1 1 (const head)
foldrBraun :: Tree a -> (a -> b -> b) -> b -> b
foldrBraun tr c n =
case tr of
Leaf -> n
_ -> tol [tr]
where tol [] = n
tol xs = foldr (c . root) (tol (children xs id)) xs
children [] k = k []
children (Node _ Leaf _:_) k = k []
children (Node _ l Leaf:ts) k =
l : foldr leftChildren (k []) ts
children (Node _ l r:ts) k = l : children ts (k . (:) r)
children _ _ =
errorWithoutStackTrace "Data.Tree.Braun.toList: bug!"
leftChildren (Node _ Leaf _) _ = []
leftChildren (Node _ l _) a = l : a
leftChildren _ _ =
errorWithoutStackTrace "Data.Tree.Braun.toList: bug!"
root (Node x _ _) = x
root _ = errorWithoutStackTrace "Data.Tree.Braun.toList: bug!"
toList :: Tree a -> [a]
toList tr = build (foldrBraun tr)
size :: Tree a -> Int
size Leaf = 0
size (Node _ l r) = 1 + 2 * m + diff l m where
m = size r
diff Leaf 0 = 0
diff (Node _ Leaf Leaf) 0 = 1
diff (Node _ s t) k
| odd k = diff s (k `div` 2)
| otherwise = diff t ((k `div` 2) 1)
diff Leaf _ = errorWithoutStackTrace "Data.Tree.Braun.size: bug!"
replicate :: Int -> a -> Tree a
replicate m x = go m (const id)
where
go 0 k = k (Node x Leaf Leaf) Leaf
go n k
| odd n = go (pred n `div` 2) $ \s t -> k (Node x s t) (Node x t t)
| otherwise = go (pred n `div` 2) $ \s t -> k (Node x s s) (Node x s t)
(!) :: HasCallStack => Tree a -> Int -> a
(!) (Node x _ _) 0 = x
(!) (Node _ y z) i
| odd i = y ! j
| otherwise = z ! j
where j = (i1) `div` 2
(!) _ _ = error "Data.Tree.Braun.!: index out of range"
(!?) :: Tree a -> Int -> Maybe a
(!?) (Node x _ _) 0 = Just x
(!?) (Node _ y z) i
| odd i = y !? j
| otherwise = z !? j
where j = (i1) `div` 2
(!?) _ _ = Nothing
data UpperBound a = Exact a
| TooHigh Int
| Finite
ub :: (a -> b -> Ordering) -> a -> Tree b -> UpperBound b
ub f x t = go f x t 0 1
where
go _ _ Leaf !_ !_ = Finite
go _ _ (Node hd _ ev) !n !k =
case f x hd of
LT -> TooHigh n
EQ -> Exact hd
GT -> go f x ev (n+2*k) (2*k)
uncons :: Tree a -> Maybe (a, Tree a)
uncons (Node x Leaf Leaf) = Just (x, Leaf)
uncons (Node x y z) = Just (x, Node lp z q)
where
Just (lp,q) = uncons y
uncons Leaf = Nothing
uncons' :: HasCallStack => Tree a -> (a, Tree a)
uncons' (Node x Leaf Leaf) = (x, Leaf)
uncons' (Node x y z) = (x, Node lp z q)
where
(lp,q) = uncons' y
uncons' Leaf = error "Data.Tree.Braun.uncons': empty tree"
cons :: a -> Tree a -> Tree a
cons x Leaf = Node x Leaf Leaf
cons x (Node y p q) = Node x (cons y q) p
tail :: Tree a -> Tree a
tail (Node _ Leaf Leaf) = Leaf
tail (Node _ y z) = Node lp z q
where
(lp,q) = uncons' y
tail Leaf = Leaf