{-# LANGUAGE PatternSynonyms #-}

{- |
This module contains some debug utilities. It should only be used for debugging/testing purposes.
-}

module Data.RRBVector.Internal.Debug
    ( showTree
    , fromListUnbalanced
    , pattern Empty, pattern Root
    , Tree, Shift
    , pattern Balanced, pattern Unbalanced, pattern Leaf
    , Invariant, valid
    ) where

import Control.Monad.ST (runST)
import Data.Bits (shiftL)
import Data.Foldable (foldl', toList, traverse_)
import Data.List (intercalate)
import Data.Primitive.PrimArray (PrimArray, primArrayToList, indexPrimArray, sizeofPrimArray)

import Data.RRBVector.Internal hiding (Empty, Root, Balanced, Unbalanced, Leaf)
import qualified Data.RRBVector.Internal as RRB
import Data.RRBVector.Internal.Array (Array)
import qualified Data.RRBVector.Internal.Array as A
import qualified Data.RRBVector.Internal.Buffer as Buffer

-- | \(O(n)\). Show the underlying tree of a vector.
showTree :: (Show a) => Vector a -> String
showTree :: forall a. Show a => Vector a -> String
showTree Vector a
Empty = String
"Empty"
showTree (Root Int
size Int
sh Tree a
tree) = String
"Root {size = " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
size String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
", shift = " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
sh String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
", tree = " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Tree a -> String
forall {a}. Show a => Tree a -> String
debugShowTree Tree a
tree String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"}"
  where
    debugShowTree :: Tree a -> String
debugShowTree (Balanced Array (Tree a)
arr) = String
"Balanced " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Array (Tree a) -> String
debugShowArray Array (Tree a)
arr
    debugShowTree (Unbalanced Array (Tree a)
arr PrimArray Int
sizes) = String
"Unbalanced " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Array (Tree a) -> String
debugShowArray Array (Tree a)
arr String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" (" String -> String -> String
forall a. [a] -> [a] -> [a]
++ [Int] -> String
forall a. Show a => a -> String
show (PrimArray Int -> [Int]
forall a. Prim a => PrimArray a -> [a]
primArrayToList PrimArray Int
sizes) String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
")"
    debugShowTree (Leaf Array a
arr) = String
"Leaf " String -> String -> String
forall a. [a] -> [a] -> [a]
++ [a] -> String
forall a. Show a => a -> String
show (Array a -> [a]
forall a. Array a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Array a
arr)

    debugShowArray :: Array (Tree a) -> String
debugShowArray Array (Tree a)
arr = String
"[" String -> String -> String
forall a. [a] -> [a] -> [a]
++ String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
"," ((Tree a -> String) -> [Tree a] -> [String]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Tree a -> String
debugShowTree (Array (Tree a) -> [Tree a]
forall a. Array a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Array (Tree a)
arr)) String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"]"

-- | \(O(n)\). Create a new unbalanced vector from a list.
--
-- Note that it is not possbible to create an invalid 'Vector' with this function.
fromListUnbalanced :: [a] -> Vector a
fromListUnbalanced :: forall a. [a] -> Vector a
fromListUnbalanced [] = Vector a
forall a. Vector a
RRB.Empty
fromListUnbalanced [a
x] = a -> Vector a
forall a. a -> Vector a
singleton a
x
fromListUnbalanced [a]
ls = case (Array a -> Tree a) -> [a] -> [Tree a]
forall {a} {a}. (Array a -> a) -> [a] -> [a]
nodes Array a -> Tree a
forall a. Array a -> Tree a
RRB.Leaf [a]
ls of
    [Tree a
tree] -> Int -> Int -> Tree a -> Vector a
forall a. Int -> Int -> Tree a -> Vector a
RRB.Root (Int -> Tree a -> Int
forall a. Int -> Tree a -> Int
treeSize Int
0 Tree a
tree) Int
0 Tree a
tree -- tree is a single leaf
    [Tree a]
ls' -> Int -> [Tree a] -> Vector a
forall {a}. Int -> [Tree a] -> Vector a
iterateNodes Int
blockShift [Tree a]
ls'
  where
    n :: Int
n = Int
blockSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1

    nodes :: (Array a -> a) -> [a] -> [a]
nodes Array a -> a
f [a]
trees = (forall s. ST s [a]) -> [a]
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s [a]) -> [a]) -> (forall s. ST s [a]) -> [a]
forall a b. (a -> b) -> a -> b
$ do
        Buffer s a
