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

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

import Data.RandomAccessList.Class qualified as RAL

-- | 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 = Leaf a
            | Node 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)

-- | 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)

{-# 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
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

-- 1-based
safeIndexOne :: RAList a -> Word64 -> Maybe a
safeIndexOne :: forall a. RAList a -> Word64 -> Maybe a
safeIndexOne RAList a
Nil Word64
_ = Maybe a
forall a. Maybe a
Nothing
safeIndexOne (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
safeIndexOne 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
_ Word64
0 Tree a
_ = Maybe a
forall a. Maybe a
Nothing -- "index zero"
    indexTree Word64
1 Word64
1 (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
1 (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`
            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 -> Maybe a
forall a. Word64 -> Word64 -> Tree a -> Maybe a
indexTree Word64
halfSize Word64
offset' 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
halfSize) Tree a
t2

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