{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE TypeFamilies   #-}
{- |
Module: Internal.AppendVec
Description: Helpers for efficient appending to vectors.
-}
module Internal.AppendVec
    ( AppendVec
    , fromVector
    , makeEmpty
    , getVector
    , getCapacity
    , FrozenAppendVec(..)
    , grow
    , canGrowWithoutCopy
    ) where

import Control.Monad           (when)
import Control.Monad.Catch     (MonadThrow(throwM))
import Control.Monad.Primitive (PrimMonad, PrimState)

import qualified Data.Vector.Generic         as GV
import qualified Data.Vector.Generic.Mutable as GMV

import Capnp.Errors (Error(SizeError))
import Data.Mutable (Thaw (..))

-- | 'AppendVec' wraps a mutable vector, and affords amortized O(1) appending.
data AppendVec v s a = AppendVec
    { AppendVec v s a -> v s a
mutVec    :: v s a
    , AppendVec v s a -> Int
mutVecLen :: !Int
    }

-- | 'fromVector' wraps a mutable vector in an appendVector, with no initial
-- spare capacity.
fromVector :: GMV.MVector v a => v s a -> AppendVec v s a
fromVector :: v s a -> AppendVec v s a
fromVector v s a
vec = AppendVec :: forall (v :: * -> * -> *) s a. v s a -> Int -> AppendVec v s a
AppendVec
    { mutVec :: v s a
mutVec = v s a
vec
    , mutVecLen :: Int
mutVecLen = v s a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
GMV.length v s a
vec
    }

-- | 'makeEmpty' makes an initially empty 'AppendVec', using the argument
-- as allocation space for 'grow'.
makeEmpty :: GMV.MVector v a => v s a -> AppendVec v s a
makeEmpty :: v s a -> AppendVec v s a
makeEmpty v s a
vec = AppendVec :: forall (v :: * -> * -> *) s a. v s a -> Int -> AppendVec v s a
AppendVec
    { mutVec :: v s a
mutVec = v s a
vec
    , mutVecLen :: Int
mutVecLen = Int
0
    }

-- | 'getVector' returns the valid portion of the underlying mutable vector.
getVector :: GMV.MVector v a => AppendVec v s a -> v s a
getVector :: AppendVec v s a -> v s a
getVector AppendVec{v s a
mutVec :: v s a
mutVec :: forall (v :: * -> * -> *) s a. AppendVec v s a -> v s a
mutVec, Int
mutVecLen :: Int
mutVecLen :: forall (v :: * -> * -> *) s a. AppendVec v s a -> Int
mutVecLen} = Int -> Int -> v s a -> v s a
forall (v :: * -> * -> *) a s.
MVector v a =>
Int -> Int -> v s a -> v s a
GMV.slice Int
0 Int
mutVecLen v s a
mutVec

getCapacity :: GMV.MVector v a => AppendVec v s a -> Int
getCapacity :: AppendVec v s a -> Int
getCapacity AppendVec{v s a
mutVec :: v s a
mutVec :: forall (v :: * -> * -> *) s a. AppendVec v s a -> v s a
mutVec} = v s a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
GMV.length v s a
mutVec

-- | immutable version of 'AppendVec'; this is defined for the purpose of
-- implementing 'Thaw'.
newtype FrozenAppendVec v a = FrozenAppendVec { FrozenAppendVec v a -> v a
getFrozenVector :: v a }

instance GV.Vector v a => Thaw (FrozenAppendVec v a) where
    type Mutable s (FrozenAppendVec v a) = AppendVec (GV.Mutable v) s a

    thaw :: FrozenAppendVec v a -> m (Mutable s (FrozenAppendVec v a))
thaw         = (v a -> m (Mutable v s a))
-> FrozenAppendVec v a -> m (AppendVec (Mutable v) s a)
forall (m :: * -> *) (v :: * -> * -> *) a (v :: * -> *) a s.
(Monad m, MVector v a) =>
(v a -> m (v s a)) -> FrozenAppendVec v a -> m (AppendVec v s a)
thawAppend v a -> m (Mutable v s a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
v a -> m (Mutable v (PrimState m) a)
GV.thaw
    unsafeThaw :: FrozenAppendVec v a -> m (Mutable s (FrozenAppendVec v a))
unsafeThaw   = (v a -> m (Mutable v s a))
-> FrozenAppendVec v a -> m (AppendVec (Mutable v) s a)
forall (m :: * -> *) (v :: * -> * -> *) a (v :: * -> *) a s.
(Monad m, MVector v a) =>
(v a -> m (v s a)) -> FrozenAppendVec v a -> m (AppendVec v s a)
thawAppend v a -> m (Mutable v s a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
v a -> m (Mutable v (PrimState m) a)
GV.unsafeThaw
    freeze :: Mutable s (FrozenAppendVec v a) -> m (FrozenAppendVec v a)
freeze       = (Mutable v s a -> m (v a))
-> AppendVec (Mutable v) s a -> m (FrozenAppendVec v a)
forall (f :: * -> *) (v :: * -> * -> *) a s (v :: * -> *) a.
(Functor f, MVector v a) =>
(v s a -> f (v a)) -> AppendVec v s a -> f (FrozenAppendVec v a)
freezeAppend Mutable v s a -> m (v a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
GV.freeze
    unsafeFreeze :: Mutable s (FrozenAppendVec v a) -> m (FrozenAppendVec v a)
unsafeFreeze = (Mutable v s a -> m (v a))
-> AppendVec (Mutable v) s a -> m (FrozenAppendVec v a)
forall (f :: * -> *) (v :: * -> * -> *) a s (v :: * -> *) a.
(Functor f, MVector v a) =>
(v s a -> f (v a)) -> AppendVec v s a -> f (FrozenAppendVec v a)
freezeAppend Mutable v s a -> m (v a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
GV.unsafeFreeze

-- Helpers for the thaw & freeze methods for the instance above.
thawAppend :: (v a -> m (v s a)) -> FrozenAppendVec v a -> m (AppendVec v s a)
thawAppend v a -> m (v s a)
thaw FrozenAppendVec v a
frozen = do
    v s a
mvec <- v a -> m (v s a)
thaw (v a -> m (v s a)) -> v a -> m (v s a)
forall a b. (a -> b) -> a -> b
$ FrozenAppendVec v a -> v a
forall (v :: * -> *) a. FrozenAppendVec v a -> v a
getFrozenVector FrozenAppendVec v a
frozen
    AppendVec v s a -> m (AppendVec v s a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure AppendVec :: forall (v :: * -> * -> *) s a. v s a -> Int -> AppendVec v s a
AppendVec
        { mutVec :: v s a
mutVec = v s a
mvec
        , mutVecLen :: Int
mutVecLen = v s a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
GMV.length v s a
mvec
        }
freezeAppend :: (v s a -> f (v a)) -> AppendVec v s a -> f (FrozenAppendVec v a)
freezeAppend v s a -> f (v a)
freeze = (v a -> FrozenAppendVec v a) -> f (v a) -> f (FrozenAppendVec v a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap v a -> FrozenAppendVec v a
forall (v :: * -> *) a. v a -> FrozenAppendVec v a
FrozenAppendVec (f (v a) -> f (FrozenAppendVec v a))
-> (AppendVec v s a -> f (v a))
-> AppendVec v s a
-> f (FrozenAppendVec v a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v s a -> f (v a)
freeze (v s a -> f (v a))
-> (AppendVec v s a -> v s a) -> AppendVec v s a -> f (v a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AppendVec v s a -> v s a
forall (v :: * -> * -> *) a s.
MVector v a =>
AppendVec v s a -> v s a
getVector

-- | @'grow' vec amount maxSize@ grows the vector @vec@ by @amount@ elements,
-- provided the result does not exceed @maxSize@. Amortized O(@amount@). Returns
-- the new vector; the original should not be used.
-- .
-- If the result does exceed @maxSize@, throws 'SizeError'.
grow :: (MonadThrow m, PrimMonad m, s ~ PrimState m, GMV.MVector v a)
    => AppendVec v s a -> Int -> Int -> m (AppendVec v s a)
grow :: AppendVec v s a -> Int -> Int -> m (AppendVec v s a)
grow vec :: AppendVec v s a
vec@AppendVec{v s a
mutVec :: v s a
mutVec :: forall (v :: * -> * -> *) s a. AppendVec v s a -> v s a
mutVec,Int
mutVecLen :: Int
mutVecLen :: forall (v :: * -> * -> *) s a. AppendVec v s a -> Int
mutVecLen} Int
amount Int
maxSize = do
    Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
maxSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
amount Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
mutVecLen) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
        Error -> m ()
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM Error
SizeError
    v s a
mutVec <-
        if AppendVec v s a -> Int -> Bool
forall (v :: * -> * -> *) a s.
MVector v a =>
AppendVec v s a -> Int -> Bool
canGrowWithoutCopy AppendVec v s a
vec Int
amount then
            -- we have enough un-allocated space already; leave the vector
            -- itself alone.
            v s a -> m (v s a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure v s a
mutVec
        else
            -- Allocate some more space. we at least double the underlying
            -- vector's size, to make appending amortized O(1), but if the
            -- vector is small enough and the allocation is big enough, we
            -- may need to do more to satisfy the request:
            v (PrimState m) a -> Int -> m (v (PrimState m) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m (v (PrimState m) a)
GMV.grow v s a
v (PrimState m) a
mutVec (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
amount (Int
mutVecLen Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2))
    AppendVec v s a -> m (AppendVec v s a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure AppendVec :: forall (v :: * -> * -> *) s a. v s a -> Int -> AppendVec v s a
AppendVec
        { mutVec :: v s a
mutVec = v s a
mutVec
        , mutVecLen :: Int
mutVecLen = Int
mutVecLen Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
amount
        }

canGrowWithoutCopy :: (GMV.MVector v a) => AppendVec v s a -> Int -> Bool
canGrowWithoutCopy :: AppendVec v s a -> Int -> Bool
canGrowWithoutCopy AppendVec{v s a
mutVec :: v s a
mutVec :: forall (v :: * -> * -> *) s a. AppendVec v s a -> v s a
mutVec,Int
mutVecLen :: Int
mutVecLen :: forall (v :: * -> * -> *) s a. AppendVec v s a -> Int
mutVecLen} Int
amount =
    Int
mutVecLen Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
amount Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= v s a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
GMV.length v s a
mutVec