{-# 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 Capnp.Errors (Error (SizeError))
import Control.Monad (when)
import Control.Monad.Catch (MonadThrow (throwM))
import Control.Monad.Primitive (PrimMonad, PrimState)
import qualified Data.Vector.Generic.Mutable as GMV

-- | 'AppendVec' wraps a mutable vector, and affords amortized O(1) appending.
data AppendVec v s a = AppendVec
  { forall (v :: * -> * -> *) s a. AppendVec v s a -> v s a
mutVec :: v s a,
    forall (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 :: forall (v :: * -> * -> *) a s.
MVector v a =>
v s a -> AppendVec v s a
fromVector v s a
vec =
  AppendVec
    { mutVec :: v s a
mutVec = v s a
vec,
      mutVecLen :: Int
mutVecLen = 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 :: forall (v :: * -> * -> *) a s.
MVector v a =>
v s a -> AppendVec v s a
makeEmpty v s a
vec =
  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 :: forall (v :: * -> * -> *) a s.
MVector v a =>
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} = forall (v :: * -> * -> *) a s.
(HasCallStack, 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 :: forall (v :: * -> * -> *) a s.
MVector v a =>
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} = 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 :: forall (m :: * -> *) s (v :: * -> * -> *) a.
(MonadThrow m, PrimMonad m, s ~ PrimState m, MVector v a) =>
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
  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
maxSize forall a. Num a => a -> a -> a
- Int
amount forall a. Ord a => a -> a -> Bool
< Int
mutVecLen) forall a b. (a -> b) -> a -> b
$
    forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM Error
SizeError
  v s a
mutVec <-
    if 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.
        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:
        forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m (v (PrimState m) a)
GMV.grow v s a
mutVec (forall a. Ord a => a -> a -> a
max Int
amount (Int
mutVecLen forall a. Num a => a -> a -> a
* Int
2))
  forall (f :: * -> *) a. Applicative f => a -> f a
pure
    AppendVec
      { mutVec :: v s a
mutVec = v s a
mutVec,
        mutVecLen :: Int
mutVecLen = Int
mutVecLen forall a. Num a => a -> a -> a
+ Int
amount
      }

canGrowWithoutCopy :: (GMV.MVector v a) => AppendVec v s a -> Int -> Bool
canGrowWithoutCopy :: forall (v :: * -> * -> *) a s.
MVector v a =>
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 forall a. Num a => a -> a -> a
+ Int
amount forall a. Ord a => a -> a -> Bool
<= forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
GMV.length v s a
mutVec