{-# OPTIONS_HADDOCK not-home #-}
-- |
-- Copyright: (c) 2021 Xy Ren
-- License: BSD3
-- Maintainer: xy.r@outlook.com
-- Stability: unstable
-- Portability: non-portable (GHC only)
module Data.RadixVec (RadixVec, size, empty, lookup, update, snoc) where

import           Control.Monad.ST          (ST)
import           Data.Bits                 (Bits (unsafeShiftL, unsafeShiftR, (.&.)), FiniteBits (countTrailingZeros))
import           Data.Primitive.MachDeps   (sIZEOF_INT)
import           Data.Primitive.SmallArray (SmallArray, SmallMutableArray, copySmallArray, indexSmallArray,
                                            newSmallArray, readSmallArray, runSmallArray, sizeofSmallArray,
                                            thawSmallArray, writeSmallArray)
import           Prelude                   hiding (lookup)

-- | An efficient vector type, implemented as a radix tree. It has the following time complexities:
--
-- * Lookup: \( O(\log n) \)
-- * Update: \( O(\log n) \)
-- * Append: \( O(\log n) \)
--
-- The branching factor (base of log) is 32 therefore the time is close to constant in most cases. Note that in
-- practice, lookup is faster than update, and update is faster than append.
data RadixVec a = RadixVec !Int !(Tree a)

type Shift = Int

-- | The \( \log_2 \) of the branching factor. The branching factor is set to be 32 for now but may change in the
-- future.
factor :: Int
factor :: Int
factor = Int
5

-- | A mask covering one chunk of an index.
initialMask :: Int
initialMask :: Int
initialMask = (Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
factor) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1

-- | A radix tree. The tree is always left-leaning.
data Tree a
  = Tip
    {-# UNPACK #-} !(SmallArray a)
  | Node
    {-# UNPACK #-} !Shift
    {-# UNPACK #-} !(SmallArray (Tree a))

-- | Mask a portion of an index.
mask :: Shift -> Int -> Int
mask :: Int -> Int -> Int
mask Int
s Int
x = Int
initialMask Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. (Int
x Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
s)

-- | Mask the zeroth portion of the index.
mask0 :: Int -> Int
mask0 :: Int -> Int
mask0 Int
x = Int
initialMask Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
x

-- | Alter an element in a 'SmallMutableArray' by a function.
alterSmallArray :: SmallMutableArray s a -> Int -> (a -> a) -> ST s ()
alterSmallArray :: SmallMutableArray s a -> Int -> (a -> a) -> ST s ()
alterSmallArray SmallMutableArray s a
marr Int
ix a -> a
f = do
  a
x <- SmallMutableArray (PrimState (ST s)) a -> Int -> ST s a
forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> m a
readSmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
marr Int
ix
  SmallMutableArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
marr Int
ix (a -> ST s ()) -> a -> ST s ()
forall a b. (a -> b) -> a -> b
$! a -> a
f a
x

-- | Get the size of the 'RadixVec'.
size :: RadixVec a -> Int
size :: RadixVec a -> Int
size (RadixVec Int
sz Tree a
_) = Int
sz

-- | The empty 'RadixVec'.
empty :: RadixVec a
empty :: RadixVec a
empty = Int -> Tree a -> RadixVec a
forall a. Int -> Tree a -> RadixVec a
RadixVec Int
0 (Tree a -> RadixVec a) -> Tree a -> RadixVec a
forall a b. (a -> b) -> a -> b
$ SmallArray a -> Tree a
forall a. SmallArray a -> Tree a
Tip (SmallArray a -> Tree a) -> SmallArray a -> Tree a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray ((forall s. ST s (SmallMutableArray s a)) -> SmallArray a)
-> (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a b. (a -> b) -> a -> b
$ Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: Type -> Type) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
0 (a -> ST s (SmallMutableArray (PrimState (ST s)) a))
-> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall a b. (a -> b) -> a -> b
$ [Char] -> a
forall a. HasCallStack => [Char] -> a
error
  [Char]
"Cleff.Internal.RadixVec: Encountered an element in an empty RadixVec. Please report this as a bug."

-- | Lookup in a 'RadixVec' by an index. This does not perform any bounds check.
lookup :: Int -> RadixVec a -> a
lookup :: Int -> RadixVec a -> a
lookup Int
ix (RadixVec Int
_ Tree a
tree) = Tree a -> a
go Tree a
tree
  where
    go :: Tree a -> a
go (Tip SmallArray a
arr)    = SmallArray a -> Int -> a
forall a. SmallArray a -> Int -> a
indexSmallArray SmallArray a
arr (Int
initialMask Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
ix)
    go (Node Int
s SmallArray (Tree a)
arr) = Tree a -> a
go (SmallArray (Tree a) -> Int -> Tree a
forall a. SmallArray a -> Int -> a
indexSmallArray SmallArray (Tree a)
arr (Int -> Int -> Int
mask Int
s Int
ix))

-- | Update a value in a 'RadixVec' by an index. The value will be forced before installing. This does not perform any
-- bounds check.
update :: Int -> a -> RadixVec a -> RadixVec a
update :: Int -> a -> RadixVec a -> RadixVec a
update Int
ix a
x (RadixVec Int
len Tree a
tree) = Int -> Tree a -> RadixVec a
forall a. Int -> Tree a -> RadixVec a
RadixVec Int
len (Tree a -> Tree a
go Tree a
tree)
  where
    go :: Tree a -> Tree a
go (Tip SmallArray a
arr) = SmallArray a -> Tree a
forall a. SmallArray a -> Tree a
Tip (SmallArray a -> Tree a) -> SmallArray a -> Tree a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray do
      SmallMutableArray s a
marr <- SmallArray a
-> Int -> Int -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: Type -> Type) a.
PrimMonad m =>
SmallArray a -> Int -> Int -> m (SmallMutableArray (PrimState m) a)
thawSmallArray SmallArray a
arr Int
0 (SmallArray a -> Int
forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
arr)
      SmallMutableArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
marr (Int -> Int
mask0 Int
ix) (a -> ST s ()) -> a -> ST s ()
forall a b. (a -> b) -> a -> b
$! a
x
      SmallMutableArray s a -> ST s (SmallMutableArray s a)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure SmallMutableArray s a
marr
    go (Node Int
s SmallArray (Tree a)
arr) = Int -> SmallArray (Tree a) -> Tree a
forall a. Int -> SmallArray (Tree a) -> Tree a
Node Int
s (SmallArray (Tree a) -> Tree a) -> SmallArray (Tree a) -> Tree a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallMutableArray s (Tree a)))
-> SmallArray (Tree a)
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray do
      SmallMutableArray s (Tree a)
marr <- SmallArray (Tree a)
-> Int
-> Int
-> ST s (SmallMutableArray (PrimState (ST s)) (Tree a))
forall (m :: Type -> Type) a.
PrimMonad m =>
SmallArray a -> Int -> Int -> m (SmallMutableArray (PrimState m) a)
thawSmallArray SmallArray (Tree a)
arr Int
0 (SmallArray (Tree a) -> Int
forall a. SmallArray a -> Int
sizeofSmallArray SmallArray (Tree a)
arr)
      SmallMutableArray s (Tree a)
-> Int -> (Tree a -> Tree a) -> ST s ()
forall s a. SmallMutableArray s a -> Int -> (a -> a) -> ST s ()
alterSmallArray SmallMutableArray s (Tree a)
marr (Int -> Int -> Int
mask Int
s Int
ix) Tree a -> Tree a
go
      SmallMutableArray s (Tree a) -> ST s (SmallMutableArray s (Tree a))
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure SmallMutableArray s (Tree a)
marr

-- | Append a value to a 'RadixVec'. The value will be forced before installing. This does not perform any bounds check.
snoc :: RadixVec a -> a -> RadixVec a
snoc :: RadixVec a -> a -> RadixVec a
snoc (RadixVec Int
len Tree a
tree) a
x
  | Int
ins Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
topShift = Int -> Tree a -> RadixVec a
forall a. Int -> Tree a -> RadixVec a
RadixVec (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Tree a -> Tree a
go Tree a
tree)
  | Bool
otherwise = Int -> Tree a -> RadixVec a
forall a. Int -> Tree a -> RadixVec a
RadixVec (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Tree a -> RadixVec a) -> Tree a -> RadixVec a
forall a b. (a -> b) -> a -> b
$ Int -> SmallArray (Tree a) -> Tree a
forall a. Int -> SmallArray (Tree a) -> Tree a
Node (Int
topShift Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
factor) (SmallArray (Tree a) -> Tree a) -> SmallArray (Tree a) -> Tree a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallMutableArray s (Tree a)))
-> SmallArray (Tree a)
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray ((forall s. ST s (SmallMutableArray s (Tree a)))
 -> SmallArray (Tree a))
