{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}
{-# OPTIONS_HADDOCK not-home #-}
{- HLINT ignore "Avoid restricted alias" -}

-- | Vectors with support for appending elements.
module Database.LSMTree.Internal.Vector.Growing
(
    GrowingVector (GrowingVector),
    new,
    append,
    freeze,
    readMaybeLast
)
where

import           Prelude hiding (init, last, length, read)

import           Control.Monad (when)
import           Control.Monad.ST.Strict (ST)
import           Data.Primitive.PrimVar (PrimVar, newPrimVar, readPrimVar,
                     writePrimVar)
import           Data.STRef.Strict (STRef, newSTRef, readSTRef, writeSTRef)
import           Data.Vector (Vector)
import qualified Data.Vector as Mutable (freeze)
import           Data.Vector.Mutable (MVector)
import qualified Data.Vector.Mutable as Mutable (grow, length, new, read, set,
                     slice, take)

{-|
    A vector with support for appending elements.

    Internally, the elements of a growing vector are stored in a buffer. This
    buffer is automatically enlarged whenever this is needed for storing
    additional elements. On each such enlargement, the size of the buffer is
    multiplied by a power of 2, whose exponent is chosen just big enough to make
    the final buffer size at least as high as the new vector length.

    Note that, while buffer sizes and vector lengths are represented as 'Int'
    values internally, the above-described buffer enlargement scheme has the
    consequence that the largest possible buffer size and thus the largest
    possible vector length may not be the maximum 'Int' value. That said, they
    are always greater than half the maximum 'Int' value, which should be enough
    for all practical purposes.
-}
data GrowingVector s a = GrowingVector
                             !(STRef s (MVector s a)) -- Reference to the buffer
                             !(PrimVar s Int)         -- Reference to the length

-- | Creates a new, initially empty, vector.
new :: Int                      -- ^ Initial buffer size
    -> ST s (GrowingVector s a) -- ^ Construction of the vector
new :: forall s a. Int -> ST s (GrowingVector s a)
new Int
illegalInitialBufferSize | Int
illegalInitialBufferSize Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0
    = [Char] -> ST s (GrowingVector s a)
forall a. HasCallStack => [Char] -> a
error [Char]
"Initial buffer size not positive"
new Int
initialBufferSize
    = do
        MVector s a
buffer <- Int -> ST s (MVector (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (MVector (PrimState m) a)
Mutable.new Int
initialBufferSize
        STRef s (MVector s a)
bufferRef <- MVector s a -> ST s (STRef s (MVector s a))
forall a s. a -> ST s (STRef s a)
newSTRef (MVector s a -> ST s (STRef s (MVector s a)))
-> MVector s a -> ST s (STRef s (MVector s a))
forall a b. (a -> b) -> a -> b
$! MVector s a
buffer
        PrimVar s Int
lengthRef <- Int -> ST s (PrimVar (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
a -> m (PrimVar (PrimState m) a)
newPrimVar Int
0
        GrowingVector s a -> ST s (GrowingVector s a)
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return (STRef s (MVector s a) -> PrimVar s Int -> GrowingVector s a
forall s a.
STRef s (MVector s a) -> PrimVar s Int -> GrowingVector s a
GrowingVector STRef s (MVector s a)
bufferRef PrimVar s Int
lengthRef)

{-|
    Appends a value a certain number of times to a vector. If a negative number
    is provided as the count, the vector is not changed.
-}
append :: forall s a . GrowingVector s a -> Int -> a -> ST s ()
append :: forall s a. GrowingVector s a -> Int -> a -> ST s ()
append GrowingVector s a
_ Int
pseudoCount a
_ | Int
pseudoCount Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0
    = () -> ST s ()
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
append (GrowingVector STRef s (MVector s a)
bufferRef PrimVar s Int
lengthRef) Int
count a
val
    = do
          Int
length <- PrimVar (PrimState (ST s)) Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> m a
readPrimVar PrimVar s Int
PrimVar (PrimState (ST s)) Int
lengthRef
          ST s ()
makeRoom
          MVector s a
buffer' <- STRef s (MVector s a) -> ST s (MVector s a)
forall s a. STRef s a -> ST s a
readSTRef STRef s (MVector s a)
bufferRef
          MVector (PrimState (ST s)) a -> a -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> a -> m ()
Mutable.set (Int -> Int -> MVector s a -> MVector s a
forall s a. Int -> Int -> MVector s a -> MVector s a
Mutable.slice Int
length Int
count MVector s a
buffer') (a -> ST s ()) -> a -> ST s ()
forall a b. (a -> b) -> a -> b
$! a
val
    where

    makeRoom :: ST s ()
    makeRoom :: ST s ()
makeRoom = do
        Int
length <- PrimVar (PrimState (ST s)) Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> m a
readPrimVar PrimVar s Int
PrimVar (PrimState (ST s)) Int
lengthRef
        Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
count Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
forall a. Bounded a => a
maxBound Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
length) ([Char] -> ST s ()
forall a. HasCallStack => [Char] -> a
error [Char]
"New length too large")
        MVector s a
buffer <- STRef s (MVector s a) -> ST s (MVector s a)
forall s a. STRef s a -> ST s a
readSTRef STRef s (MVector s a)
bufferRef
        let

            bufferSize :: Int
            !bufferSize :: Int
bufferSize = MVector s a -> Int
forall s a. MVector s a -> Int
Mutable.length MVector s a
buffer

            length' :: Int
            !length' :: Int
length' = Int
length Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
count

        Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
bufferSize Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
length') (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
            let

                higherBufferSizes :: [Int]
                higherBufferSizes :: [Int]
higherBufferSizes = [Int] -> [Int]
forall a. HasCallStack => [a] -> [a]
tail ([Int]
init [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
last]) where

                    init :: [Int]
                    last :: Int
                    ([Int]
init, Int
last : [Int]
_) = (Int -> Bool) -> [Int] -> ([Int], [Int])
forall a. (a -> Bool) -> [a] -> ([a], [a])
span (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
forall a. Bounded a => a
maxBound Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) ([Int] -> ([Int], [Int])) -> [Int] -> ([Int], [Int])
forall a b. (a -> b) -> a -> b
$
                                       (Int -> Int) -> Int -> [Int]
forall a. (a -> a) -> a -> [a]
iterate (Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2) Int
bufferSize
                    {-NOTE:
                        In order to prevent overflow, we have to start with the
                        current buffer size here, although we know that it is
                        not sufficient.
                    -}

                sufficientBufferSizes :: [Int]
                sufficientBufferSizes :: [Int]
sufficientBufferSizes = (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
length') [Int]
higherBufferSizes

            case [Int]
sufficientBufferSizes of
                []
                    -> [Char] -> ST s ()
forall a. HasCallStack => [Char] -> a
error [Char]
"No sufficient buffer size available"
                Int
bufferSize' : [Int]
_
                    -> MVector (PrimState (ST s)) a
-> Int -> ST s (MVector (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m (MVector (PrimState m) a)
Mutable.grow MVector s a
MVector (PrimState (ST s)) a
buffer (Int
bufferSize' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
bufferSize) ST s (MVector s a) -> (MVector s a -> ST s ()) -> ST s ()
forall a b. ST s a -> (a -> ST s b) -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>=
                       (STRef s (MVector s a) -> MVector s a -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s (MVector s a)
bufferRef $!)
        PrimVar (PrimState (ST s)) Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> a -> m ()
writePrimVar PrimVar s Int
PrimVar (PrimState (ST s)) Int
lengthRef Int
length'

-- | Turns a growing vector into an ordinary vector.
freeze :: GrowingVector s a -> ST s (Vector a)
freeze :: forall s a. GrowingVector s a -> ST s (Vector a)
freeze (GrowingVector STRef s (MVector s a)
bufferRef PrimVar s Int
lengthRef) = do
    MVector s a
buffer <- STRef s (MVector s a) -> ST s (MVector s a)
forall s a. STRef s a -> ST s a
readSTRef STRef s (MVector s a)
bufferRef
    Int
length <- PrimVar (PrimState (ST s)) Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> m a
readPrimVar PrimVar s Int
PrimVar (PrimState (ST s)) Int
lengthRef
    MVector (PrimState (ST s)) a -> ST s (Vector a)
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> m (Vector a)
Mutable.freeze (Int -> MVector s a -> MVector s a
forall s a. Int -> MVector s a -> MVector s a
Mutable.take Int
length MVector s a
buffer)

-- | Reads the last element of a growing vector if it exists.
readMaybeLast :: GrowingVector s a -> ST s (Maybe a)
readMaybeLast :: forall s a. GrowingVector s a -> ST s (Maybe a)
readMaybeLast (GrowingVector STRef s (MVector s a)
bufferRef PrimVar s Int
lengthRef) = do
    Int
length <- PrimVar (PrimState (ST s)) Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> m a
readPrimVar PrimVar s Int
PrimVar (PrimState (ST s)) Int
lengthRef
    if Int
length Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
        then Maybe a -> ST s (Maybe a)
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing
        else do
                 MVector s a
buffer <- STRef s (MVector s a) -> ST s (MVector s a)
forall s a. STRef s a -> ST s a
readSTRef STRef s (MVector s a)
bufferRef
                 a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> ST s a -> ST s (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) a -> Int -> ST s a
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
Mutable.read MVector s a
MVector (PrimState (ST s)) a
buffer (Int -> Int
forall a. Enum a => a -> a
pred Int
length)