{-# LANGUAGE BangPatterns         #-}
{-# LANGUAGE LambdaCase           #-}
{-# LANGUAGE PatternSynonyms      #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE TypeFamilies         #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns         #-}
module Data.RandomAccessList.SkewBinary
    ( RAList(Cons,Nil)
    , safeIndexZero
    , unsafeIndexZero
    , safeIndexOne
    , safeIndexOneCont
    , unsafeIndexOne
    , Data.RandomAccessList.SkewBinary.null
    , uncons
    ) where

import Data.Bits (unsafeShiftR)
import Data.Word
import GHC.Exts

import Data.RandomAccessList.Class qualified as RAL

-- 'Node' appears first to make it more likely for GHC to reorder pattern matches to make the 'Node'
-- one appear first (which makes it more efficient).
-- | A complete binary tree.
-- Note: the size of the tree is not stored/cached,
-- unless it appears as a root tree in 'RAList', which the size is stored inside the Cons.
data Tree a = Node a !(Tree a) !(Tree a)
            | Leaf 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)

-- | A strict list of complete binary trees accompanied by their size.
-- The trees appear in >=-size order.
-- Note: this list is strict in its spine, unlike the Prelude list
data RAList a = BHead
               {-# UNPACK #-} !Word64 -- ^ the size of the head tree
               !(Tree a) -- ^ the head tree
               !(RAList a) -- ^ the tail trees
             | Nil
             -- the derived Eq instance is correct,
             -- because binary skew numbers have unique representation
             -- and hence all trees of the same size will have the same structure
             deriving stock (RAList a -> RAList a -> Bool
(RAList a -> RAList a -> Bool)
-> (RAList a -> RAList a -> Bool) -> Eq (RAList a)
forall a. Eq a => RAList a -> RAList a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => RAList a -> RAList a -> Bool
== :: RAList a -> RAList a -> Bool
$c/= :: forall a. Eq a => RAList a -> RAList a -> Bool
/= :: RAList a -> RAList a -> Bool
Eq, 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)

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

{-# 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
cons :: a -> RAList a -> RAList a
cons :: forall a. a -> RAList a -> RAList a
cons a
x = \case
    (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 -> Word64 -> Tree a -> RAList a -> RAList a
forall a. Word64 -> Tree a -> RAList a -> RAList a
BHead (Word64
2Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
*Word64
w1Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+Word64
1) (a -> Tree a -> Tree a -> Tree a
forall a. a -> Tree a -> Tree a -> Tree a
Node 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 (a -> Tree a
forall a. a -> Tree a
Leaf a
x) RAList a
ts

-- /O(1)/
uncons :: RAList a -> Maybe (a, RAList a)
uncons :: forall a. RAList a -> Maybe (a, RAList a)
uncons = \case
    BHead Word64
_ (Leaf a
x) RAList a
ts -> (a, RAList a) -> Maybe (a, RAList a)
forall a. a -> Maybe a
Just (a
x, RAList a
ts)
    BHead Word64
treeSize (Node 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
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
w Tree a
t RAList a
ts) !Word64
i  =
    if Word64
i Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
w
    then Word64 -> Word64 -> Tree a -> a
forall a. Word64 -> Word64 -> Tree a -> a
indexTree Word64
w Word64
i Tree a
t
    else RAList a -> Word64 -> a
forall a. RAList a -> Word64 -> a
unsafeIndexZero RAList a
ts (Word64
iWord64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
-Word64
w)
  where
    indexTree :: Word64 -> Word64 -> Tree a -> a
    indexTree :: forall a. Word64 -> Word64 -> Tree a -> a
indexTree Word64
1 Word64
0 (Leaf a
x) = a
x
    indexTree Word64
_ Word64
_ (Leaf a
_) = String -> a
forall a. HasCallStack => String -> a
error String
"out of bounds"
    indexTree Word64
_ Word64
0 (Node a
x Tree a
_ Tree a
_) = a
x
    indexTree Word64
treeSize Word64
offset (Node a
_ Tree a
t1 Tree a
t2) =
        let halfSize :: Word64
halfSize = Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
unsafeShiftR Word64
treeSize Int
1 -- probably faster than `div w 2`
        in if Word64
offset Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
<= Word64
halfSize
           then Word64 -> Word64 -> Tree a -> a
forall a. Word64 -> Word64 -> Tree a -> a
indexTree Word64
halfSize (Word64
offset Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
1) Tree a
t1
           else Word64 -> Word64 -> Tree a -> a
forall a. Word64 -> Word64 -> Tree a -> a
indexTree Word64
halfSize (Word64
offset Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
1 Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
halfSize) Tree a
t2

-- 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
w Tree a
t RAList a
ts) !Word64
i  =
    if Word64
i Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
w
    then Word64 -> Word64 -> Tree a -> Maybe a
forall a. Word64 -> Word64 -> Tree a -> Maybe a
indexTree Word64
w Word64
i Tree a
t
    else RAList a -> Word64 -> Maybe a
forall a. RAList a -> Word64 -> Maybe a
safeIndexZero RAList a
ts (Word64
iWord64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
-Word64
w)
  where
    indexTree :: Word64 -> Word64 -> Tree a -> Maybe a
    indexTree :: forall a. Word64 -> Word64 -> Tree a -> Maybe a
indexTree Word64
1 Word64
0 (Leaf a
x) = a -> Maybe a
forall a. a -> Maybe a
Just a
x
    indexTree Word64
_ Word64
_ (Leaf a
_) = Maybe a
forall a. Maybe a
Nothing
    indexTree Word64
_ Word64
0 (Node a
x Tree a
_ Tree a
_) = a -> Maybe a
forall a. a -> Maybe a
Just a
x
    indexTree Word64
treeSize Word64
offset (Node a
_ Tree a
t1 Tree a
t2) =
        let halfSize :: Word64
halfSize = Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
unsafeShiftR Word64
treeSize Int
1 -- probably faster than `div w 2`
        in if Word64
offset Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
<= Word64
halfSize
           then Word64 -> Word64 -> Tree a -> Maybe a
forall a. Word64 -> Word64 -> Tree a -> Maybe a
indexTree Word64
halfSize (Word64
offset Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
1) Tree a
t1
           else Word64 -> Word64 -> Tree a -> Maybe a
forall a. Word64 -> Word64 -> Tree a -> Maybe a
indexTree Word64
halfSize (Word64
offset Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
1 Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
halfSize) Tree a
t2

-- 1-based
unsafeIndexOne :: RAList a -> Word64 -> a
unsafeIndexOne :: forall a. RAList a -> Word64 -> a
unsafeIndexOne RAList a
Nil Word64
_ = String -> a
forall a. HasCallStack => String -> a
error String
"out of bounds"
unsafeIndexOne (BHead Word64
w Tree a
t RAList a
ts) !Word64
i =
    if Word64
i Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
<= Word64
w
    then Word64 -> Word64 -> Tree a -> a
forall a. Word64 -> Word64 -> Tree a -> a
indexTree Word64
w Word64
i Tree a
t
    else RAList a -> Word64 -> a
forall a. RAList a -> Word64 -> a
unsafeIndexOne RAList a
ts (Word64
iWord64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
-Word64
w)
  where
    indexTree :: Word64 -> Word64 -> Tree a -> a
    indexTree :: forall a. Word64 -> Word64 -> Tree a -> a
indexTree Word64
_ Word64
0 Tree a
_ = String -> a
forall a. HasCallStack => String -> a
error String
"index zero"
    indexTree Word64
1 Word64
1 (Leaf a
x) = a
x
    indexTree Word64
_ Word64
_ (Leaf a
_) = String -> a
forall a. HasCallStack => String -> a
error String
"out of bounds"
    indexTree Word64
_ Word64
1 (Node a
x Tree a
_ Tree a
_) = a
x
    indexTree Word64
treeSize Word64
offset (Node a
_ Tree a
t1 Tree a
t2) =
        let halfSize :: Word64
halfSize = Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
unsafeShiftR Word64
treeSize Int
1 -- probably faster than `div w 2`
            offset' :: Word64
offset' = Word64
offset Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
1
        in if Word64
offset' Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
<= Word64
halfSize
           then Word64 -> Word64 -> Tree a -> a
forall a. Word64 -> Word64 -> Tree a -> a
indexTree Word64
halfSize Word64
offset' Tree a
t1
           else Word64 -> Word64 -> Tree a -> a
forall a. Word64 -> Word64 -> Tree a -> a
indexTree Word64
halfSize (Word64
offset' Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
halfSize) Tree a
t2

{- Note [Optimizations of safeIndexOneCont]
Bangs in the local definitions of 'safeIndexOneCont' are needed to tell GHC that the functions are
strict in the 'Word64' argument, so that GHC produces workers operating on @Word64#@.

The function itself is CPS-ed, so that the arguments force the local definitions to be retained
within 'safeIndexOneCont' instead of being pulled out via full-laziness or some other optimization
pass. This ensures that when 'safeIndexOneCont' gets inlined, the local definitions appear directly
in the GHC Core, allowing GHC to inline the arguments of 'safeIndexOneCont' and transform the whole
thing into a beautiful recursive join point full of @Word64#@s, i.e. allocating very little if
anything at all.
-}

-- See Note [Optimizations of safeIndexOneCont].
safeIndexOneCont :: forall a b. b -> (a -> b) -> RAList a -> Word64 -> b
safeIndexOneCont :: forall a b. b -> (a -> b) -> RAList a -> Word64 -> b
safeIndexOneCont b
z a -> b
f = RAList a -> Word64 -> b
findTree where
    findTree :: RAList a -> Word64 -> b
    -- See Note [Optimizations of safeIndexOneCont].
    findTree :: RAList a -> Word64 -> b
findTree RAList a
Nil !Word64
_ = b
z
    findTree (BHead Word64
w Tree a
t RAList a
ts) Word64
i =
        if Word64
i Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
<= Word64
w
        then Word64 -> Word64 -> Tree a -> b
indexTree Word64
w Word64
i Tree a
t
        else RAList a -> Word64 -> b
findTree RAList a
ts (Word64
iWord64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
-Word64
w)

    indexTree :: Word64 -> Word64 -> Tree a -> b
    -- See Note [Optimizations of safeIndexOneCont].
    indexTree :: Word64 -> Word64 -> Tree a -> b
indexTree !Word64
w Word64
1 Tree a
t = case Tree a
t of
        Node a
x Tree a
_ Tree a
_ -> a -> b
f a
x
        Leaf a
x     -> if Word64
w Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Word64
1 then a -> b
f a
x else b
z
    indexTree Word64
_ Word64
0 Tree a
_ = b
z -- "index zero"
    indexTree Word64
_ Word64
_ (Leaf a
_) = b
z
    indexTree Word64
treeSize Word64
offset (Node a
_ Tree a
t1 Tree a
t2) =
        let halfSize :: Word64
halfSize = Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
unsafeShiftR Word64
treeSize Int
1 -- probably faster than `div w 2`
            offset' :: Word64
offset' = Word64
offset Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
1
        in if Word64
offset' Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
<= Word64
halfSize
           then Word64 -> Word64 -> Tree a -> b
indexTree Word64
halfSize Word64
offset' Tree a
t1
           else Word64 -> Word64 -> Tree a -> b
indexTree Word64
halfSize (Word64
offset' Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
halfSize) Tree a
t2
{-# INLINE safeIndexOneCont #-}

-- 1-based
safeIndexOne :: RAList a -> Word64 -> Maybe a
safeIndexOne :: forall a. RAList a -> Word64 -> Maybe a
safeIndexOne = Maybe a -> (a -> Maybe a) -> RAList a -> Word64 -> Maybe a
forall a b. b -> (a -> b) -> RAList a -> Word64 -> b
safeIndexOneCont Maybe a
forall a. Maybe a
Nothing a -> Maybe a
forall a. a -> Maybe a
Just

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
sz Tree a
_ RAList a
tl) = Word64
sz 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 indexZero #-}
    indexZero :: RAList a -> Word64 -> Maybe (Element (RAList a))
indexZero = RAList a -> Word64 -> Maybe a
RAList a -> Word64 -> Maybe (Element (RAList a))
forall a. RAList a -> Word64 -> Maybe a
safeIndexZero
    {-# INLINABLE indexOne #-}
    indexOne :: RAList a -> Word64 -> Maybe (Element (RAList a))
indexOne = RAList a -> Word64 -> Maybe a
RAList a -> Word64 -> Maybe (Element (RAList a))
forall a. RAList a -> Word64 -> Maybe a
safeIndexOne
    {-# INLINABLE unsafeIndexZero #-}
    unsafeIndexZero :: RAList a -> Word64 -> Element (RAList a)
unsafeIndexZero = RAList a -> Word64 -> a
RAList a -> Word64 -> Element (RAList a)
forall a. RAList a -> Word64 -> a
unsafeIndexZero
    {-# INLINABLE unsafeIndexOne #-}
    unsafeIndexOne :: RAList a -> Word64 -> Element (RAList a)
unsafeIndexOne = RAList a -> Word64 -> a
RAList a -> Word64 -> Element (RAList a)
forall a. RAList a -> Word64 -> a
unsafeIndexOne