content: {-# LANGUAGE NoImplicitPrelude #-}
  module RIO.Deque
      ( -- * Types
        Deque
      , UDeque
      , SDeque
      , BDeque
        -- * Operations
      , newDeque
      , getDequeSize
      , popFrontDeque
      , popBackDeque
      , pushFrontDeque
      , pushBackDeque
      , foldlDeque
      , foldrDeque
      , dequeToList
      , dequeToVector
      , freezeDeque
        -- * Inference helpers
      , asUDeque
      , asSDeque
      , asBDeque
      ) where
  
  import           RIO.Prelude.Reexports
  import           Control.Exception            (assert)
  import           Control.Monad                (liftM)
  import qualified Data.Vector.Generic          as VG
  import qualified Data.Vector.Generic.Mutable  as V
  import qualified Data.Vector.Mutable          as B
  import qualified Data.Vector.Storable.Mutable as S
  import qualified Data.Vector.Unboxed.Mutable  as U
  import           Data.Primitive.MutVar
  
  data DequeState v s a = DequeState
      !(v s a)
      {-# UNPACK #-} !Int -- start
      {-# UNPACK #-} !Int -- size
  
  -- | A double-ended queue supporting any underlying vector type and any monad.
  --
  -- This implements a circular double-ended queue with exponential growth.
  --
  -- @since 0.1.9.0
  newtype Deque v s a = Deque (MutVar s (DequeState v s a))
  
  -- | A 'Deque' specialized to unboxed vectors.
  --
  -- @since 0.1.9.0
  type UDeque = Deque U.MVector
  
  -- | A 'Deque' specialized to storable vectors.
  --
  -- @since 0.1.9.0
  type SDeque = Deque S.MVector
  
  -- | A 'Deque' specialized to boxed vectors.
  --
  -- @since 0.1.9.0
  type BDeque = Deque B.MVector
  
  -- | Helper function to assist with type inference, forcing usage of
  -- an unboxed vector.
  --
  -- @since 0.1.9.0
  asUDeque :: UDeque s a -> UDeque s a
  asUDeque = id
  
  -- | Helper function to assist with type inference, forcing usage of a
  -- storable vector.
  --
  -- @since 0.1.9.0
  asSDeque :: SDeque s a -> SDeque s a
  asSDeque = id
  
  -- | Helper function to assist with type inference, forcing usage of a
  -- boxed vector.
  --
  -- @since 0.1.9.0
  asBDeque :: BDeque s a -> BDeque s a
  asBDeque = id
  
  -- | Create a new, empty 'Deque'
  --
  -- @since 0.1.9.0
  newDeque
    :: (V.MVector v a, PrimMonad m)
    => m (Deque v (PrimState m) a)
  newDeque = do
      v <- V.new baseSize
      Deque <$> newMutVar (DequeState v 0 0)
    where
      baseSize = 32
  {-# INLINE newDeque #-}
  
  
  -- | /O(1)/ - Get the number of elements that is currently in the `Deque`
  --
  -- @since 0.1.9.0
  getDequeSize :: PrimMonad m => Deque v (PrimState m) a -> m Int
  getDequeSize (Deque var) = do
    DequeState _ _ size <- readMutVar var
    pure size
  {-# INLINE getDequeSize #-}
  
  
  -- | Pop the first value from the beginning of the 'Deque'
  --
  -- @since 0.1.9.0
  popFrontDeque
    :: (V.MVector v a, PrimMonad m)
    => Deque v (PrimState m) a
    -> m (Maybe a)
  popFrontDeque (Deque var) = do
      DequeState v start size <- readMutVar var
      if size == 0
          then return Nothing
          else do
              x <- V.unsafeRead v start
              let start' = start + 1
                  start''
                      | start' >= V.length v = 0
                      | otherwise = start'
              writeMutVar var $! DequeState v start'' (size - 1)
              return (Just x)
  {-# INLINE popFrontDeque #-}
  
  -- | Pop the first value from the end of the 'Deque'
  --
  -- @since 0.1.9.0
  popBackDeque
    :: (V.MVector v a, PrimMonad m)
    => Deque v (PrimState m) a
    -> m (Maybe a)
  popBackDeque (Deque var) = do
      DequeState v start size <- readMutVar var
      if size == 0
          then return Nothing
          else do
              let size' = size - 1
                  end = start + size'
                  end'
                      | end >= V.length v = end - V.length v
                      | otherwise = end
              x <- V.unsafeRead v end'
              writeMutVar var $! DequeState v start size'
              return (Just x)
  {-# INLINE popBackDeque #-}
  
  -- | Push a new value to the beginning of the 'Deque'
  --
  -- @since 0.1.9.0
  pushFrontDeque
    :: (V.MVector v a, PrimMonad m)
    => Deque v (PrimState m) a
    -> a
    -> m ()
  pushFrontDeque (Deque var) x = do
      DequeState v start size <- readMutVar var
      inner v start size
    where
      inner v start size = do
          if size >= V.length v
              then newVector v start size inner
              else do
                  let size' = size + 1
                      start' = (start - 1) `rem` V.length v
                      start''
                          | start' < 0 = V.length v + start'
                          | otherwise = start'
                  V.unsafeWrite v start'' x
                  writeMutVar var $! DequeState v start'' size'
  {-# INLINE pushFrontDeque #-}
  
  -- | Push a new value to the end of the 'Deque'
  --
  -- @since 0.1.9.0
  pushBackDeque
    :: (V.MVector v a, PrimMonad m)
    => Deque v (PrimState m) a
    -> a
    -> m ()
  pushBackDeque (Deque var) x = do
      DequeState v start size <- readMutVar var
      inner v start size
    where
      inner v start size = do
          if size >= V.length v
              then newVector v start size inner
              else do
                  let end = start + size
                      end'
                          | end >= V.length v = end - V.length v
                          | otherwise = end
                  V.unsafeWrite v end' x
                  writeMutVar var $! DequeState v start (size + 1)
  {-# INLINE pushBackDeque #-}
  
  -- | Fold over a 'Deque', starting at the beginning. Does not modify the 'Deque'.
  --
  -- @since 0.1.9.0
  foldlDeque
    :: (V.MVector v a, PrimMonad m)
    => (acc -> a -> m acc)
    -> acc
    -> Deque v (PrimState m) a
    -> m acc
  foldlDeque f acc0 (Deque var) = do
    DequeState v start size <- readMutVar var
    let loop idx acc
          | idx >= size = pure acc
          | otherwise = do
              let idxPlusStart = idx + start
                  idx'
                    | idxPlusStart >= V.length v = idxPlusStart - V.length v
                    | otherwise = idxPlusStart
              a <- V.unsafeRead v idx'
              acc' <- f acc a
              loop (idx + 1) $! acc'
    loop 0 acc0
  
  
  -- | Fold over a 'Deque', starting at the end. Does not modify the 'Deque'.
  --
  -- @since 0.1.9.0
  foldrDeque
    :: (V.MVector v a, PrimMonad m)
    => (a -> acc -> m acc)
    -> acc
    -> Deque v (PrimState m) a
    -> m acc
  foldrDeque f acc0 (Deque var) = do
    DequeState v start size <- readMutVar var
    let loop idx acc
          | idx < 0 = pure acc
          | otherwise = do
              let idxPlusStart = idx + start
                  idx'
                    | idxPlusStart >= V.length v = idxPlusStart - V.length v
                    | otherwise = idxPlusStart
              a <- V.unsafeRead v idx'
              acc' <- f a acc
              loop (idx - 1) $! acc'
    loop (size - 1) acc0
  
  -- | Convert a 'Deque' into a list. Does not modify the 'Deque'.
  --
  -- @since 0.1.9.0
  dequeToList
    :: (V.MVector v a, PrimMonad m)
    => Deque v (PrimState m) a
    -> m [a]
  dequeToList = foldrDeque (\a rest -> pure $ a : rest) []
  {-# INLINE dequeToList #-}
  
  
  -- | Convert to an immutable vector of any type. If resulting pure vector corresponds to the mutable
  -- one used by the `Deque`, it will be more efficient to use `freezeDeque` instead.
  --
  -- ==== __Example__
  --
  -- >>> :set -XTypeApplications
  -- >>> import qualified RIO.Vector.Unboxed as U
  -- >>> import qualified RIO.Vector.Storable as S
  -- >>> d <- newDeque @U.MVector @Int
  -- >>> mapM_ (pushFrontDeque d) [0..10]
  -- >>> dequeToVector @S.Vector d
  -- [10,9,8,7,6,5,4,3,2,1,0]
  --
  -- @since 0.1.9.0
  dequeToVector :: (VG.Vector v' a, V.MVector v a, PrimMonad m)
                => Deque v (PrimState m) a -> m (v' a)
  dequeToVector dq = do
      size <- getDequeSize dq
      mv <- V.unsafeNew size
      foldlDeque (\i e -> V.unsafeWrite mv i e >> pure (i+1)) 0 dq
      VG.unsafeFreeze mv
  
  
  newVector :: (PrimMonad m, V.MVector v a)
            => v (PrimState m) a
            -> Int
            -> Int
            -> (v (PrimState m) a -> Int -> Int -> m b)
            -> m b
  newVector v size2 sizeOrig f = assert (sizeOrig == V.length v) $ do
      v' <- V.unsafeNew (V.length v * 2)
      let size1 = V.length v - size2
      V.unsafeCopy
          (V.unsafeTake size1 v')
          (V.unsafeSlice size2 size1 v)
      V.unsafeCopy
          (V.unsafeSlice size1 size2 v')
          (V.unsafeTake size2 v)
      f v' 0 sizeOrig
  {-# INLINE newVector #-}
  
  
  -- | Yield an immutable copy of the underlying mutable vector. The difference from `dequeToVector`
  -- is that the the copy will be performed with a more efficient @memcpy@, rather than element by
  -- element. The downside is that the resulting vector type must be the one that corresponds to the
  -- mutable one that is used in the `Deque`.
  --
  -- ==== __Example__
  --
  -- >>> :set -XTypeApplications
  -- >>> import qualified RIO.Vector.Unboxed as U
  -- >>> d <- newDeque @U.MVector @Int
  -- >>> mapM_ (pushFrontDeque d) [0..10]
  -- >>> freezeDeque @U.Vector d
  -- [10,9,8,7,6,5,4,3,2,1,0]
  --
  -- @since 0.1.9.0
  freezeDeque ::
       (VG.Vector v a, PrimMonad m)
    => Deque (VG.Mutable v) (PrimState m) a
    -> m (v a)
  freezeDeque (Deque var) = do
      state@(DequeState v _ size) <- readMutVar var
      v' <- V.unsafeNew size
      makeCopy v' state
      VG.unsafeFreeze v'
  
  
  makeCopy ::
       (V.MVector v a, PrimMonad m)
    => v (PrimState m) a
    -> DequeState v (PrimState m) a
    -> m ()
  makeCopy v' (DequeState v start size) = do
      let size1 = min size (V.length v - start)
          size2 = size - size1
      V.unsafeCopy
          (V.unsafeTake size1 v')
          (V.unsafeSlice start size1 v)
      when (size > size1) $ V.unsafeCopy
          (V.unsafeSlice size1 size2 v')
          (V.unsafeTake size2 v)
  {-# INLINE makeCopy #-}
  