buffer <- Int -> ST s (Buffer s a)
forall s a. Int -> ST s (Buffer s a)
Buffer.new Int
n
        let loop :: [a] -> ST s [a]
loop [] = do
                Array a
result <- Buffer s a -> ST s (Array a)
forall s a. Buffer s a -> ST s (Array a)
Buffer.get Buffer s a
buffer
                [a] -> ST s [a]
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Array a -> a
f Array a
result]
            loop (a
t : [a]
ts) = do
                Int
size <- Buffer s a -> ST s Int
forall s a. Buffer s a -> ST s Int
Buffer.size Buffer s a
buffer
                if Int
size Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n then do
                    Array a
result <- Buffer s a -> ST s (Array a)
forall s a. Buffer s a -> ST s (Array a)
Buffer.get Buffer s a
buffer
                    Buffer s a -> a -> ST s ()
forall s a. Buffer s a -> a -> ST s ()
Buffer.push Buffer s a
buffer a
t
                    [a]
rest <- [a] -> ST s [a]
loop [a]
ts
                    [a] -> ST s [a]
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array a -> a
f Array a
result a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
rest)
                else do
                    Buffer s a -> a -> ST s ()
forall s a. Buffer s a -> a -> ST s ()
Buffer.push Buffer s a
buffer a
t
                    [a] -> ST s [a]
loop [a]
ts
        [a] -> ST s [a]
