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

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

import qualified Data.Vector.Generic.Mutable as GMV

import Capnp.Errors (Error(SizeError))

-- | '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

-- | @'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