module Data.Tree.Braun.Sized
(
Braun(..)
,fromList
,empty
,singleton
,Builder
,consB
,nilB
,runB
,snoc
,unsnoc
,unsnoc'
,cons
,uncons
,uncons'
,insertBy
,deleteBy
,glb
,cmpRoot
,ltRoot
)
where
import Data.Tree.Binary (Tree (..))
import Data.Tree.Braun (UpperBound (..))
import qualified Data.Tree.Braun as Unsized
import Data.Tree.Braun.Internal (zipLevels)
import Control.DeepSeq (NFData (rnf))
import Data.Data (Data)
import Data.Typeable (Typeable)
import GHC.Generics (Generic, Generic1)
import Control.Applicative hiding (empty)
import GHC.Stack
import Data.Foldable
data Braun a = Braun
{ size :: !Int
, tree :: Tree a
} deriving (Show,Read,Eq,Ord,Functor,Typeable,Generic,Generic1
,Data)
instance NFData a => NFData (Braun a) where
rnf (Braun _ tr) = rnf tr
instance Foldable Braun where
foldr f b (Braun _ xs) = Unsized.foldrBraun xs f b
length = size
toList (Braun _ xs) = Unsized.toList xs
instance Traversable Braun where
traverse f (Braun n tr) = fmap k (Unsized.foldrBraun tr c b)
where
c = liftA2 Unsized.consB . f
b = pure Unsized.nilB
k = Braun n . Unsized.runB
snoc :: a -> Braun a -> Braun a
snoc x (Braun 0 Leaf) = Braun 1 (Node x Leaf Leaf)
snoc x (Braun n (Node y z w))
| even n = Braun (n + 1) (Node y z (tree (snoc x (Braun (m 1) w))))
| otherwise = Braun (n + 1) (Node y (tree (snoc x (Braun m z))) w)
where
m = n `div` 2
snoc _ (Braun _ Leaf) = errorWithoutStackTrace "Data.Tree.Braun.Sized.snoc: bug!"
type Builder a b = (Int -> Int -> Int -> (([Tree a] -> [Tree a] -> [Tree a]) -> [Tree a] -> Int -> b) -> b)
consB :: a -> Builder a b -> Builder a b
consB e a !k 1 !s p = a (k*2) k (s+1) (\ys zs -> p (\_ _ -> []) (zipLevels e ys zs (drop k zs)))
consB e a !k !m !s p = a k (m1) (s+1) (p . zipLevels e)
nilB :: Builder a b
nilB _ _ s p = p (\_ _ -> []) [Leaf] s
runB :: Builder a (Braun a) -> Braun a
runB xs = xs 1 1 0 (const (flip Braun . head))
fromList :: [a] -> Braun a
fromList xs = runB (foldr consB nilB xs)
empty :: Braun a
empty = Braun 0 Leaf
singleton :: a -> Braun a
singleton x = Braun 1 (Node x Leaf Leaf)
insertBy :: (a -> a -> Ordering) -> a -> Braun a -> Braun a
insertBy cmp x b@(Braun s xs) =
case break
(\y ->
cmp x y /= GT)
(Unsized.toList xs) of
(_,[]) -> snoc x b
(lt,gte@(y:_)) ->
if cmp x y == EQ
then b
else Braun
(s + 1)
(Unsized.runB
(foldr
Unsized.consB
(Unsized.consB
x
(foldr Unsized.consB Unsized.nilB gte))
lt))
deleteBy :: (a -> a -> Ordering) -> a -> Braun a -> Braun a
deleteBy cmp x b@(Braun s xs) =
case break
(\y -> cmp x y /= GT)
(Unsized.toList xs) of
(_,[]) -> b
(lt,y:gt) ->
if cmp x y /= EQ
then b
else Braun
(s 1)
(Unsized.runB (foldr Unsized.consB (foldr Unsized.consB Unsized.nilB gt) lt))
glb :: (a -> b -> Ordering) -> a -> Braun b -> Maybe b
glb _ _ (Braun _ Leaf) = Nothing
glb cmp x (Braun n ys@(Node h _ _)) =
case cmp x h of
LT -> Nothing
EQ -> Just h
GT ->
case Unsized.ub cmp x ys of
Exact ans -> Just ans
Finite
| cmp x final == LT -> go 0 (n 1)
| otherwise -> Just final
where final = ys Unsized.! (n 1)
TooHigh m -> go 0 m
where
go _ 0 = Nothing
go i j
| j <= i = Just $ ys Unsized.! (j 1)
| i + 1 == j = Just $ ys Unsized.! i
| otherwise =
case cmp x middle of
LT -> go i k
EQ -> Just middle
GT -> go k j
where
k = (i + j) `div` 2
middle = ys Unsized.! k
cons :: a -> Braun a -> Braun a
cons x (Braun n xs) = Braun (n+1) (Unsized.cons x xs)
uncons :: Braun a -> Maybe (a, Braun a)
uncons (Braun n tr) = (fmap.fmap) (Braun (n1)) (Unsized.uncons tr)
uncons' :: HasCallStack => Braun a -> (a, Braun a)
uncons' (Braun n tr) = fmap (Braun (n1)) (Unsized.uncons' tr)
cmpRoot :: (a -> b -> Ordering) -> a -> Braun b -> Ordering
cmpRoot cmp x (Braun _ (Node y _ _)) = cmp x y
cmpRoot _ _ _ = error "Data.Tree.Braun.Sized.compRoot: empty tree"
ltRoot :: (a -> b -> Ordering) -> a -> Braun b -> Bool
ltRoot cmp x (Braun _ (Node y _ _)) = cmp x y == LT
ltRoot _ _ _ = error "Data.Tree.Braun.Sized.ltRoot: empty tree"
unsnoc :: Braun a -> Maybe (a, Braun a)
unsnoc (Braun _ (Node x Leaf Leaf)) = Just (x, Braun 0 Leaf)
unsnoc (Braun n (Node x y z))
| odd n =
let Just (p,Braun _ q) = unsnoc (Braun m z)
in Just (p, Braun (n 1) (Node x y q))
| otherwise =
let Just (p,Braun _ q) = unsnoc (Braun m y)
in Just (p, Braun (n 1) (Node x q z))
where
m = n `div` 2
unsnoc (Braun _ Leaf) = Nothing
unsnoc' :: HasCallStack => Braun a -> (a, Braun a)
unsnoc' (Braun _ (Node x Leaf Leaf)) = (x, Braun 0 Leaf)
unsnoc' (Braun n (Node x y z))
| odd n =
let (p,Braun _ q) = unsnoc' (Braun m z)
in (p, Braun (n 1) (Node x y q))
| otherwise =
let (p,Braun _ q) = unsnoc' (Braun m y)
in (p, Braun (n 1) (Node x q z))
where
m = n `div` 2
unsnoc' (Braun _ Leaf) = error "Data.Tree.Braun.Sized.unsnoc': empty tree"