{-# LANGUAGE RecordWildCards #-}
module AtCoder.Internal.Queue
(
Queue,
new,
capacity,
length,
null,
pushBack,
pushFront,
popFront,
popFront_,
clear,
freeze,
unsafeFreeze,
)
where
import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)
import Control.Monad.ST (ST)
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)
data Queue s a = Queue
{
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)
}
{-# 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 = ST (PrimState m) (Queue (PrimState m) a)
-> m (Queue (PrimState m) a)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (Queue (PrimState m) a)
-> m (Queue (PrimState m) a))
-> ST (PrimState m) (Queue (PrimState m) a)
-> m (Queue (PrimState m) a)
forall a b. (a -> b) -> a -> b
$ Int -> ST (PrimState m) (Queue (PrimState m) a)
forall a s. Unbox a => Int -> ST s (Queue s a)
newST Int
n
{-# 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
{-# 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 (PrimState m) a
que = ST (PrimState m) Int -> m Int
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) Int -> m Int) -> ST (PrimState m) Int -> m Int
forall a b. (a -> b) -> a -> b
$ Queue (PrimState m) a -> ST (PrimState m) Int
forall a s. Unbox a => Queue s a -> ST s Int
lengthST Queue (PrimState m) a
que
{-# 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
{-# 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 (PrimState m) a
que a
e = ST (PrimState m) () -> m ()
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) () -> m ()) -> ST (PrimState m) () -> m ()
forall a b. (a -> b) -> a -> b
$ Queue (PrimState m) a -> a -> ST (PrimState m) ()
forall a s. (HasCallStack, Unbox a) => Queue s a -> a -> ST s ()
pushBackST Queue (PrimState m) a
que a
e
{-# 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 (PrimState m) a
que a
e = ST (PrimState m) () -> m ()
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) () -> m ()) -> ST (PrimState m) () -> m ()
forall a b. (a -> b) -> a -> b
$ Queue (PrimState m) a -> a -> ST (PrimState m) ()
forall a s. (HasCallStack, Unbox a) => Queue s a -> a -> ST s ()
pushFrontST Queue (PrimState m) a
que a
e
{-# 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 (PrimState m) a
que = ST (PrimState m) (Maybe a) -> m (Maybe a)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (Maybe a) -> m (Maybe a))
-> ST (PrimState m) (Maybe a) -> m (Maybe a)
forall a b. (a -> b) -> a -> b
$ Queue (PrimState m) a -> ST (PrimState m) (Maybe a)
forall a s. Unbox a => Queue s a -> ST s (Maybe a)
popFrontST Queue (PrimState m) a
que
{-# 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 = ST (PrimState m) () -> m ()
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) () -> m ()) -> ST (PrimState m) () -> m ()
forall a b. (a -> b) -> a -> b
$ Queue (PrimState m) a -> ST (PrimState m) ()
forall a s. Unbox a => Queue s a -> ST s ()
popFrontST_ Queue (PrimState m) a
que
{-# 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
{-# 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 (PrimState m) a
que = ST (PrimState m) (Vector a) -> m (Vector a)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (Vector a) -> m (Vector a))
-> ST (PrimState m) (Vector a) -> m (Vector a)
forall a b. (a -> b) -> a -> b
$ Queue (PrimState m) a -> ST (PrimState m) (Vector a)
forall a s. Unbox a => Queue s a -> ST s (Vector a)
freezeST Queue (PrimState m) a
que
{-# 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 (PrimState m) a
que = ST (PrimState m) (Vector a) -> m (Vector a)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (Vector a) -> m (Vector a))
-> ST (PrimState m) (Vector a) -> m (Vector a)
forall a b. (a -> b) -> a -> b
$ Queue (PrimState m) a -> ST (PrimState m) (Vector a)
forall a s. Unbox a => Queue s a -> ST s (Vector a)
unsafeFreezeST Queue (PrimState m) a
que
{-# INLINEABLE newST #-}
newST :: (VU.Unbox a) => Int -> ST s (Queue s a)
newST :: forall a s. Unbox a => Int -> ST s (Queue s a)
newST Int
n = do
MVector s Int
posQ <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
2 (Int
0 :: Int)
MVector s a
vecQ <- Int -> ST s (MVector (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
n
Queue s a -> ST s (Queue s a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Queue {MVector s a
MVector s Int
vecQ :: MVector s a
posQ :: MVector s Int
posQ :: MVector s Int
vecQ :: MVector s a
..}
{-# INLINEABLE lengthST #-}
lengthST :: (VU.Unbox a) => Queue s a -> ST s Int
lengthST :: forall a s. Unbox a => Queue s a -> ST s Int
lengthST Queue {MVector s a
MVector s Int
vecQ :: forall s a. Queue s a -> MVector s a
posQ :: forall s a. Queue s a -> MVector s Int
posQ :: MVector s Int
vecQ :: MVector s a
..} = do
Int
l <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector s Int
MVector (PrimState (ST s)) Int
posQ Int
0
Int
r <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector s Int
MVector (PrimState (ST s)) Int
posQ Int
1
Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> ST s Int) -> Int -> ST s Int
forall a b. (a -> b) -> a -> b
$ Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l
{-# INLINEABLE pushBackST #-}
pushBackST :: (HasCallStack, VU.Unbox a) => Queue s a -> a -> ST s ()
pushBackST :: forall a s. (HasCallStack, Unbox a) => Queue s a -> a -> ST s ()
pushBackST Queue {MVector s a
MVector s Int
vecQ :: forall s a. Queue s a -> MVector s a
posQ :: forall s a. Queue s a -> MVector s Int
posQ :: MVector s Int
vecQ :: MVector s a
..} a
e = do
MVector (PrimState (ST s)) Int
-> (Int -> ST s Int) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> m a) -> Int -> m ()
VGM.unsafeModifyM
MVector s Int
MVector (PrimState (ST s)) Int
posQ
( \Int
r -> do
MVector (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s a
MVector (PrimState (ST s)) a
vecQ Int
r a
e
Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> ST s Int) -> Int -> ST s Int
forall a b. (a -> b) -> a -> b
$ Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
)
Int
1
{-# INLINEABLE pushFrontST #-}
pushFrontST :: (HasCallStack, VU.Unbox a) => Queue s a -> a -> ST s ()
pushFrontST :: forall a s. (HasCallStack, Unbox a) => Queue s a -> a -> ST s ()
pushFrontST Queue {MVector s a
MVector s Int
vecQ :: forall s a. Queue s a -> MVector s a
posQ :: forall s a. Queue s a -> MVector s Int
posQ :: MVector s Int
vecQ :: MVector s a
..} a
e = do
Int
l0 <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector s Int
MVector (PrimState (ST s)) Int
posQ Int
0
if Int
l0 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
then [Char] -> ST s ()
forall a. HasCallStack => [Char] -> a
error [Char]
"AtCoder.Internal.Queue.pushFront: no empty front space"
else do
MVector (PrimState (ST s)) Int
-> (Int -> ST s Int) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> m a) -> Int -> m ()
VGM.unsafeModifyM
MVector s Int
MVector (PrimState (ST s)) Int
posQ
( \Int
l -> do
MVector (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s a
MVector (PrimState (ST s)) a
vecQ (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) a
e
Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> ST s Int) -> Int -> ST s Int
forall a b. (a -> b) -> a -> b
$ Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
)
Int
0
{-# INLINEABLE popFrontST #-}
popFrontST :: (VU.Unbox a) => Queue s a -> ST s (Maybe a)
popFrontST :: forall a s. Unbox a => Queue s a -> ST s (Maybe a)
popFrontST Queue {MVector s a
MVector s Int
vecQ :: forall s a. Queue s a -> MVector s a
posQ :: forall s a. Queue s a -> MVector s Int
posQ :: MVector s Int
vecQ :: MVector s a
..} = do
Int
l <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector s Int
MVector (PrimState (ST s)) Int
posQ Int
0
Int
r <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector s Int
MVector (PrimState (ST s)) Int
posQ Int
1
if Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
r
then Maybe a -> ST s (Maybe a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing
else do
a
x <- MVector (PrimState (ST s)) a -> Int -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s a
MVector (PrimState (ST s)) a
vecQ Int
l
MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
posQ Int
0 (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
Maybe a -> ST s (Maybe a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe a -> ST s (Maybe a)) -> Maybe a -> ST s (Maybe a)
forall a b. (a -> b) -> a -> b
$ a -> Maybe a
forall a. a -> Maybe a
Just a
x
{-# INLINEABLE popFrontST_ #-}
popFrontST_ :: (VU.Unbox a) => Queue s a -> ST s ()
popFrontST_ :: forall a s. Unbox a => Queue s a -> ST s ()
popFrontST_ Queue s a
que = do
Maybe a
_ <- Queue (PrimState (ST s)) a -> ST s (Maybe a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Queue (PrimState m) a -> m (Maybe a)
popFront Queue s a
Queue (PrimState (ST s)) a
que
() -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
{-# INLINEABLE clearST #-}
clearST :: (VU.Unbox a) => Queue s a -> ST s ()
clearST :: forall a s. Unbox a => Queue s a -> ST s ()
clearST Queue {MVector s a
MVector s Int
vecQ :: forall s a. Queue s a -> MVector s a
posQ :: forall s a. Queue s a -> MVector s Int
posQ :: MVector s Int
vecQ :: MVector s a
..} = do
MVector (PrimState (ST s)) Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> a -> m ()
VGM.set MVector s Int
MVector (PrimState (ST s)) Int
posQ Int
0
{-# INLINEABLE freezeST #-}
freezeST :: (VU.Unbox a) => Queue s a -> ST s (VU.Vector a)
freezeST :: forall a s. Unbox a => Queue s a -> ST s (Vector a)
freezeST Queue {MVector s a
MVector s Int
vecQ :: forall s a. Queue s a -> MVector s a
posQ :: forall s a. Queue s a -> MVector s Int
posQ :: MVector s Int
vecQ :: MVector s a
..} = do
Int
l <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector s Int
MVector (PrimState (ST s)) Int
posQ Int
0
Int
r <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector s Int
MVector (PrimState (ST s)) Int
posQ Int
1
MVector (PrimState (ST s)) a -> ST s (Vector a)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.freeze (MVector (PrimState (ST s)) a -> ST s (Vector a))
-> MVector (PrimState (ST s)) a -> ST s (Vector a)
forall a b. (a -> b) -> a -> b
$ Int -> MVector (PrimState (ST s)) a -> MVector (PrimState (ST s)) 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 (ST s)) a -> MVector (PrimState (ST s)) a)
-> MVector (PrimState (ST s)) a -> MVector (PrimState (ST s)) a
forall a b. (a -> b) -> a -> b
$ Int -> MVector s a -> MVector s a
forall a s. Unbox a => Int -> MVector s a -> MVector s a
VUM.drop Int
l MVector s a
vecQ
{-# INLINEABLE unsafeFreezeST #-}
unsafeFreezeST :: (VU.Unbox a) => Queue s a -> ST s (VU.Vector a)
unsafeFreezeST :: forall a s. Unbox a => Queue s a -> ST s (Vector a)
unsafeFreezeST Queue {MVector s a
MVector s Int
vecQ :: forall s a. Queue s a -> MVector s a
posQ :: forall s a. Queue s a -> MVector s Int
posQ :: MVector s Int
vecQ :: MVector s a
..} = do
Int
l <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector s Int
MVector (PrimState (ST s)) Int
posQ Int
0
Int
r <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector s Int
MVector (PrimState (ST s)) Int
posQ Int
1
MVector (PrimState (ST s)) a -> ST s (Vector a)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze (MVector (PrimState (ST s)) a -> ST s (Vector a))
-> MVector (PrimState (ST s)) a -> ST s (Vector a)
forall a b. (a -> b) -> a -> b
$ Int -> MVector (PrimState (ST s)) a -> MVector (PrimState (ST s)) 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 (ST s)) a -> MVector (PrimState (ST s)) a)
-> MVector (PrimState (ST s)) a -> MVector (PrimState (ST s)) a
forall a b. (a -> b) -> a -> b
$ Int -> MVector s a -> MVector s a
forall a s. Unbox a => Int -> MVector s a -> MVector s a
VUM.drop Int
l MVector s a
vecQ