{-# LANGUAGE BangPatterns         #-}
{-# LANGUAGE LambdaCase           #-}
{-# LANGUAGE PatternSynonyms      #-}
{-# LANGUAGE TypeFamilies         #-}
{-# LANGUAGE TypeOperators        #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns         #-}
module Data.RandomAccessList.SkewBinarySlab
    ( RAList(Cons,Nil)
    , safeIndexZero
    , unsafeIndexZero
    , Data.RandomAccessList.SkewBinarySlab.null
    , uncons
    , consSlab
    ) where

import Data.Bits (unsafeShiftR)
import Data.Vector.NonEmpty qualified as NEV
import Data.Word
import GHC.Exts (IsList, toList)

import Data.RandomAccessList.Class qualified as RAL

{- Note [Skew binary slab lists]
This module implements a very similar structure to the one in 'SkewBinary', but
instead of storing a single value at each node, it instead stores potentially many
values.

The advantage of this is that we can rapidly cons on a collection of values, and that
if we do this regularly then the size of the structure will grow more slowly than
the number of values stored, giving us a discount on our lookup performance (which
depends on the size of the structure!).

The disadvantages are several:
- It's more complex.
- We need another intermediary type, which means more indirect lookups.
- We need to store another size in the spine of the list *and* in the tree nodes,
since a) the structure size no longer tells us the element count, and b) as we
traverse a tree it's no longer true that the size on each side is always half of
the overall size.

Benchmarking suggests that it *is* slightly slower than the normal version
on the non-slab-based workflows, but it's much faster on the slab-based workflows.

So it's not an unqualified win, but it may be better in some cases.
-}

-- Why not just store `NonEmptyVector`s and add singleton values by making singleton
-- vectors? The answer is that using only vectors makes simple consing significantly
-- slower, and doesn't obviously make the other code paths faster.
-- | The values that can be stored in a node. Either a single value, or a non-empty vector of
-- values.
data Values a = One a | Many {-# UNPACK #-} !(NEV.NonEmptyVector a)
    deriving stock (Values a -> Values a -> Bool
(Values a -> Values a -> Bool)
-> (Values a -> Values a -> Bool) -> Eq (Values a)
forall a. Eq a => Values a -> Values a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => Values a -> Values a -> Bool
== :: Values a -> Values a -> Bool
$c/= :: forall a. Eq a => Values a -> Values a -> Bool
/= :: Values a -> Values a -> Bool
Eq, Int -> Values a -> ShowS
[Values a] -> ShowS
Values a -> String
(Int -> Values a -> ShowS)
-> (Values a -> String) -> ([Values a] -> ShowS) -> Show (Values a)
forall a. Show a => Int -> Values a -> ShowS
forall a. Show a => [Values a] -> ShowS
forall a. Show a => Values a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> Values a -> ShowS
showsPrec :: Int -> Values a -> ShowS
$cshow :: forall a. Show a => Values a -> String
show :: Values a -> String
$cshowList :: forall a. Show a => [Values a] -> ShowS
showList :: [Values a] -> ShowS
Show)

valuesCount :: Values a -> Word64
valuesCount :: forall a. Values a -> Word64
valuesCount (One a
_)  = Word64
1
valuesCount (Many NonEmptyVector a
v) = Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word64) -> Int -> Word64
forall a b. (a -> b) -> a -> b
$ NonEmptyVector a -> Int
forall a. NonEmptyVector a -> Int
NEV.length NonEmptyVector a
v

unsafeIndexValues :: Word64 -> Values a -> a
unsafeIndexValues :: forall a. Word64 -> Values a -> a
unsafeIndexValues Word64
0 (One a
a)  = a
a
unsafeIndexValues Word64
_ (One a
_)  = String -> a
forall a. HasCallStack => String -> a
error String
"out of bounds"
unsafeIndexValues Word64
i (Many NonEmptyVector a
v) = NonEmptyVector a
v NonEmptyVector a -> Int -> a
forall a. NonEmptyVector a -> Int -> a
NEV.! Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
i

safeIndexValues :: Word64 -> Values a -> Maybe a
safeIndexValues :: forall a. Word64 -> Values a -> Maybe a
safeIndexValues Word64
0 (One a
a)  = a -> Maybe a
forall a. a -> Maybe a
Just a
a
safeIndexValues Word64
_ (One a
_)  = Maybe a
forall a. Maybe a
Nothing
safeIndexValues Word64
i (Many NonEmptyVector a
v) = NonEmptyVector a
v NonEmptyVector a -> Int -> Maybe a
forall a. NonEmptyVector a -> Int -> Maybe a
NEV.!? Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
i

-- O(1)
unconsValues :: Values a -> RAList a -> (a, RAList a)
unconsValues :: forall a. Values a -> RAList a -> (a, RAList a)
unconsValues (One a
x) RAList a
l = (a
x, RAList a
l)
unconsValues (Many NonEmptyVector a
v) RAList a
l =
    -- unconsing vectors is actually O(1), which is important!
    let (a
x, Vector a
xs) = NonEmptyVector a -> (a, Vector a)
forall a. NonEmptyVector a -> (a, Vector a)
NEV.uncons NonEmptyVector a
v
        remaining :: RAList a
remaining = case Vector a -> Maybe (NonEmptyVector a)
forall a. Vector a -> Maybe (NonEmptyVector a)
NEV.fromVector Vector a
xs of
            Just NonEmptyVector a
v' -> NonEmptyVector a -> RAList a -> RAList a
forall a. NonEmptyVector a -> RAList a -> RAList a
consSlab NonEmptyVector a
v' RAList a
l
            Maybe (NonEmptyVector a)
Nothing -> RAList a
l
    in (a
x, RAList a
remaining)

-- | A complete binary tree.
data Tree a = Leaf !(Values a)
            -- Nodes track the number of elements in the tree (including those in the node)
            | Node {-# UNPACK #-} !Word64 !(Values a) !(Tree a) !(Tree a)
            deriving stock (Tree a -> Tree a -> Bool
(Tree a -> Tree a -> Bool)
-> (Tree a -> Tree a -> Bool) -> Eq (Tree a)
forall a. Eq a => Tree a -> Tree a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => Tree a -> Tree a -> Bool
== :: Tree a -> Tree a -> Bool
$c/= :: forall a. Eq a => Tree a -> Tree a -> Bool
/= :: Tree a -> Tree a -> Bool
Eq, Int -> Tree a -> ShowS
[Tree a] -> ShowS
Tree a -> String
(Int -> Tree a -> ShowS)
-> (Tree a -> String) -> ([Tree a] -> ShowS) -> Show (Tree a)
forall a. Show a => Int -> Tree a -> ShowS
forall a. Show a => [Tree a] -> ShowS
forall a. Show a => Tree a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> Tree a -> ShowS
showsPrec :: Int -> Tree a -> ShowS
$cshow :: forall a. Show a => Tree a -> String
show :: Tree a -> String
$cshowList :: forall a. Show a => [Tree a] -> ShowS
showList :: [Tree a] -> ShowS
Show)

treeCount :: Tree a -> Word64
treeCount :: forall a. Tree a -> Word64
treeCount (Leaf Values a
v)       = Values a -> Word64
forall a. Values a -> Word64
valuesCount Values a
v
treeCount (Node Word64
s Values a
_ Tree a
_ Tree a
_) = Word64
s

unsafeIndexTree :: Word64 -> Tree a -> a
unsafeIndexTree :: forall a. Word64 -> Tree a -> a
unsafeIndexTree Word64
offset (Leaf Values a
v) = Word64 -> Values a -> a
forall a. Word64 -> Values a -> a
unsafeIndexValues Word64
offset Values a
v
unsafeIndexTree Word64
offset (Node Word64
_ Values a
v Tree a
t1 Tree a
t2) =
    let nCount :: Word64
nCount = Values a -> Word64
forall a. Values a -> Word64
valuesCount Values a
v
    in if Word64
offset Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
nCount
    then Word64 -> Values a -> a
forall a. Word64 -> Values a -> a
unsafeIndexValues Word64
offset Values a
v
    else
        let offset' :: Word64
offset' = Word64
offset Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
nCount
            lCount :: Word64
lCount = Tree a -> Word64
forall a. Tree a -> Word64
treeCount Tree a
t1
        in if Word64
offset' Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
lCount
        then Word64 -> Tree a -> a
forall a. Word64 -> Tree a -> a
unsafeIndexTree Word64
offset' Tree a
t1
        else Word64 -> Tree a -> a
forall a. Word64 -> Tree a -> a
unsafeIndexTree (Word64
offset' Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
lCount) Tree a
t2

safeIndexTree :: Word64 -> Tree a -> Maybe a
safeIndexTree :: forall a. Word64 -> Tree a -> Maybe a
safeIndexTree Word64
offset (Leaf Values a
v) = Word64 -> Values a -> Maybe a
forall a. Word64 -> Values a -> Maybe a
safeIndexValues Word64
offset Values a
v
safeIndexTree Word64
offset (Node Word64
_ Values a
v Tree a
t1 Tree a
t2) =
    let nCount :: Word64
nCount = Values a -> Word64
forall a. Values a -> Word64
valuesCount Values a
v
    in if Word64
offset Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
nCount
    then Word64 -> Values a -> Maybe a
forall a. Word64 -> Values a -> Maybe a
safeIndexValues Word64
offset Values a
v
    else
        let offset' :: Word64
offset' = Word64
offset Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
nCount
            lCount :: Word64
lCount = Tree a -> Word64
forall a. Tree a -> Word64
treeCount Tree a
t1
        in if Word64
offset' Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
lCount
        then Word64 -> Tree a -> Maybe a
forall a. Word64 -> Tree a -> Maybe a
safeIndexTree Word64
offset' Tree a
t1
        else Word64 -> Tree a -> Maybe a
forall a. Word64 -> Tree a -> Maybe a
safeIndexTree (Word64
offset' Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
lCount) Tree a
t2

-- | A strict list of complete binary trees accompanied by their node size.
-- The trees appear in >=-node size order.
-- Note: this list is strict in its spine, unlike the Prelude list
data RAList a = BHead
               {-# UNPACK #-} !Word64 -- ^ the number of nodes in the head tree
               !(Tree a) -- ^ the head tree
               !(RAList a) -- ^ the tail trees
             | Nil
             deriving stock (Int -> RAList a -> ShowS
[RAList a] -> ShowS
RAList a -> String
(Int -> RAList a -> ShowS)
-> (RAList a -> String) -> ([RAList a] -> ShowS) -> Show (RAList a)
forall a. Show a => Int -> RAList a -> ShowS
forall a. Show a => [RAList a] -> ShowS
forall a. Show a => RAList a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> RAList a -> ShowS
showsPrec :: Int -> RAList a -> ShowS
$cshow :: forall a. Show a => RAList a -> String
show :: RAList a -> String
$cshowList :: forall a. Show a => [RAList a] -> ShowS
showList :: [RAList a] -> ShowS
Show)
             deriving (Int -> [Item (RAList a)] -> RAList a
[Item (RAList a)] -> RAList a
RAList a -> [Item (RAList a)]
([Item (RAList a)] -> RAList a)
-> (Int -> [Item (RAList a)] -> RAList a)
-> (RAList a -> [Item (RAList a)])
-> IsList (RAList a)
forall a. Int -> [Item (RAList a)] -> RAList a
forall a. [Item (RAList a)] -> RAList a
forall a. RAList a -> [Item (RAList a)]
forall l.
([Item l] -> l)
-> (Int -> [Item l] -> l) -> (l -> [Item l]) -> IsList l
$cfromList :: forall a. [Item (RAList a)] -> RAList a
fromList :: [Item (RAList a)] -> RAList a
$cfromListN :: forall a. Int -> [Item (RAList a)] -> RAList a
fromListN :: Int -> [Item (RAList a)] -> RAList a
$ctoList :: forall a. RAList a -> [Item (RAList a)]
toList :: RAList a -> [Item (RAList a)]
IsList) via RAL.AsRAL (RAList a)

-- Can't use the derived instance because it's no longer the case that lists with
-- the same contents have to have the same structure! Could definitely write a
-- faster implementation if it matters, though.
instance Eq a => Eq (RAList a) where
    RAList a
l == :: RAList a -> RAList a -> Bool
== RAList a
l' = RAList a -> [Item (RAList a)]
forall l. IsList l => l -> [Item l]
toList RAList a
l [a] -> [a] -> Bool
forall a. Eq a => a -> a -> Bool
== RAList a -> [Item (RAList a)]
forall l. IsList l => l -> [Item l]
toList RAList a
l'

{-# INLINABLE null #-}
null :: RAList a -> Bool
null :: forall a. RAList a -> Bool
null RAList a
Nil = Bool
True
null RAList a
_   = Bool
False

{-# complete Cons, Nil #-}
{-# complete BHead, Nil #-}

-- /O(1)/
pattern Cons :: a -> RAList a -> RAList a
pattern $mCons :: forall {r} {a}.
RAList a -> (a -> RAList a -> r) -> ((# #) -> r) -> r
$bCons :: forall a. a -> RAList a -> RAList a
Cons x xs <- (uncons -> Just (x, xs)) where
  Cons a
x RAList a
xs = a -> RAList a -> RAList a
forall a. a -> RAList a -> RAList a
cons a
x RAList a
xs

-- O(1) worst-case
consValues :: Values a -> RAList a -> RAList a
consValues :: forall a. Values a -> RAList a -> RAList a
consValues Values a
x RAList a
l = case RAList a
l of
    (BHead Word64
w1 Tree a
t1 (BHead Word64
w2 Tree a
t2 RAList a
ts')) | Word64
w1 Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Word64
w2 ->
        let ts :: Word64
ts = Word64
w1 Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
w2 Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
1
            ec :: Word64
ec = Tree a -> Word64
forall a. Tree a -> Word64
treeCount Tree a
t1 Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Tree a -> Word64
forall a. Tree a -> Word64
treeCount Tree a
t2 Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Values a -> Word64
forall a. Values a -> Word64
valuesCount Values a
x
        in Word64 -> Tree a -> RAList a -> RAList a
forall a. Word64 -> Tree a -> RAList a -> RAList a
BHead Word64
ts (Word64 -> Values a -> Tree a -> Tree a -> Tree a
forall a. Word64 -> Values a -> Tree a -> Tree a -> Tree a
Node Word64
ec Values a
x Tree a
t1 Tree a
t2) RAList a
ts'
    RAList a
ts -> Word64 -> Tree a -> RAList a -> RAList a
forall a. Word64 -> Tree a -> RAList a -> RAList a
BHead Word64
1 (Values a -> Tree a
forall a. Values a -> Tree a
Leaf Values a
x) RAList a
ts

-- O(1) worst-case
{-# INLINE cons #-}
cons :: a -> RAList a -> RAList a
cons :: forall a. a -> RAList a -> RAList a
cons a
x = Values a -> RAList a -> RAList a
forall a. Values a -> RAList a -> RAList a
consValues (a -> Values a
forall a. a -> Values a
One a
x)

-- O(1) worst-case
{-# INLINE consSlab #-}
consSlab :: NEV.NonEmptyVector a -> RAList a -> RAList a
consSlab :: forall a. NonEmptyVector a -> RAList a -> RAList a
consSlab NonEmptyVector a
x = Values a -> RAList a -> RAList a
forall a. Values a -> RAList a -> RAList a
consValues (NonEmptyVector a -> Values a
forall a. NonEmptyVector a -> Values a
Many NonEmptyVector a
x)

-- /O(1)/
-- 'uncons' is a bit funny: if we uncons a vector of values
-- initially, we will then uncons the front of *that* and possibly
-- cons the rest back on! Fortunately all these operations are O(1),
-- so it adds up to being okay.
uncons :: RAList a -> Maybe (a, RAList a)
uncons :: forall a. RAList a -> Maybe (a, RAList a)
uncons = \case
    BHead Word64
_ (Leaf Values a
v) RAList a
ts -> (a, RAList a) -> Maybe (a, RAList a)
forall a. a -> Maybe a
Just ((a, RAList a) -> Maybe (a, RAList a))
-> (a, RAList a) -> Maybe (a, RAList a)
forall a b. (a -> b) -> a -> b
$ Values a -> RAList a -> (a, RAList a)
forall a. Values a -> RAList a -> (a, RAList a)
unconsValues Values a
v RAList a
ts
    BHead Word64
_ (Node Word64
treeSize Values a
x Tree a
t1 Tree a
t2) RAList a
ts ->
        -- probably faster than `div w 2`
        let halfSize :: Word64
halfSize = Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
unsafeShiftR Word64
treeSize Int
1
            -- split the node in two)
        in (a, RAList a) -> Maybe (a, RAList a)
forall a. a -> Maybe a
Just ((a, RAList a) -> Maybe (a, RAList a))
-> (a, RAList a) -> Maybe (a, RAList a)
forall a b. (a -> b) -> a -> b
$ Values a -> RAList a -> (a, RAList a)
forall a. Values a -> RAList a -> (a, RAList a)
unconsValues Values a
x (Word64 -> Tree a -> RAList a -> RAList a
forall a. Word64 -> Tree a -> RAList a -> RAList a
BHead Word64
halfSize Tree a
t1 (RAList a -> RAList a) -> RAList a -> RAList a
forall a b. (a -> b) -> a -> b
$ Word64 -> Tree a -> RAList a -> RAList a
forall a. Word64 -> Tree a -> RAList a -> RAList a
BHead Word64
halfSize Tree a
t2 RAList a
ts)
    RAList a
Nil -> Maybe (a, RAList a)
forall a. Maybe a
Nothing

-- 0-based
unsafeIndexZero :: RAList a -> Word64 -> a
unsafeIndexZero :: forall a. RAList a -> Word64 -> a
unsafeIndexZero RAList a
Nil Word64
_  = String -> a
forall a. HasCallStack => String -> a
error String
"out of bounds"
unsafeIndexZero (BHead Word64
_ Tree a
t RAList a
ts) !Word64
i  =
    let tCount :: Word64
tCount = Tree a -> Word64
forall a. Tree a -> Word64
treeCount Tree a
t
    in if Word64
i Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
tCount
    then Word64 -> Tree a -> a
forall a. Word64 -> Tree a -> a
unsafeIndexTree Word64
i Tree a
t
    else RAList a -> Word64 -> a
forall a. RAList a -> Word64 -> a
unsafeIndexZero RAList a
ts (Word64
i Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
tCount)

-- 0-based
safeIndexZero :: RAList a -> Word64 -> Maybe a
safeIndexZero :: forall a. RAList a -> Word64 -> Maybe a
safeIndexZero RAList a
Nil Word64
_  = Maybe a
forall a. Maybe a
Nothing
safeIndexZero (BHead Word64
_ Tree a
t RAList a
ts) !Word64
i  =
    let tCount :: Word64
tCount = Tree a -> Word64
forall a. Tree a -> Word64
treeCount Tree a
t
    in if Word64
i Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
tCount
    then Word64 -> Tree a -> Maybe a
forall a. Word64 -> Tree a -> Maybe a
safeIndexTree Word64
i Tree a
t
    else RAList a -> Word64 -> Maybe a
forall a. RAList a -> Word64 -> Maybe a
safeIndexZero RAList a
ts (Word64
i Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
tCount)

instance RAL.RandomAccessList (RAList a) where
    type Element (RAList a) = a

    {-# INLINABLE empty #-}
    empty :: RAList a
empty = RAList a
forall a. RAList a
Nil
    {-# INLINABLE cons #-}
    cons :: Element (RAList a) -> RAList a -> RAList a
cons = a -> RAList a -> RAList a
Element (RAList a) -> RAList a -> RAList a
forall a. a -> RAList a -> RAList a
Cons
    {-# INLINABLE uncons #-}
    uncons :: RAList a -> Maybe (Element (RAList a), RAList a)
uncons = RAList a -> Maybe (a, RAList a)
RAList a -> Maybe (Element (RAList a), RAList a)
forall a. RAList a -> Maybe (a, RAList a)
uncons
    {-# INLINABLE length #-}
    length :: RAList a -> Word64
length RAList a
Nil            = Word64
0
    length (BHead Word64
_ Tree a
t RAList a
tl) = Tree a -> Word64
forall a. Tree a -> Word64
treeCount Tree a
t Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ RAList a -> Word64
forall e. RandomAccessList e => e -> Word64
RAL.length RAList a
tl
    {-# INLINABLE consSlab #-}
    consSlab :: NonEmptyVector (Element (RAList a)) -> RAList a -> RAList a
consSlab = NonEmptyVector a -> RAList a -> RAList a
NonEmptyVector (Element (RAList a)) -> RAList a -> RAList a
forall a. NonEmptyVector a -> RAList a -> RAList a
consSlab
    {-# INLINABLE indexZero #-}
    indexZero :: RAList a -> Word64 -> Maybe (Element (RAList a))
indexZero RAList a
l Word64
i = RAList a -> Word64 -> Maybe a
forall a. RAList a -> Word64 -> Maybe a
safeIndexZero RAList a
l Word64
i
    {-# INLINABLE unsafeIndexZero #-}
    unsafeIndexZero :: RAList a -> Word64 -> Element (RAList a)
unsafeIndexZero RAList a
l Word64
i = RAList a -> Word64 -> a
forall a. RAList a -> Word64 -> a
unsafeIndexZero RAList a
l Word64
i