{-# LANGUAGE NamedFieldPuns #-} module Data.Buffer where import Control.Applicative (Applicative (..)) import Control.Exception (assert) import Control.Monad.Primitive import qualified Data.Vector.Unboxed as U import qualified Data.Vector.Unboxed.Mutable as UM import Prelude hiding (Applicative (..)) data Buffer s a = Buffer { forall s a. Buffer s a -> MVector s Int bufferVars :: !(UM.MVector s Int) , forall s a. Buffer s a -> MVector s a internalBuffer :: !(UM.MVector s a) , forall s a. Buffer s a -> Int internalBufferSize :: !Int } _bufferFrontPos :: Int _bufferFrontPos :: Int _bufferFrontPos = Int 0 _bufferBackPos :: Int _bufferBackPos :: Int _bufferBackPos = Int 1 newBuffer :: (U.Unbox a, PrimMonad m) => Int -> m (Buffer (PrimState m) a) newBuffer :: forall a (m :: * -> *). (Unbox a, PrimMonad m) => Int -> m (Buffer (PrimState m) a) newBuffer Int n = MVector (PrimState m) Int -> MVector (PrimState m) a -> Int -> Buffer (PrimState m) a forall s a. MVector s Int -> MVector s a -> Int -> Buffer s a Buffer (MVector (PrimState m) Int -> MVector (PrimState m) a -> Int -> Buffer (PrimState m) a) -> m (MVector (PrimState m) Int) -> m (MVector (PrimState m) a -> Int -> Buffer (PrimState m) a) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b <$> Int -> Int -> m (MVector (PrimState m) Int) forall (m :: * -> *) a. (PrimMonad m, Unbox a) => Int -> a -> m (MVector (PrimState m) a) UM.replicate Int 2 Int 0 m (MVector (PrimState m) a -> Int -> Buffer (PrimState m) a) -> m (MVector (PrimState m) a) -> m (Int -> Buffer (PrimState m) a) forall a b. m (a -> b) -> m a -> m b forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b <*> Int -> m (MVector (PrimState m) a) forall (m :: * -> *) a. (PrimMonad m, Unbox a) => Int -> m (MVector (PrimState m) a) UM.unsafeNew Int n m (Int -> Buffer (PrimState m) a) -> m Int -> m (Buffer (PrimState m) a) forall a b. m (a -> b) -> m a -> m b forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b <*> Int -> m Int forall a. a -> m a forall (f :: * -> *) a. Applicative f => a -> f a pure Int n type Stack s a = Buffer s a newBufferAsStack :: (U.Unbox a, PrimMonad m) => Int -> m (Buffer (PrimState m) a) newBufferAsStack :: forall a (m :: * -> *). (Unbox a, PrimMonad m) => Int -> m (Buffer (PrimState m) a) newBufferAsStack Int n = MVector (PrimState m) Int -> MVector (PrimState m) a -> Int -> Buffer (PrimState m) a forall s a. MVector s Int -> MVector s a -> Int -> Buffer s a Buffer (MVector (PrimState m) Int -> MVector (PrimState m) a -> Int -> Buffer (PrimState m) a) -> m (MVector (PrimState m) Int) -> m (MVector (PrimState m) a -> Int -> Buffer (PrimState m) a) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b <$> Int -> Int -> m (MVector (PrimState m) Int) forall (m :: * -> *) a. (PrimMonad m, Unbox a) => Int -> a -> m (MVector (PrimState m) a) UM.replicate Int 2 Int 0 m (MVector (PrimState m) a -> Int -> Buffer (PrimState m) a) -> m (MVector (PrimState m) a) -> m (Int -> Buffer (PrimState m) a) forall a b. m (a -> b) -> m a -> m b forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b <*> Int -> m (MVector (PrimState m) a) forall (m :: * -> *) a. (PrimMonad m, Unbox a) => Int -> m (MVector (PrimState m) a) UM.unsafeNew Int n m (Int -> Buffer (PrimState m) a) -> m Int -> m (Buffer (PrimState m) a) forall a b. m (a -> b) -> m a -> m b forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b <*> Int -> m Int forall a. a -> m a forall (f :: * -> *) a. Applicative f => a -> f a pure Int n type Queue s a = Buffer s a newBufferAsQueue :: (U.Unbox a, PrimMonad m) => Int -> m (Buffer (PrimState m) a) newBufferAsQueue :: forall a (m :: * -> *). (Unbox a, PrimMonad m) => Int -> m (Buffer (PrimState m) a) newBufferAsQueue Int n = MVector (PrimState m) Int -> MVector (PrimState m) a -> Int -> Buffer (PrimState m) a forall s a. MVector s Int -> MVector s a -> Int -> Buffer s a Buffer (MVector (PrimState m) Int -> MVector (PrimState m) a -> Int -> Buffer (PrimState m) a) -> m (MVector (PrimState m) Int) -> m (MVector (PrimState m) a -> Int -> Buffer (PrimState m) a) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b <$> Int -> Int -> m (MVector (PrimState m) Int) forall (m :: * -> *) a. (PrimMonad m, Unbox a) => Int -> a -> m (MVector (PrimState m) a) UM.replicate Int 2 Int 0 m (MVector (PrimState m) a -> Int -> Buffer (PrimState m) a) -> m (MVector (PrimState m) a) -> m (Int -> Buffer (PrimState m) a) forall a b. m (a -> b) -> m a -> m b forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b <*> Int -> m (MVector (PrimState m) a) forall (m :: * -> *) a. (PrimMonad m, Unbox a) => Int -> m (MVector (PrimState m) a) UM.unsafeNew Int n m (Int -> Buffer (PrimState m) a) -> m Int -> m (Buffer (PrimState m) a) forall a b. m (a -> b) -> m a -> m b forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b <*> Int -> m Int forall a. a -> m a forall (f :: * -> *) a. Applicative f => a -> f a pure Int n type Deque s a = Buffer s a newBufferAsDeque :: (U.Unbox a, PrimMonad m) => Int -> m (Buffer (PrimState m) a) newBufferAsDeque :: forall a (m :: * -> *). (Unbox a, PrimMonad m) => Int -> m (Buffer (PrimState m) a) newBufferAsDeque Int n = MVector (PrimState m) Int -> MVector (PrimState m) a -> Int -> Buffer (PrimState m) a forall s a. MVector s Int -> MVector s a -> Int -> Buffer s a Buffer (MVector (PrimState m) Int -> MVector (PrimState m) a -> Int -> Buffer (PrimState m) a) -> m (MVector (PrimState m) Int) -> m (MVector (PrimState m) a -> Int -> Buffer (PrimState m) a) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b <$> Int -> Int -> m (MVector (PrimState m) Int) forall (m :: * -> *) a. (PrimMonad m, Unbox a) => Int -> a -> m (MVector (PrimState m) a) UM.replicate Int 2 Int n m (MVector (PrimState m) a -> Int -> Buffer (PrimState m) a) -> m (MVector (PrimState m) a) -> m (Int -> Buffer (PrimState m) a) forall a b. m (a -> b) -> m a -> m b forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b <*> Int -> m (MVector (PrimState m) a) forall (m :: * -> *) a. (PrimMonad m, Unbox a) => Int -> m (MVector (PrimState m) a) UM.unsafeNew (Int 2 Int -> Int -> Int forall a. Num a => a -> a -> a * Int n) m (Int -> Buffer (PrimState m) a) -> m Int -> m (Buffer (PrimState m) a) forall a b. m (a -> b) -> m a -> m b forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b <*> Int -> m Int forall a. a -> m a forall (f :: * -> *) a. Applicative f => a -> f a pure (Int 2 Int -> Int -> Int forall a. Num a => a -> a -> a * Int n) lengthBuffer :: (PrimMonad m) => Buffer (PrimState m) a -> m Int lengthBuffer :: forall (m :: * -> *) a. PrimMonad m => Buffer (PrimState m) a -> m Int lengthBuffer Buffer{MVector (PrimState m) Int bufferVars :: forall s a. Buffer s a -> MVector s Int bufferVars :: MVector (PrimState m) Int bufferVars} = (Int -> Int -> Int) -> m Int -> m Int -> m Int forall a b c. (a -> b -> c) -> m a -> m b -> m c forall (f :: * -> *) a b c. Applicative f => (a -> b -> c) -> f a -> f b -> f c liftA2 (-) (MVector (PrimState m) Int -> Int -> m Int forall (m :: * -> *) a. (PrimMonad m, Unbox a) => MVector (PrimState m) a -> Int -> m a UM.unsafeRead MVector (PrimState m) Int bufferVars Int _bufferBackPos) (MVector (PrimState m) Int -> Int -> m Int forall (m :: * -> *) a. (PrimMonad m, Unbox a) => MVector (PrimState m) a -> Int -> m a UM.unsafeRead MVector (PrimState m) Int bufferVars Int _bufferFrontPos) {-# INLINE lengthBuffer #-} clearBuffer :: (PrimMonad m) => Buffer (PrimState m) a -> m () clearBuffer :: forall (m :: * -> *) a. PrimMonad m => Buffer (PrimState m) a -> m () clearBuffer Buffer{MVector (PrimState m) Int bufferVars :: forall s a. Buffer s a -> MVector s Int bufferVars :: MVector (PrimState m) Int bufferVars} = do MVector (PrimState m) Int -> Int -> Int -> m () forall (m :: * -> *) a. (PrimMonad m, Unbox a) => MVector (PrimState m) a -> Int -> a -> m () UM.unsafeWrite MVector (PrimState m) Int bufferVars Int _bufferFrontPos Int 0 MVector (PrimState m) Int -> Int -> Int -> m () forall (m :: * -> *) a. (PrimMonad m, Unbox a) => MVector (PrimState m) a -> Int -> a -> m () UM.unsafeWrite MVector (PrimState m) Int bufferVars Int _bufferBackPos Int 0 freezeBuffer :: (U.Unbox a, PrimMonad m) => Buffer (PrimState m) a -> m (U.Vector a) freezeBuffer :: forall a (m :: * -> *). (Unbox a, PrimMonad m) => Buffer (PrimState m) a -> m (Vector a) freezeBuffer Buffer{MVector (PrimState m) Int bufferVars :: forall s a. Buffer s a -> MVector s Int bufferVars :: MVector (PrimState m) Int bufferVars, MVector (PrimState m) a internalBuffer :: forall s a. Buffer s a -> MVector s a internalBuffer :: MVector (PrimState m) a internalBuffer} = do f <- MVector (PrimState m) Int -> Int -> m Int forall (m :: * -> *) a. (PrimMonad m, Unbox a) => MVector (PrimState m) a -> Int -> m a UM.unsafeRead MVector (PrimState m) Int bufferVars Int _bufferFrontPos b <- UM.unsafeRead bufferVars _bufferBackPos U.freeze $ UM.unsafeSlice f (b - f) internalBuffer unsafeFreezeBuffer :: (U.Unbox a, PrimMonad m) => Buffer (PrimState m) a -> m (U.Vector a) unsafeFreezeBuffer :: forall a (m :: * -> *). (Unbox a, PrimMonad m) => Buffer (PrimState m) a -> m (Vector a) unsafeFreezeBuffer Buffer{MVector (PrimState m) Int bufferVars :: forall s a. Buffer s a -> MVector s Int bufferVars :: MVector (PrimState m) Int bufferVars, MVector (PrimState m) a internalBuffer :: forall s a. Buffer s a -> MVector s a internalBuffer :: MVector (PrimState m) a internalBuffer} = do f <- MVector (PrimState m) Int -> Int -> m Int forall (m :: * -> *) a. (PrimMonad m, Unbox a) => MVector (PrimState m) a -> Int -> m a UM.unsafeRead MVector (PrimState m) Int bufferVars Int _bufferFrontPos b <- UM.unsafeRead bufferVars _bufferBackPos U.unsafeFreeze $ UM.unsafeSlice f (b - f) internalBuffer freezeInternalBuffer :: (U.Unbox a, PrimMonad m) => Buffer (PrimState m) a -> m (U.Vector a) freezeInternalBuffer :: forall a (m :: * -> *). (Unbox a, PrimMonad m) => Buffer (PrimState m) a -> m (Vector a) freezeInternalBuffer Buffer{MVector (PrimState m) Int bufferVars :: forall s a. Buffer s a -> MVector s Int bufferVars :: MVector (PrimState m) Int bufferVars, MVector (PrimState m) a internalBuffer :: forall s a. Buffer s a -> MVector s a internalBuffer :: MVector (PrimState m) a internalBuffer} = do b <- MVector (PrimState m) Int -> Int -> m Int forall (m :: * -> *) a. (PrimMonad m, Unbox a) => MVector (PrimState m) a -> Int -> m a UM.unsafeRead MVector (PrimState m) Int bufferVars Int _bufferBackPos U.freeze $ UM.unsafeSlice 0 b internalBuffer unsafeFreezeInternalBuffer :: (U.Unbox a, PrimMonad m) => Buffer (PrimState m) a -> m (U.Vector a) unsafeFreezeInternalBuffer :: forall a (m :: * -> *). (Unbox a, PrimMonad m) => Buffer (PrimState m) a -> m (Vector a) unsafeFreezeInternalBuffer Buffer{MVector (PrimState m) Int bufferVars :: forall s a. Buffer s a -> MVector s Int bufferVars :: MVector (PrimState m) Int bufferVars, MVector (PrimState m) a internalBuffer :: forall s a. Buffer s a -> MVector s a internalBuffer :: MVector (PrimState m) a internalBuffer} = do b <- MVector (PrimState m) Int -> Int -> m Int forall (m :: * -> *) a. (PrimMonad m, Unbox a) => MVector (PrimState m) a -> Int -> m a UM.unsafeRead MVector (PrimState m) Int bufferVars Int _bufferBackPos U.unsafeFreeze $ UM.unsafeSlice 0 b internalBuffer popFront :: (U.Unbox a, PrimMonad m) => Buffer (PrimState m) a -> m (Maybe a) popFront :: forall a (m :: * -> *). (Unbox a, PrimMonad m) => Buffer (PrimState m) a -> m (Maybe a) popFront Buffer{MVector (PrimState m) Int bufferVars :: forall s a. Buffer s a -> MVector s Int bufferVars :: MVector (PrimState m) Int bufferVars, MVector (PrimState m) a internalBuffer :: forall s a. Buffer s a -> MVector s a internalBuffer :: MVector (PrimState m) a internalBuffer} = do f <- MVector (PrimState m) Int -> Int -> m Int forall (m :: * -> *) a. (PrimMonad m, Unbox a) => MVector (PrimState m) a -> Int -> m a UM.unsafeRead MVector (PrimState m) Int bufferVars Int _bufferFrontPos b <- UM.unsafeRead bufferVars _bufferBackPos if f < b then do UM.unsafeWrite bufferVars _bufferFrontPos (f + 1) pure <$> UM.unsafeRead internalBuffer f else return Nothing {-# INLINE popFront #-} viewFront :: (U.Unbox a, PrimMonad m) => Buffer (PrimState m) a -> m (Maybe a) viewFront :: forall a (m :: * -> *). (Unbox a, PrimMonad m) => Buffer (PrimState m) a -> m (Maybe a) viewFront Buffer{MVector (PrimState m) Int bufferVars :: forall s a. Buffer s a -> MVector s Int bufferVars :: MVector (PrimState m) Int bufferVars, MVector (PrimState m) a internalBuffer :: forall s a. Buffer s a -> MVector s a internalBuffer :: MVector (PrimState m) a internalBuffer} = do f <- MVector (PrimState m) Int -> Int -> m Int forall (m :: * -> *) a. (PrimMonad m, Unbox a) => MVector (PrimState m) a -> Int -> m a UM.unsafeRead MVector (PrimState m) Int bufferVars Int _bufferFrontPos b <- UM.unsafeRead bufferVars _bufferBackPos if f < b then pure <$> UM.unsafeRead internalBuffer f else return Nothing {-# INLINE viewFront #-} popBack :: (U.Unbox a, PrimMonad m) => Buffer (PrimState m) a -> m (Maybe a) popBack :: forall a (m :: * -> *). (Unbox a, PrimMonad m) => Buffer (PrimState m) a -> m (Maybe a) popBack Buffer{MVector (PrimState m) Int bufferVars :: forall s a. Buffer s a -> MVector s Int bufferVars :: MVector (PrimState m) Int bufferVars, MVector (PrimState m) a internalBuffer :: forall s a. Buffer s a -> MVector s a internalBuffer :: MVector (PrimState m) a internalBuffer} = do f <- MVector (PrimState m) Int -> Int -> m Int forall (m :: * -> *) a. (PrimMonad m, Unbox a) => MVector (PrimState m) a -> Int -> m a UM.unsafeRead MVector (PrimState m) Int bufferVars Int _bufferFrontPos b <- UM.unsafeRead bufferVars _bufferBackPos if f < b then do UM.unsafeWrite bufferVars _bufferBackPos (b - 1) pure <$> UM.unsafeRead internalBuffer (b - 1) else return Nothing {-# INLINE popBack #-} viewBack :: (U.Unbox a, PrimMonad m) => Buffer (PrimState m) a -> m (Maybe a) viewBack :: forall a (m :: * -> *). (Unbox a, PrimMonad m) => Buffer (PrimState m) a -> m (Maybe a) viewBack Buffer{MVector (PrimState m) Int bufferVars :: forall s a. Buffer s a -> MVector s Int bufferVars :: MVector (PrimState m) Int bufferVars, MVector (PrimState m) a internalBuffer :: forall s a. Buffer s a -> MVector s a internalBuffer :: MVector (PrimState m) a internalBuffer} = do f <- MVector (PrimState m) Int -> Int -> m Int forall (m :: * -> *) a. (PrimMonad m, Unbox a) => MVector (PrimState m) a -> Int -> m a UM.unsafeRead MVector (PrimState m) Int bufferVars Int _bufferFrontPos b <- UM.unsafeRead bufferVars _bufferBackPos if f < b then pure <$> UM.unsafeRead internalBuffer (b - 1) else return Nothing {-# INLINE viewBack #-} pushFront :: (U.Unbox a, PrimMonad m) => a -> Buffer (PrimState m) a -> m () pushFront :: forall a (m :: * -> *). (Unbox a, PrimMonad m) => a -> Buffer (PrimState m) a -> m () pushFront a x Buffer{MVector (PrimState m) Int bufferVars :: forall s a. Buffer s a -> MVector s Int bufferVars :: MVector (PrimState m) Int bufferVars, MVector (PrimState m) a internalBuffer :: forall s a. Buffer s a -> MVector s a internalBuffer :: MVector (PrimState m) a internalBuffer} = do f <- MVector (PrimState m) Int -> Int -> m Int forall (m :: * -> *) a. (PrimMonad m, Unbox a) => MVector (PrimState m) a -> Int -> m a UM.unsafeRead MVector (PrimState m) Int bufferVars Int _bufferFrontPos UM.unsafeWrite bufferVars _bufferFrontPos (f - 1) assert (f > 0) $ do UM.unsafeWrite internalBuffer (f - 1) x {-# INLINE pushFront #-} pushBack :: (U.Unbox a, PrimMonad m) => a -> Buffer (PrimState m) a -> m () pushBack :: forall a (m :: * -> *). (Unbox a, PrimMonad m) => a -> Buffer (PrimState m) a -> m () pushBack a x Buffer{MVector (PrimState m) Int bufferVars :: forall s a. Buffer s a -> MVector s Int bufferVars :: MVector (PrimState m) Int bufferVars, MVector (PrimState m) a internalBuffer :: forall s a. Buffer s a -> MVector s a internalBuffer :: MVector (PrimState m) a internalBuffer, Int internalBufferSize :: forall s a. Buffer s a -> Int internalBufferSize :: Int internalBufferSize} = do b <- MVector (PrimState m) Int -> Int -> m Int forall (m :: * -> *) a. (PrimMonad m, Unbox a) => MVector (PrimState m) a -> Int -> m a UM.unsafeRead MVector (PrimState m) Int bufferVars Int _bufferBackPos UM.unsafeWrite bufferVars _bufferBackPos (b + 1) assert (b < internalBufferSize) $ do UM.unsafeWrite internalBuffer b x {-# INLINE pushBack #-} pushFronts :: (U.Unbox a, PrimMonad m) => U.Vector a -> Buffer (PrimState m) a -> m () pushFronts :: forall a (m :: * -> *). (Unbox a, PrimMonad m) => Vector a -> Buffer (PrimState m) a -> m () pushFronts Vector a vec Buffer{MVector (PrimState m) Int bufferVars :: forall s a. Buffer s a -> MVector s Int bufferVars :: MVector (PrimState m) Int bufferVars, MVector (PrimState m) a internalBuffer :: forall s a. Buffer s a -> MVector s a internalBuffer :: MVector (PrimState m) a internalBuffer} = do let n :: Int n = Vector a -> Int forall a. Unbox a => Vector a -> Int U.length Vector a vec f <- MVector (PrimState m) Int -> Int -> m Int forall (m :: * -> *) a. (PrimMonad m, Unbox a) => MVector (PrimState m) a -> Int -> m a UM.unsafeRead MVector (PrimState m) Int bufferVars Int _bufferFrontPos UM.unsafeWrite bufferVars _bufferFrontPos (f - n) assert (n <= f) $ do U.unsafeCopy (UM.unsafeSlice (f - n) n internalBuffer) vec {-# INLINE pushFronts #-} pushBacks :: (U.Unbox a, PrimMonad m) => U.Vector a -> Buffer (PrimState m) a -> m () pushBacks :: forall a (m :: * -> *). (Unbox a, PrimMonad m) => Vector a -> Buffer (PrimState m) a -> m () pushBacks Vector a vec Buffer{MVector (PrimState m) Int bufferVars :: forall s a. Buffer s a -> MVector s Int bufferVars :: MVector (PrimState m) Int bufferVars, MVector (PrimState m) a internalBuffer :: forall s a. Buffer s a -> MVector s a internalBuffer :: MVector (PrimState m) a internalBuffer, Int internalBufferSize :: forall s a. Buffer s a -> Int internalBufferSize :: Int internalBufferSize} = do let n :: Int n = Vector a -> Int forall a. Unbox a => Vector a -> Int U.length Vector a vec b <- MVector (PrimState m) Int -> Int -> m Int forall (m :: * -> *) a. (PrimMonad m, Unbox a) => MVector (PrimState m) a -> Int -> m a UM.unsafeRead MVector (PrimState m) Int bufferVars Int _bufferBackPos UM.unsafeWrite bufferVars _bufferBackPos (b + n) assert (b + n - 1 < internalBufferSize) $ do U.unsafeCopy (UM.unsafeSlice b n internalBuffer) vec {-# INLINE pushBacks #-}