-> (forall s. ST s (SmallMutableArray s (Tree a)))
-> SmallArray (Tree a)
forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray s (Tree a)
marr <- Int
-> Tree a -> ST s (SmallMutableArray (PrimState (ST s)) (Tree a))
forall (m :: Type -> Type) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
2 (Tree a -> ST s (SmallMutableArray s (Tree a)))
-> Tree a -> ST s (SmallMutableArray s (Tree a))
forall a b. (a -> b) -> a -> b
$! Tree a
tree
    SmallMutableArray (PrimState (ST s)) (Tree a)
-> Int -> Tree a -> ST s ()
forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s (Tree a)
SmallMutableArray (PrimState (ST s)) (Tree a)
marr Int
1 (Tree a -> ST s ()) -> Tree a -> ST s ()
forall a b. (a -> b) -> a -> b
$! Int -> Tree a
branch Int
topShift
    SmallMutableArray s (Tree a) -> ST s (SmallMutableArray s (Tree a))
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure SmallMutableArray s (Tree a)
marr
  where
    topShift :: Int
topShift = case Tree a
tree of
      Tip SmallArray a
_    -> Int
0
      Node Int
s SmallArray (Tree a)
_ -> Int
s
    ins :: Int
ins = (Int -> Int
forall b. FiniteBits b => b -> Int
countTrailingZeros Int
len Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
sIZEOF_INT Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
factor) Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
factor
    branch :: Int -> Tree a
