{-# LANGUAGE FlexibleContexts #-}

module Data.IntervalIntMap.Internal.GrowableVector
    ( GrowableVector
    , GrowableVectorData(..)
    , new
    , pushBack
    , unsafeFreeze
    , length
    ) where

import Prelude hiding (length)

import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Storable.Mutable as VSM
import           Foreign.Storable (Storable(..))
import           Control.Monad.Primitive (PrimMonad, PrimState)
import           Data.Primitive.MutVar (MutVar, newMutVar, readMutVar, writeMutVar)

{--| This is a growable vector (i.e., one that includes the 'pushBack'
 - function) which must exist in a 'PrimMonad'. It only supports Storable data
 - items.
 -}


type GrowableVector s a = MutVar s (GrowableVectorData s a)

data GrowableVectorData s a = GrowableVectorData !Int !(VSM.MVector s a)

-- | Empty vector
new :: (PrimMonad m, Storable a) => m (GrowableVector (PrimState m) a)
new :: m (GrowableVector (PrimState m) a)
new = do
    GrowableVectorData (PrimState m) a
vd <- Int
-> MVector (PrimState m) a -> GrowableVectorData (PrimState m) a
forall s a. Int -> MVector s a -> GrowableVectorData s a
GrowableVectorData Int
0 (MVector (PrimState m) a -> GrowableVectorData (PrimState m) a)
-> m (MVector (PrimState m) a)
-> m (GrowableVectorData (PrimState m) a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
Int -> m (MVector (PrimState m) a)
VSM.unsafeNew Int
16
    GrowableVectorData (PrimState m) a
-> m (GrowableVector (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
a -> m (MutVar (PrimState m) a)
newMutVar GrowableVectorData (PrimState m) a
vd

-- | Insert an element at the end of the vector
pushBack :: (PrimMonad m, Storable a) => a -> GrowableVector (PrimState m) a -> m ()
pushBack :: a -> GrowableVector (PrimState m) a -> m ()
pushBack a
val GrowableVector (PrimState m) a
gv =
    GrowableVector (PrimState m) a
-> m (GrowableVectorData (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar GrowableVector (PrimState m) a
gv m (GrowableVectorData (PrimState m) a)
-> (GrowableVectorData (PrimState m) a
    -> m (GrowableVectorData (PrimState m) a))
-> m (GrowableVectorData (PrimState m) a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a
-> GrowableVectorData (PrimState m) a
-> m (GrowableVectorData (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
a
-> GrowableVectorData (PrimState m) a
-> m (GrowableVectorData (PrimState m) a)
pushBack' a
val m (GrowableVectorData (PrimState m) a)
-> (GrowableVectorData (PrimState m) a -> m ()) -> m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= GrowableVector (PrimState m) a
-> GrowableVectorData (PrimState m) a -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> a -> m ()
writeMutVar GrowableVector (PrimState m) a
gv

pushBack' :: (PrimMonad m, Storable a) => a -> GrowableVectorData (PrimState m) a -> m (GrowableVectorData (PrimState m) a)
pushBack' :: a
-> GrowableVectorData (PrimState m) a
-> m (GrowableVectorData (PrimState m) a)
pushBack' a
val (GrowableVectorData Int
used MVector (PrimState m) a
vec)
    | Int
used Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== MVector (PrimState m) a -> Int
forall a s. Storable a => MVector s a -> Int
VSM.length MVector (PrimState m) a
vec = do
        MVector (PrimState m) a
vec' <- MVector (PrimState m) a -> Int -> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> m (MVector (PrimState m) a)
VSM.grow MVector (PrimState m) a
vec (MVector (PrimState m) a -> Int
forall a s. Storable a => MVector s a -> Int
VSM.length MVector (PrimState m) a
vec Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) -- multiplying by 1.5 is close to optimal
        a
-> GrowableVectorData (PrimState m) a
-> m (GrowableVectorData (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
a
-> GrowableVectorData (PrimState m) a
-> m (GrowableVectorData (PrimState m) a)
pushBack' a
val (Int
-> MVector (PrimState m) a -> GrowableVectorData (PrimState m) a
forall s a. Int -> MVector s a -> GrowableVectorData s a
GrowableVectorData  Int
used MVector (PrimState m) a
vec')
    | Bool
otherwise = do
        MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VSM.write MVector (PrimState m) a
vec Int
used a
val
        GrowableVectorData (PrimState m) a
-> m (GrowableVectorData (PrimState m) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (GrowableVectorData (PrimState m) a
 -> m (GrowableVectorData (PrimState m) a))
-> GrowableVectorData (PrimState m) a
-> m (GrowableVectorData (PrimState m) a)
forall a b. (a -> b) -> a -> b
$! Int
-> MVector (PrimState m) a -> GrowableVectorData (PrimState m) a
forall s a. Int -> MVector s a -> GrowableVectorData s a
GrowableVectorData (Int
usedInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) MVector (PrimState m) a
vec

-- | This operation is unsafe as original vector should not be used again!
unsafeFreeze :: (PrimMonad m, Storable a) => GrowableVector (PrimState m) a -> m (VS.Vector a)
unsafeFreeze :: GrowableVector (PrimState m) a -> m (Vector a)
unsafeFreeze GrowableVector (PrimState m) a
gv = do
    GrowableVectorData Int
used MVector (PrimState m) a
vec <- GrowableVector (PrimState m) a
-> m (GrowableVectorData (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar GrowableVector (PrimState m) a
gv
    Int -> Vector a -> Vector a
forall a. Storable a => Int -> Vector a -> Vector a
VS.take Int
used (Vector a -> Vector a) -> m (Vector a) -> m (Vector a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState m) a -> m (Vector a)
forall a (m :: * -> *).
(Storable a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VS.unsafeFreeze MVector (PrimState m) a
vec

-- | Return the current number of stored elements
length :: (PrimMonad m, Storable a) => GrowableVector (PrimState m) a -> m Int
length :: GrowableVector (PrimState m) a -> m Int
length GrowableVector (PrimState m) a
gv = do
    GrowableVectorData Int
len MVector (PrimState m) a
_ <- GrowableVector (PrimState m) a
-> m (GrowableVectorData (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar GrowableVector (PrimState m) a
gv
    Int -> m Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
len