{-# 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
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
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 =
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)
data Tree a = Leaf !(Values a)
| 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
data RAList a = BHead
{-# UNPACK #-} !Word64
!(Tree a)
!(RAList a)
| 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)
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 #-}
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
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
{-# 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)
{-# 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)
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 ->
let halfSize :: Word64
halfSize = Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
unsafeShiftR Word64
treeSize Int
1
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
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)
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