module Data.Vector.Persistent.Internal.Buffer where import Control.Monad.Primitive import Data.Primitive.SmallArray import Data.Vector.Persistent.Internal.Array (shrinkSmallMutableArray_) import Prelude hiding (length) data Buffer s a = Buffer { forall s a. Buffer s a -> Int offset :: !Int, forall s a. Buffer s a -> SmallMutableArray s a marr :: !(SmallMutableArray s a) } new :: (PrimMonad m, s ~ PrimState m) => m (Buffer s a) new :: forall (m :: * -> *) s a. (PrimMonad m, s ~ PrimState m) => m (Buffer s a) new = do SmallMutableArray s a marr <- forall (m :: * -> *) a. PrimMonad m => Int -> a -> m (SmallMutableArray (PrimState m) a) newSmallArray Int 0 forall a. a undefinedElem forall (f :: * -> *) a. Applicative f => a -> f a pure Buffer {$sel:offset:Buffer :: Int offset = Int 0, SmallMutableArray s a marr :: SmallMutableArray s a $sel:marr:Buffer :: SmallMutableArray s a marr} {-# INLINE new #-} newWithCapacity :: (PrimMonad m, s ~ PrimState m) => Int -> m (Buffer s a) newWithCapacity :: forall (m :: * -> *) s a. (PrimMonad m, s ~ PrimState m) => Int -> m (Buffer s a) newWithCapacity Int cap = do SmallMutableArray s a marr <- forall (m :: * -> *) a. PrimMonad m => Int -> a -> m (SmallMutableArray (PrimState m) a) newSmallArray Int cap forall a. a undefinedElem forall (f :: * -> *) a. Applicative f => a -> f a pure Buffer {$sel:offset:Buffer :: Int offset = Int 0, SmallMutableArray s a marr :: SmallMutableArray s a $sel:marr:Buffer :: SmallMutableArray s a marr} {-# INLINE newWithCapacity #-} push :: (PrimMonad m, s ~ PrimState m) => a -> Buffer s a -> m (Buffer s a) push :: forall (m :: * -> *) s a. (PrimMonad m, s ~ PrimState m) => a -> Buffer s a -> m (Buffer s a) push a a Buffer s a buffer = do Buffer s a buffer' <- if forall s a. Buffer s a -> Int length Buffer s a buffer forall a. Eq a => a -> a -> Bool == forall s a. Buffer s a -> Int capacity Buffer s a buffer then forall (m :: * -> *) s a. (PrimMonad m, s ~ PrimState m) => Buffer s a -> m (Buffer s a) resize Buffer s a buffer else forall (f :: * -> *) a. Applicative f => a -> f a pure Buffer s a buffer forall (m :: * -> *) a. PrimMonad m => SmallMutableArray (PrimState m) a -> Int -> a -> m () writeSmallArray (forall s a. Buffer s a -> SmallMutableArray s a marr Buffer s a buffer') (forall s a. Buffer s a -> Int length Buffer s a buffer) a a forall (f :: * -> *) a. Applicative f => a -> f a pure Buffer s a buffer' {$sel:offset:Buffer :: Int offset = forall s a. Buffer s a -> Int offset Buffer s a buffer' forall a. Num a => a -> a -> a + Int 1} {-# INLINE push #-} read :: (PrimMonad m, s ~ PrimState m) => Int -> Buffer s a -> m a read :: forall (m :: * -> *) s a. (PrimMonad m, s ~ PrimState m) => Int -> Buffer s a -> m a read Int i Buffer {SmallMutableArray s a marr :: SmallMutableArray s a $sel:marr:Buffer :: forall s a. Buffer s a -> SmallMutableArray s a marr} = forall (m :: * -> *) a. PrimMonad m => SmallMutableArray (PrimState m) a -> Int -> m a readSmallArray SmallMutableArray s a marr Int i {-# INLINE read #-} write :: (PrimMonad m, s ~ PrimState m) => Int -> a -> Buffer s a -> m () write :: forall (m :: * -> *) s a. (PrimMonad m, s ~ PrimState m) => Int -> a -> Buffer s a -> m () write Int i a a Buffer {SmallMutableArray s a marr :: SmallMutableArray s a $sel:marr:Buffer :: forall s a. Buffer s a -> SmallMutableArray s a marr} = forall (m :: * -> *) a. PrimMonad m => SmallMutableArray (PrimState m) a -> Int -> a -> m () writeSmallArray SmallMutableArray s a marr Int i a a {-# INLINE write #-} clear :: Buffer s a -> Buffer s a clear :: forall s a. Buffer s a -> Buffer s a clear = forall s a. Int -> Buffer s a -> Buffer s a shrink Int 0 {-# INLINE clear #-} shrink :: Int -> Buffer s a -> Buffer s a shrink :: forall s a. Int -> Buffer s a -> Buffer s a shrink Int i Buffer s a buffer = Buffer s a buffer {$sel:offset:Buffer :: Int offset = Int i} {-# INLINE shrink #-} unsafeShrink :: (PrimMonad m, s ~ PrimState m) => Int -> Buffer s a -> m (Buffer s a) unsafeShrink :: forall (m :: * -> *) s a. (PrimMonad m, s ~ PrimState m) => Int -> Buffer s a -> m (Buffer s a) unsafeShrink Int i Buffer {SmallMutableArray s a marr :: SmallMutableArray s a $sel:marr:Buffer :: forall s a. Buffer s a -> SmallMutableArray s a marr} = do SmallMutableArray s a marr <- forall (m :: * -> *) a. PrimMonad m => MArray (PrimState m) a -> Int -> m (MArray (PrimState m) a) shrinkSmallMutableArray_ SmallMutableArray s a marr Int i forall (f :: * -> *) a. Applicative f => a -> f a pure Buffer {SmallMutableArray s a marr :: SmallMutableArray s a $sel:marr:Buffer :: SmallMutableArray s a marr, $sel:offset:Buffer :: Int offset = Int i} {-# INLINE unsafeShrink #-} capacity :: Buffer s a -> Int capacity :: forall s a. Buffer s a -> Int capacity Buffer {SmallMutableArray s a marr :: SmallMutableArray s a $sel:marr:Buffer :: forall s a. Buffer s a -> SmallMutableArray s a marr} = forall s a. SmallMutableArray s a -> Int sizeofSmallMutableArray SmallMutableArray s a marr {-# INLINE capacity #-} null :: Buffer s a -> Bool null :: forall s a. Buffer s a -> Bool null = (Int 0 forall a. Eq a => a -> a -> Bool ==) forall b c a. (b -> c) -> (a -> b) -> a -> c . forall s a. Buffer s a -> Int length length :: Buffer s a -> Int length :: forall s a. Buffer s a -> Int length = forall s a. Buffer s a -> Int offset {-# INLINE length #-} undefinedElem :: forall a. a undefinedElem :: forall a. a undefinedElem = forall a. HasCallStack => [Char] -> a error [Char] "undefined element" {-# NOINLINE undefinedElem #-} resize :: (PrimMonad m, s ~ PrimState m) => Buffer s a -> m (Buffer s a) resize :: forall (m :: * -> *) s a. (PrimMonad m, s ~ PrimState m) => Buffer s a -> m (Buffer s a) resize Buffer s a buffer = do if forall s a. Buffer s a -> Int capacity Buffer s a buffer forall a. Eq a => a -> a -> Bool == Int 0 then forall (m :: * -> *) s a. (PrimMonad m, s ~ PrimState m) => Int -> Buffer s a -> m (Buffer s a) grow Int 32 Buffer s a buffer else forall (m :: * -> *) s a. (PrimMonad m, s ~ PrimState m) => Int -> Buffer s a -> m (Buffer s a) grow (forall s a. Buffer s a -> Int capacity Buffer s a buffer) Buffer s a buffer {-# INLINE resize #-} grow :: (PrimMonad m, s ~ PrimState m) => Int -> Buffer s a -> m (Buffer s a) grow :: forall (m :: * -> *) s a. (PrimMonad m, s ~ PrimState m) => Int -> Buffer s a -> m (Buffer s a) grow Int more buffer :: Buffer s a buffer@Buffer {SmallMutableArray s a marr :: SmallMutableArray s a $sel:marr:Buffer :: forall s a. Buffer s a -> SmallMutableArray s a marr, Int offset :: Int $sel:offset:Buffer :: forall s a. Buffer s a -> Int offset} = do SmallMutableArray s a marr' <- forall (m :: * -> *) a. PrimMonad m => Int -> a -> m (SmallMutableArray (PrimState m) a) newSmallArray (forall s a. SmallMutableArray s a -> Int sizeofSmallMutableArray SmallMutableArray s a marr forall a. Num a => a -> a -> a + Int more) forall a. a undefinedElem forall (m :: * -> *) a. PrimMonad m => SmallMutableArray (PrimState m) a -> Int -> SmallMutableArray (PrimState m) a -> Int -> Int -> m () copySmallMutableArray SmallMutableArray s a marr' Int 0 SmallMutableArray s a marr Int 0 Int offset forall (f :: * -> *) a. Applicative f => a -> f a pure Buffer s a buffer {$sel:marr:Buffer :: SmallMutableArray s a marr = SmallMutableArray s a marr'} {-# INLINE grow #-} freeze :: (PrimMonad m, s ~ PrimState m) => Buffer s a -> m (SmallArray a) freeze :: forall (m :: * -> *) s a. (PrimMonad m, s ~ PrimState m) => Buffer s a -> m (SmallArray a) freeze Buffer {SmallMutableArray s a marr :: SmallMutableArray s a $sel:marr:Buffer :: forall s a. Buffer s a -> SmallMutableArray s a marr, Int offset :: Int $sel:offset:Buffer :: forall s a. Buffer s a -> Int offset} = forall (m :: * -> *) a. PrimMonad m => SmallMutableArray (PrimState m) a -> Int -> Int -> m (SmallArray a) freezeSmallArray SmallMutableArray s a marr Int 0 Int offset {-# INLINE freeze #-} unsafeFreeze :: (PrimMonad m, s ~ PrimState m) => Buffer s a -> m (SmallArray a) unsafeFreeze :: forall (m :: * -> *) s a. (PrimMonad m, s ~ PrimState m) => Buffer s a -> m (SmallArray a) unsafeFreeze Buffer {SmallMutableArray s a marr :: SmallMutableArray s a $sel:marr:Buffer :: forall s a. Buffer s a -> SmallMutableArray s a marr, Int offset :: Int $sel:offset:Buffer :: forall s a. Buffer s a -> Int offset} = do SmallMutableArray s a marr <- forall (m :: * -> *) a. PrimMonad m => MArray (PrimState m) a -> Int -> m (MArray (PrimState m) a) shrinkSmallMutableArray_ SmallMutableArray s a marr Int offset forall (m :: * -> *) a. PrimMonad m => SmallMutableArray (PrimState m) a -> m (SmallArray a) unsafeFreezeSmallArray SmallMutableArray s a marr {-# INLINE unsafeFreeze #-}