loop [a]
trees
    {-# INLINE nodes #-}

    iterateNodes :: Int -> [Tree a] -> Vector a
iterateNodes Int
sh [Tree a]
trees = case (Array (Tree a) -> Tree a) -> [Tree a] -> [Tree a]
forall {a} {a}. (Array a -> a) -> [a] -> [a]
nodes (Int -> Array (Tree a) -> Tree a
forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
sh) [Tree a]
trees of
        [Tree a
tree] -> Int -> Int -> Tree a -> Vector a
forall a. Int -> Int -> Tree a -> Vector a
RRB.Root (Int -> Tree a -> Int
forall a. Int -> Tree a -> Int
treeSize Int
sh Tree a
tree) Int
sh Tree a
tree
        [Tree a]
trees' -> Int -> [Tree a] -> Vector a
iterateNodes (Int -> Int
up Int
sh) [Tree a]
trees'

pattern Empty :: Vector a
pattern $mEmpty :: forall {r} {a}. Vector a -> ((# #) -> r) -> ((# #) -> r) -> r
Empty <- RRB.Empty

pattern Root :: Int -> Shift -> Tree a -> Vector a
pattern $mRoot :: forall {r} {a}.
Vector a -> (Int -> Int -> Tree a -> r) -> ((# #) -> r) -> r
Root size sh tree <- RRB.Root size sh tree

{-# COMPLETE Empty, Root #-}

pattern Balanced :: Array (Tree a) -> Tree a
pattern $mBalanced :: forall {r} {a}.
Tree a -> (Array (Tree a) -> r) -> ((# #) -> r) -> r
Balanced arr <- RRB.Balanced arr

pattern Unbalanced :: Array (Tree a) -> PrimArray Int -> Tree a
pattern $mUnbalanced :: forall {r} {a}.
Tree a
-> (Array (Tree a) -> PrimArray Int -> r) -> ((# #) -> r) -> r
Unbalanced arr sizes <- RRB.Unbalanced arr sizes

pattern Leaf :: Array a -> Tree a
pattern $mLeaf :: forall {r} {a}. Tree a -> (Array a -> r) -> ((# #) -> r) -> r
Leaf arr <- RRB.Leaf arr

{-# COMPLETE Balanced, Unbalanced, Leaf #-}

-- | Structural invariants a vector is expected to hold.
data Invariant
    = RootSizeGt0      -- Root: Size > 0
    | RootShiftDiv     -- Root: The shift at the root is divisible by blockShift
    | RootSizeCorrect  -- Root: The size at the root is correct
    | RootGt1Child     -- Root: The root has more than 1 child if not a Leaf
    | BalShiftGt0      -- Balanced: Shift > 0
    | BalNumChildren   -- Balanced: The number of children is blockSize unless
                       -- the parent is unbalanced or the node is on the right
                       -- edge in which case it is in [1,blockSize]
    | BalFullChildren  -- Balanced: All children are full, except for the last
                       -- if the node is on the right edge
    | UnbalShiftGt0    -- Unbalanced: Shift > 0
    | UnbalParentUnbal -- Unbalanced: Parent is Unbalanced
    | UnbalNumChildren -- Unbalanced: The number of children is in [1,blockSize]
    | UnbalSizes       -- Unbalanced: The sizes array is correct
    | UnbalNotBal      -- Unbalanced: The tree is not full enough to be a
                       -- Balanced
    | LeafShift0       -- Leaf: Shift == 0
    | LeafNumElems     -- Leaf: The number of elements is in [1,blockSize]
    deriving Int -> Invariant -> String -> String
[Invariant] -> String -> String
Invariant -> String
(Int -> Invariant -> String -> String)
-> (Invariant -> String)
-> ([Invariant] -> String -> String)
-> Show Invariant
forall a.
(Int -> a -> String -> String)
-> (a -> String) -> ([a] -> String -> String) -> Show a
$cshowsPrec :: Int -> Invariant -> String -> String
showsPrec :: Int -> Invariant -> String -> String
$cshow :: Invariant -> String
show :: Invariant -> String
$cshowList :: [Invariant] -> String -> String
showList :: [Invariant] -> String -> String
Show

assert :: Invariant -> Bool -> Either Invariant ()
assert :: Invariant -> Bool -> Either Invariant ()
assert Invariant
i Bool
False = Invariant -> Either Invariant ()
forall a b. a -> Either a b
Left Invariant
i
assert Invariant
_ Bool
True = () -> Either Invariant ()
forall a. a -> Either Invariant a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | Check tree invariants. Returns @Left@ on finding a violated invariant.
valid :: Vector a -> Either Invariant ()
valid :: forall a. Vector a -> Either Invariant ()
valid Vector a
RRB.Empty = () -> Either Invariant ()
forall a. a -> Either Invariant a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
valid (RRB.Root Int
size Int
sh Tree a
tree) = do
    Invariant -> Bool -> Either Invariant ()
assert Invariant
RootSizeGt0 (Bool -> Either Invariant ()) -> Bool -> Either Invariant ()
forall a b. (a -> b) -> a -> b
$ Int
size Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
    Invariant -> Bool -> Either Invariant ()
assert Invariant
RootShiftDiv (Bool -> Either Invariant ()) -> Bool -> Either Invariant ()
forall a b. (a -> b) -> a -> b
$ Int
sh Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
blockShift Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
    Invariant -> Bool -> Either Invariant ()
assert Invariant
RootSizeCorrect (Bool -> Either Invariant ()) -> Bool -> Either Invariant ()
forall a b. (a -> b) -> a -> b
$ Int
size Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Tree a -> Int
forall a. Tree a -> Int
countElems Tree a
tree
    Invariant -> Bool -> Either Invariant ()
assert Invariant
RootGt1Child (Bool -> Either Invariant ()) -> Bool -> Either Invariant ()
forall a b. (a -> b) -> a -> b
$ case Tree a
tree of
        Balanced Array (Tree a)
arr -> Array (Tree a) -> Int
forall a. Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1
        Unbalanced Array (Tree a)
arr PrimArray Int
_ -> Array (Tree a) -> Int
forall a. Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1
        Leaf Array a
_ -> Bool
True
    NodeDesc -> Int -> Tree a -> Either Invariant ()
forall a. NodeDesc -> Int -> Tree a -> Either Invariant ()
validTree NodeDesc
Unbal Int
sh Tree a
tree

data NodeDesc
    = Bal           -- parent is Balanced
    | BalRightEdge  -- parent is Balanced and this node is on the right edge
    | Unbal         -- parent is Unbalanced

validTree :: NodeDesc -> Shift -> Tree a -> Either Invariant ()
validTree :: forall a. NodeDesc -> Int -> Tree a -> Either Invariant ()
validTree NodeDesc
desc Int
sh (RRB.Balanced Array (Tree a)
arr) = do
    Invariant -> Bool -> Either Invariant ()
assert Invariant
BalShiftGt0 (Bool -> Either Invariant ()) -> Bool -> Either Invariant ()
forall a b. (a -> b) -> a -> b
$ Int
sh Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
    Invariant -> Bool -> Either Invariant ()
assert Invariant
BalNumChildren (Bool -> Either Invariant ()) -> Bool -> Either Invariant ()
forall a b. (a -> b) -> a -> b
$ case NodeDesc
desc of
        NodeDesc
Bal -> Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
blockSize
        NodeDesc
BalRightEdge -> Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1 Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
blockSize
        NodeDesc
Unbal -> Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1 Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
blockSize
    Invariant -> Bool -> Either Invariant ()
assert Invariant
BalFullChildren (Bool -> Either Invariant ()) -> Bool -> Either Invariant ()
forall a b. (a -> b) -> a -> b
$
        (Tree a -> Bool) -> Array (Tree a) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\Tree a
t -> Tree a -> Int
forall a. Tree a -> Int
countElems Tree a
t Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
sh) Array (Tree a)
expectedFullChildren
    (Tree a -> Either Invariant ())
-> Array (Tree a) -> Either Invariant ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (NodeDesc -> Int -> Tree a -> Either Invariant ()
forall a. NodeDesc -> Int -> Tree a -> Either Invariant ()
validTree NodeDesc
Bal (Int -> Int
down Int
sh)) Array (Tree a)
arrInit
    NodeDesc -> Int -> Tree a -> Either Invariant ()
forall a. NodeDesc -> Int -> Tree a -> Either Invariant ()
validTree NodeDesc
descLast (Int -> Int
down Int
sh) (Array (Tree a) -> Tree a
forall a. Array a -> a
A.last Array (Tree a)
arr)
  where
    n :: Int
n = Array (Tree a) -> Int
forall a. Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr
    arrInit :: Array (Tree a)
arrInit = Array (Tree a) -> Int -> Array (Tree a)
forall a. Array a -> Int -> Array a
A.take Array (Tree a)
arr (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
    expectedFullChildren :: Array (Tree a)
expectedFullChildren = case NodeDesc
desc of
        NodeDesc
Bal -> Array (Tree a)
arr
        NodeDesc
BalRightEdge -> Array (Tree a)
arrInit
        NodeDesc
Unbal -> Array (Tree a)
arrInit
    descLast :: NodeDesc
descLast = case NodeDesc
desc of
        NodeDesc
Bal -> NodeDesc
Bal
        NodeDesc
BalRightEdge -> NodeDesc
BalRightEdge
        NodeDesc
Unbal -> NodeDesc
BalRightEdge
validTree NodeDesc
desc Int
sh (RRB.Unbalanced Array (Tree a)
arr PrimArray Int
sizes) = do
    Invariant -> Bool -> Either Invariant ()
assert Invariant
UnbalShiftGt0 (Bool -> Either Invariant ()) -> Bool -> Either Invariant ()
forall a b. (a -> b) -> a -> b
$ Int
sh Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
    case NodeDesc
desc of
        NodeDesc
Bal -> Invariant -> Bool -> Either Invariant ()
assert Invariant
UnbalParentUnbal Bool
False
        NodeDesc
BalRightEdge -> Invariant -> Bool -> Either Invariant ()
assert Invariant
UnbalParentUnbal Bool
False
        NodeDesc
Unbal -> Invariant -> Bool -> Either Invariant ()
assert Invariant
UnbalNumChildren (Bool -> Either Invariant ()) -> Bool -> Either Invariant ()
forall a b. (a -> b) -> a -> b
$ Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1 Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
blockSize
    Invariant -> Bool -> Either Invariant ()
assert Invariant
UnbalSizes (Bool -> Either Invariant ()) -> Bool -> Either Invariant ()
forall a b. (a -> b) -> a -> b
$ Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== PrimArray Int -> Int
forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray Int
sizes
    Invariant -> Bool -> Either Invariant ()
assert Invariant
UnbalSizes (Bool -> Either Invariant ()) -> Bool -> Either Invariant ()
forall a b. (a -> b) -> a -> b
$
        (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\Int
i -> Tree a -> Int
forall a. Tree a -> Int
countElems (Array (Tree a) -> Int -> Tree a
forall a. Array a -> Int -> a
A.index Array (Tree a)
arr Int
i) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== PrimArray Int -> Int -> Int
getSize PrimArray Int
sizes Int
i) [Int
0 .. Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1]
    Invariant -> Bool -> Either Invariant ()
assert Invariant
UnbalNotBal (Bool -> Either Invariant ()) -> Bool -> Either Invariant ()
forall a b. (a -> b) -> a -> b
$ Bool -> Bool
not (Int -> Array (Tree a) -> PrimArray Int -> Bool
forall a. Int -> Array (Tree a) -> PrimArray Int -> Bool
couldBeBalanced Int
sh Array (Tree a)
arr PrimArray Int
sizes)
    (Tree a -> Either Invariant ())
-> Array (Tree a) -> Either Invariant ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (NodeDesc -> Int -> Tree a -> Either Invariant ()
forall a. NodeDesc -> Int -> Tree a -> Either Invariant ()
validTree NodeDesc
Unbal (Int -> Int
down Int
sh)) Array (Tree a)
arr
  where
    n :: Int
n = Array (Tree a) -> Int
forall a. Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr
validTree NodeDesc
desc Int
sh (RRB.Leaf Array a
arr) = do
    Invariant -> Bool -> Either Invariant ()
assert Invariant
LeafShift0 (Bool -> Either Invariant ()) -> Bool -> Either Invariant ()
forall a b. (a -> b) -> a -> b
$ Int
sh Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
    Invariant -> Bool -> Either Invariant ()
assert Invariant
LeafNumElems (Bool -> Either Invariant ()) -> Bool -> Either Invariant ()
forall a b. (a -> b) -> a -> b
$ case NodeDesc
desc of
        NodeDesc
Bal -> Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
blockSize
        NodeDesc
BalRightEdge -> Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1 Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
blockSize
        NodeDesc
Unbal -> Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1 Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
blockSize
  where
    n :: Int
n = Array a -> Int
forall a. Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array a
arr

-- | Check whether an Unbalanced node could be Balanced.
couldBeBalanced :: Shift -> A.Array (Tree a) -> PrimArray Int -> Bool
couldBeBalanced :: forall a. Int -> Array (Tree a) -> PrimArray Int -> Bool
couldBeBalanced Int
sh Array (Tree a)
arr PrimArray Int
sizes =
   (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\Int
i -> PrimArray Int -> Int -> Int
getSize PrimArray Int
sizes Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
sh) [Int
0 .. Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
2] Bool -> Bool -> Bool
&&
   (case Array (Tree a) -> Tree a
forall a. Array a -> a
A.last Array (Tree a)
arr of
       Balanced Array (Tree a)
_ -> Bool
True
       Unbalanced Array (Tree a)
arr' PrimArray Int
sizes' -> Int -> Array (Tree a) -> PrimArray Int -> Bool
forall a. Int -> Array (Tree a) -> PrimArray Int -> Bool
couldBeBalanced (Int -> Int
down Int
sh) Array (Tree a)
arr' PrimArray Int
sizes'
       Leaf Array a
_ -> Bool
True)
  where
    n :: Int
n = Array (Tree a) -> Int
forall a. Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr

getSize :: PrimArray Int -> Int -> Int
getSize :: PrimArray Int -> Int -> Int
getSize PrimArray Int
sizes Int
0 = PrimArray Int -> Int -> Int
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
sizes Int
0
getSize PrimArray Int
sizes Int
i = PrimArray Int -> Int -> Int
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
sizes Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- PrimArray Int -> Int -> Int
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
sizes (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)

countElems :: Tree a -> Int
countElems :: forall a. Tree a -> Int
countElems (RRB.Balanced Array (Tree a)
arr) =
    (Int -> Tree a -> Int) -> Int -> Array (Tree a) -> Int
forall b a. (b -> a -> b) -> b -> Array a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Int
acc Tree a
tree -> Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Tree a -> Int
forall a. Tree a -> Int
countElems Tree a
tree) Int
0 Array (Tree a)
arr
countElems (RRB.Unbalanced Array (Tree a)
arr PrimArray Int
_) =
    (Int -> Tree a -> Int) -> Int -> Array (Tree a) -> Int
forall b a. (b -> a -> b) -> b -> Array a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Int
acc Tree a
tree -> Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Tree a -> Int
forall a. Tree a -> Int
countElems Tree a
tree) Int
0 Array (Tree a)
arr
countElems (RRB.Leaf Array a
arr) = Array a -> Int
forall a. Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array a
arr