{-# LANGUAGE RecordWildCards #-}

-- | Fixed-sized queue. Internally it has an \([l, r)\) pair of valid element bounds.
--
-- ==== __Example__
-- >>> import AtCoder.Internal.Queue qualified as Q
-- >>> que <- Q.new @_ @Int 3
-- >>> Q.capacity que
-- 3
--
-- >>> Q.pushBack que 0 -- [0  _  _  _]
-- >>> Q.pushBack que 1 -- [0, 1  _  _]
-- >>> Q.pushBack que 2 -- [0, 1, 2  _]
-- >>> Q.length que
-- 3
--
-- >>> Q.popFront que   -- [_  1, 2  _]
-- Just 0
--
-- >>> Q.freeze que     -- [_  1, 2  _]
-- [1,2]
--
-- >>> Q.pushFront que 10   -- [10, 1, 2  _]
-- >>> Q.pushFront que 1000
-- *** Exception: AtCoder.Internal.Queue.pushFront: no empty front space
-- ...
--
-- >>> Q.unsafeFreeze que -- [10, 1, 2  _]
-- [10,1,2]
--
-- >>> Q.clear que      -- [_  _  _  _]
-- >>> Q.pushBack que 0 -- [0  _  _  _]
-- >>> Q.pushBack que 1 -- [0, 1  _  _]
-- >>> Q.pushBack que 2 -- [0, 1, 2  _]
-- >>> Q.freeze que
-- [0,1,2]
--
-- @since 1.0.0.0
module AtCoder.Internal.Queue
  ( -- * Queue
    Queue,

    -- * Constructor
    new,

    -- * Metadata
    capacity,
    length,
    null,

    -- * Modifications

    -- ** Push/pop
    pushBack,
    pushFront,
    popFront,
    popFront_,

    -- ** Reset
    clear,

    -- * Conversions
    freeze,
    unsafeFreeze,
  )
where

import Control.Monad.Primitive (PrimMonad, PrimState)
import Data.Vector.Generic.Mutable qualified as VGM
import Data.Vector.Unboxed qualified as VU
import Data.Vector.Unboxed.Mutable qualified as VUM
import GHC.Stack (HasCallStack)
import Prelude hiding (length, null)

-- | Fixed-sized queue. Internally it has an \([l, r)\) pair of valid element bounds.
--
-- @since 1.0.0.0
data Queue s a = Queue
  { -- | Stores [l, r) range in the `vecQ`.
    forall s a. Queue s a -> MVector s Int
posQ :: !(VUM.MVector s Int),
    forall s a. Queue s a -> MVector s a
vecQ :: !(VUM.MVector s a)
  }

-- | \(O(n)\) Creates a `Queue` with capacity \(n\).
--
-- @since 1.0.0.0
{-# INLINE new #-}
new :: (PrimMonad m, VU.Unbox a) => Int -> m (Queue (PrimState m) a)
new :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (Queue (PrimState m) a)
new Int
n = do
  MVector (PrimState m) Int
posQ <- Int -> Int -> m (MVector (PrimState m) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
2 (Int
0 :: Int)
  MVector (PrimState m) a
vecQ <- Int -> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
n
  Queue (PrimState m) a -> m (Queue (PrimState m) a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Queue {MVector (PrimState m) a
MVector (PrimState m) Int
vecQ :: MVector (PrimState m) a
posQ :: MVector (PrimState m) Int
posQ :: MVector (PrimState m) Int
vecQ :: MVector (PrimState m) a
..}

-- | \(O(1)\) Returns the array size.
--
-- @since 1.0.0.0
{-# INLINE capacity #-}
capacity :: (VU.Unbox a) => Queue s a -> Int
capacity :: forall a s. Unbox a => Queue s a -> Int
capacity = MVector s a -> Int
forall a s. Unbox a => MVector s a -> Int
VUM.length (MVector s a -> Int)
-> (Queue s a -> MVector s a) -> Queue s a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Queue s a -> MVector s a
forall s a. Queue s a -> MVector s a
vecQ

-- | \(O(1)\) Returns the number of elements in the queue.
--
-- @since 1.0.0.0
{-# INLINE length #-}
length :: (PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> m Int
length :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Queue (PrimState m) a -> m Int
length Queue {MVector (PrimState m) a
MVector (PrimState m) Int
vecQ :: forall s a. Queue s a -> MVector s a
posQ :: forall s a. Queue s a -> MVector s Int
posQ :: MVector (PrimState m) Int
vecQ :: MVector (PrimState m) a
..} = do
  Int
l <- MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector (PrimState m) Int
posQ Int
0
  Int
r <- MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector (PrimState m) Int
posQ Int
1
  Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> m Int) -> Int -> m Int
forall a b. (a -> b) -> a -> b
$ Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l

-- | \(O(1)\) Returns `True` if the buffer is empty.
--
-- @since 1.0.0.0
{-# INLINE null #-}
null :: (PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> m Bool
null :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Queue (PrimState m) a -> m Bool
null = (Int -> Bool) -> m Int -> m Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
(<$>) (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) (m Int -> m Bool)
-> (Queue (PrimState m) a -> m Int)
-> Queue (PrimState m) a
-> m Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Queue (PrimState m) a -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Queue (PrimState m) a -> m Int
length

-- | \(O(1)\) Appends an element to the back. Will throw an exception if the index is out of range.
--
-- @since 1.0.0.0
{-# INLINE pushBack #-}
pushBack :: (HasCallStack, PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> a -> m ()
pushBack :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Queue (PrimState m) a -> a -> m ()
pushBack Queue {MVector (PrimState m) a
MVector (PrimState m) Int
vecQ :: forall s a. Queue s a -> MVector s a
posQ :: forall s a. Queue s a -> MVector s Int
posQ :: MVector (PrimState m) Int
vecQ :: MVector (PrimState m) a
..} a
e = do
  MVector (PrimState m) Int -> (Int -> m Int) -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> m a) -> Int -> m ()
VGM.unsafeModifyM
    MVector (PrimState m) Int
posQ
    ( \Int
r -> do
        MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState m) a
vecQ Int
r a
e
        Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> m Int) -> Int -> m Int
forall a b. (a -> b) -> a -> b
$ Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
    )
    Int
1

-- | \(O(1)\) Appends an element to the back. Will throw an exception if the index is out of range.
--
-- @since 1.0.0.0
{-# INLINE pushFront #-}
pushFront :: (HasCallStack, PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> a -> m ()
pushFront :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Queue (PrimState m) a -> a -> m ()
pushFront Queue {MVector (PrimState m) a
MVector (PrimState m) Int
vecQ :: forall s a. Queue s a -> MVector s a
posQ :: forall s a. Queue s a -> MVector s Int
posQ :: MVector (PrimState m) Int
vecQ :: MVector (PrimState m) a
..} a
e = do
  Int
l0 <- MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector (PrimState m) Int
posQ Int
0
  if Int
l0 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
    then [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"AtCoder.Internal.Queue.pushFront: no empty front space"
    else do
      MVector (PrimState m) Int -> (Int -> m Int) -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> m a) -> Int -> m ()
VGM.unsafeModifyM
        MVector (PrimState m) Int
posQ
        ( \Int
l -> do
            MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState m) a
vecQ (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) a
e
            Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> m Int) -> Int -> m Int
forall a b. (a -> b) -> a -> b
$ Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
        )
        Int
0

-- | \(O(1)\) Removes the first element from the queue and returns it, or `Nothing` if it is empty.
--
-- @since 1.0.0.0
{-# INLINE popFront #-}
popFront :: (PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> m (Maybe a)
popFront :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Queue (PrimState m) a -> m (Maybe a)
popFront Queue {MVector (PrimState m) a
MVector (PrimState m) Int
vecQ :: forall s a. Queue s a -> MVector s a
posQ :: forall s a. Queue s a -> MVector s Int
posQ :: MVector (PrimState m) Int
vecQ :: MVector (PrimState m) a
..} = do
  Int
l <- MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector (PrimState m) Int
posQ Int
0
  Int
r <- MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector (PrimState m) Int
posQ Int
1
  if Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
r
    then Maybe a -> m (Maybe a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing
    else do
      a
x <- MVector (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) a
vecQ Int
l
      MVector (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.unsafeWrite MVector (PrimState m) Int
posQ Int
0 (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
      Maybe a -> m (Maybe a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe a -> m (Maybe a)) -> Maybe a -> m (Maybe a)
forall a b. (a -> b) -> a -> b
$ a -> Maybe a
forall a. a -> Maybe a
Just a
x

-- | \(O(1)\) `popFront` with the return value discarded.
--
-- @since 1.0.0.0
{-# INLINE popFront_ #-}
popFront_ :: (PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> m ()
popFront_ :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Queue (PrimState m) a -> m ()
popFront_ Queue (PrimState m) a
que = do
  Maybe a
_ <- Queue (PrimState m) a -> m (Maybe a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Queue (PrimState m) a -> m (Maybe a)
popFront Queue (PrimState m) a
que
  () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | \(O(1)\) Sets the `length` to zero.
--
-- @since 1.0.0.0
{-# INLINE clear #-}
clear :: (PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> m ()
clear :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Queue (PrimState m) a -> m ()
clear Queue {MVector (PrimState m) a
MVector (PrimState m) Int
vecQ :: forall s a. Queue s a -> MVector s a
posQ :: forall s a. Queue s a -> MVector s Int
posQ :: MVector (PrimState m) Int
vecQ :: MVector (PrimState m) a
..} = do
  MVector (PrimState m) Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> a -> m ()
VGM.set MVector (PrimState m) Int
posQ Int
0

-- | \(O(n)\) Yields an immutable copy of the mutable vector.
--
-- @since 1.0.0.0
{-# INLINE freeze #-}
freeze :: (PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> m (VU.Vector a)
freeze :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Queue (PrimState m) a -> m (Vector a)
freeze Queue {MVector (PrimState m) a
MVector (PrimState m) Int
vecQ :: forall s a. Queue s a -> MVector s a
posQ :: forall s a. Queue s a -> MVector s Int
posQ :: MVector (PrimState m) Int
vecQ :: MVector (PrimState m) a
..} = do
  Int
l <- MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector (PrimState m) Int
posQ Int
0
  Int
r <- MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector (PrimState m) Int
posQ Int
1
  MVector (PrimState m) a -> m (Vector a)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.freeze (MVector (PrimState m) a -> m (Vector a))
-> MVector (PrimState m) a -> m (Vector a)
forall a b. (a -> b) -> a -> b
$ Int -> MVector (PrimState m) a -> MVector (PrimState m) a
forall a s. Unbox a => Int -> MVector s a -> MVector s a
VUM.take (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l) (MVector (PrimState m) a -> MVector (PrimState m) a)
-> MVector (PrimState m) a -> MVector (PrimState m) a
forall a b. (a -> b) -> a -> b
$ Int -> MVector (PrimState m) a -> MVector (PrimState m) a
forall a s. Unbox a => Int -> MVector s a -> MVector s a
VUM.drop Int
l MVector (PrimState m) a
vecQ

-- | \(O(1)\) Unsafely converts a mutable vector to an immutable one without copying. The mutable
-- vector may not be used after this operation.
--
-- @since 1.0.0.0
{-# INLINE unsafeFreeze #-}
unsafeFreeze :: (PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> m (VU.Vector a)
unsafeFreeze :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Queue (PrimState m) a -> m (Vector a)
unsafeFreeze Queue {MVector (PrimState m) a
MVector (PrimState m) Int
vecQ :: forall s a. Queue s a -> MVector s a
posQ :: forall s a. Queue s a -> MVector s Int
posQ :: MVector (PrimState m) Int
vecQ :: MVector (PrimState m) a
..} = do
  Int
l <- MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector (PrimState m) Int
posQ Int
0
  Int
r <- MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector (PrimState m) Int
posQ Int
1
  MVector (PrimState m) a -> m (Vector a)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze (MVector (PrimState m) a -> m (Vector a))
-> MVector (PrimState m) a -> m (Vector a)
forall a b. (a -> b) -> a -> b
$ Int -> MVector (PrimState m) a -> MVector (PrimState m) a
forall a s. Unbox a => Int -> MVector s a -> MVector s a
VUM.take (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l) (MVector (PrimState m) a -> MVector (PrimState m) a)
-> MVector (PrimState m) a -> MVector (PrimState m) a
forall a b. (a -> b) -> a -> b
$ Int -> MVector (PrimState m) a -> MVector (PrimState m) a
forall a s. Unbox a => Int -> MVector s a -> MVector s a
VUM.drop Int
l MVector (PrimState m) a
vecQ