branch Int
0 = SmallArray a -> Tree a
forall a. SmallArray a -> Tree a
Tip (SmallArray a -> Tree a) -> SmallArray a -> Tree a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray ((forall s. ST s (SmallMutableArray s a)) -> SmallArray a)
-> (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a b. (a -> b) -> a -> b
$ Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: Type -> Type) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
1 (a -> ST s (SmallMutableArray s a))
-> a -> ST s (SmallMutableArray s a)
forall a b. (a -> b) -> a -> b
$! a
x
    branch Int
s = Int -> SmallArray (Tree a) -> Tree a
forall a. Int -> SmallArray (Tree a) -> Tree a
Node Int
s (SmallArray (Tree a) -> Tree a) -> SmallArray (Tree a) -> Tree a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallMutableArray s (Tree a)))
-> SmallArray (Tree a)
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray ((forall s. ST s (SmallMutableArray s (Tree a)))
 -> SmallArray (Tree a))
-> (forall s. ST s (SmallMutableArray s (Tree a)))
-> SmallArray (Tree a)
forall a b. (a -> b) -> a -> b
$ Int
-> Tree a -> ST s (SmallMutableArray (PrimState (ST s)) (Tree a))
forall (m :: Type -> Type) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
1 (Tree a -> ST s (SmallMutableArray s (Tree a)))
-> Tree a -> ST s (SmallMutableArray s (Tree a))
forall a b. (a -> b) -> a -> b
$! Int -> Tree a
branch (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
factor)
    enlarge :: SmallArray a -> a -> SmallArray a
enlarge SmallArray a
arr a
new = (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray do
      let sz :: Int
sz = SmallArray a -> Int
forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
arr
      SmallMutableArray s a
marr <- Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: Type -> Type) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray (Int
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (a -> ST s (SmallMutableArray s a))
-> a -> ST s (SmallMutableArray s a)
forall a b. (a -> b) -> a -> b
$! a
new
      SmallMutableArray (PrimState (ST s)) a
-> Int -> SmallArray a -> Int -> Int -> ST s ()
forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallArray a -> Int -> Int -> m ()
copySmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
marr Int
0 SmallArray a
arr Int
0 Int
sz
      SmallMutableArray s a -> ST s (SmallMutableArray s a)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure SmallMutableArray s a
marr
    go :: Tree a -> Tree a
go (Tip SmallArray a
arr) = SmallArray a -> Tree a
forall a. SmallArray a -> Tree a
Tip (SmallArray a -> Tree a) -> SmallArray a -> Tree a
forall a b. (a -> b) -> a -> b
$ SmallArray a -> a -> SmallArray a
forall a. SmallArray a -> a -> SmallArray a
enlarge SmallArray a
arr a
x
    go (Node Int
s SmallArray (Tree a)
arr)
      | Int
ins Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
s = Int -> SmallArray (Tree a) -> Tree a
forall a. Int -> SmallArray (Tree a) -> Tree a
Node Int
s (SmallArray (Tree a) -> Tree a) -> SmallArray (Tree a) -> Tree a
forall a b. (a -> b) -> a -> b
$ SmallArray (Tree a) -> Tree a -> SmallArray (Tree a)
forall a. SmallArray a -> a -> SmallArray a
enlarge SmallArray (Tree a)
arr (Tree a -> SmallArray (Tree a)) -> Tree a -> SmallArray (Tree a)
forall a b. (a -> b) -> a -> b
$ Int -> Tree a
branch (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
factor)
      | Bool
otherwise = Int -> SmallArray (Tree a) -> Tree a
forall a. Int -> SmallArray (Tree a) -> Tree a
Node Int
s (SmallArray (Tree a) -> Tree a) -> SmallArray (Tree a) -> Tree a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallMutableArray s (Tree a)))
-> SmallArray (Tree a)
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray do
        SmallMutableArray s (Tree a)
marr <- SmallArray (Tree a)
-> Int
-> Int
-> ST s (SmallMutableArray (PrimState (ST s)) (Tree a))
forall (m :: Type -> Type) a.
PrimMonad m =>
SmallArray a -> Int -> Int -> m (SmallMutableArray (PrimState m) a)
thawSmallArray SmallArray (Tree a)
arr Int
0 (SmallArray (Tree a) -> Int
forall a. SmallArray a -> Int
sizeofSmallArray SmallArray (Tree a)
arr)
        SmallMutableArray s (Tree a)
-> Int -> (Tree a -> Tree a) -> ST s ()
forall s a. SmallMutableArray s a -> Int -> (a -> a) -> ST s ()
alterSmallArray SmallMutableArray s (Tree a)
marr (Int -> Int -> Int
mask Int
s Int
len) Tree a -> Tree a
go
        SmallMutableArray s (Tree a) -> ST s (SmallMutableArray s (Tree a))
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure SmallMutableArray s (Tree a)
marr