-- | A mutable vector that grows in size.
--
-- Example usage:
--
-- > import qualified Data.ProtoLens.Encoding.Growing as Growing
-- > import qualified Data.Vector.Unboxed as V
-- > test :: IO (V.Vector Int)
-- > test = do
-- >     v <- Growing.new
-- >     v' <- Growing.append v 1
-- >     v'' <- Growing.append v' 2
-- >     v''' <- Growing.append v'' 3
-- >     unsafeFreeze v'''
module Data.ProtoLens.Encoding.Growing (
    Growing,
    new,
    append,
    unsafeFreeze,
    RealWorld,
    ) where

import Control.Monad.Primitive (PrimMonad, PrimState, RealWorld)

import qualified Data.Vector.Generic as V
import qualified Data.Vector.Generic.Mutable as MV

-- | A mutable vector which can increase in capacity.
data Growing v s a = Growing
    {-# UNPACK #-} !Int
        -- The number of elements in the mutable vector
        -- that have already been set.
    !(V.Mutable v s a)
        -- TODOs for efficiency:
        -- - Try unpacking this.  It's difficult as-is because
        --   V.Mutable is a type function.
        -- - MVectors support slicing, but we're not using that
        --   functionality, so we're passing around an extra unnecessary
        --   Int.

-- | Create a new empty growing vector.
new :: (PrimMonad m, V.Vector v a) => m (Growing v (PrimState m) a)
new :: forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
m (Growing v (PrimState m) a)
new = forall (v :: * -> *) s a. Int -> Mutable v s a -> Growing v s a
Growing Int
0 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
MV.new Int
0

-- | Unsafely convert a growing vector to an immutable one without
-- copying.  After this call, you may not use the growing vector
-- nor any other growing vectors that were used to produce this one.
unsafeFreeze
    :: (PrimMonad m, V.Vector v a)
    => Growing v (PrimState m) a -> m (v a)
unsafeFreeze :: forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Growing v (PrimState m) a -> m (v a)
unsafeFreeze (Growing Int
len Mutable v (PrimState m) a
m) = forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
V.unsafeFreeze (forall (v :: * -> * -> *) a s. MVector v a => Int -> v s a -> v s a
MV.take Int
len Mutable v (PrimState m) a
m)

-- | Returns a new growing vector with a new element at the end.
-- Note that the return value may share storage with the input value.
-- Furthermore, calling @append@ twice on the same input may result
-- in two vectors that share the same storage.
append
    :: (PrimMonad m, V.Vector v a)
    => Growing v (PrimState m) a
    -> a
    -> m (Growing v (PrimState m) a)
append :: forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Growing v (PrimState m) a -> a -> m (Growing v (PrimState m) a)
append (Growing Int
len Mutable v (PrimState m) a
v) a
x
    | Int
len forall a. Ord a => a -> a -> Bool
< forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
MV.length Mutable v (PrimState m) a
v = do
        forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MV.unsafeWrite Mutable v (PrimState m) a
v Int
len a
x
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall (v :: * -> *) s a. Int -> Mutable v s a -> Growing v s a
Growing (Int
len forall a. Num a => a -> a -> a
+ Int
1) Mutable v (PrimState m) a
v
    | Bool
otherwise = do
        let len' :: Int
len' = Int
2 forall a. Num a => a -> a -> a
* Int
len forall a. Num a => a -> a -> a
+ Int
1
        Mutable v (PrimState m) a
v' <- forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m (v (PrimState m) a)
MV.unsafeGrow Mutable v (PrimState m) a
v Int
len'
        forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MV.unsafeWrite Mutable v (PrimState m) a
v' Int
len a
x
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall (v :: * -> *) s a. Int -> Mutable v s a -> Growing v s a
Growing (Int
len forall a. Num a => a -> a -> a
+ Int
1) Mutable v (PrimState m) a
v'
{-# INLINE append #-}