{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeFamilies #-}

-- | Fixed-sized array for \(O(1)\) allocation and \(O(1)\) clearing after \(O(n)\) construction.
module AtCoder.Extra.Pool
  ( -- * Pool
    Pool (..),
    Index (..),
    undefIndex,
    nullIndex,

    -- * Constructors
    new,
    clear,

    -- * Metadata
    capacity,
    size,

    -- * Allocations
    alloc,
    free,

    -- * Read/write
    read,
    write,
    modify,
    exchange,

    -- * Handle
    Handle (..),
    newHandle,
    nullHandle,
    invalidateHandle,
  )
where

import AtCoder.Internal.Buffer qualified as B
import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)
import Control.Monad.ST (ST)
import Data.Coerce
import Data.Vector.Generic qualified as VG
import Data.Vector.Generic.Mutable qualified as VGM
import Data.Vector.Primitive qualified as VP
import Data.Vector.Unboxed qualified as VU
import Data.Vector.Unboxed.Mutable qualified as VUM
import GHC.Stack (HasCallStack)
import Prelude hiding (read)

-- | Fixed-sized array for \(O(1)\) allocation and \(O(1)\) clearing after \(O(n)\) construction.
data Pool s a = Pool
  { -- | Data array.
    forall s a. Pool s a -> MVector s a
dataPool :: !(VUM.MVector s a),
    -- | Free slot indices pushed on free.
    forall s a. Pool s a -> Buffer s Index
freePool :: !(B.Buffer s Index),
    -- | Next index when `freePool` is empty.
    forall s a. Pool s a -> MVector s Index
nextPool :: !(VUM.MVector s Index)
  }

-- | Strongly typed index of pool items. User has to explicitly @corece@ on raw index use.
newtype Index = Index {Index -> Int
unIndex :: Int}
  deriving (Index -> Index -> Bool
(Index -> Index -> Bool) -> (Index -> Index -> Bool) -> Eq Index
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Index -> Index -> Bool
== :: Index -> Index -> Bool
$c/= :: Index -> Index -> Bool
/= :: Index -> Index -> Bool
Eq, Addr# -> Int# -> Index
ByteArray# -> Int# -> Index
Proxy Index -> Int#
Index -> Int#
(Proxy Index -> Int#)
-> (Index -> Int#)
-> (Proxy Index -> Int#)
-> (Index -> Int#)
-> (ByteArray# -> Int# -> Index)
-> (forall s.
    MutableByteArray# s -> Int# -> State# s -> (# State# s, Index #))
-> (forall s.
    MutableByteArray# s -> Int# -> Index -> State# s -> State# s)
-> (forall s.
    MutableByteArray# s
    -> Int# -> Int# -> Index -> State# s -> State# s)
-> (Addr# -> Int# -> Index)
-> (forall s. Addr# -> Int# -> State# s -> (# State# s, Index #))
-> (forall s. Addr# -> Int# -> Index -> State# s -> State# s)
-> (forall s.
    Addr# -> Int# -> Int# -> Index -> State# s -> State# s)
-> Prim Index
forall s. Addr# -> Int# -> Int# -> Index -> State# s -> State# s
forall s. Addr# -> Int# -> State# s -> (# State# s, Index #)
forall s. Addr# -> Int# -> Index -> State# s -> State# s
forall s.
MutableByteArray# s
-> Int# -> Int# -> Index -> State# s -> State# s
forall s.
MutableByteArray# s -> Int# -> State# s -> (# State# s, Index #)
forall s.
MutableByteArray# s -> Int# -> Index -> State# s -> State# s
forall a.
(Proxy a -> Int#)
-> (a -> Int#)
-> (Proxy a -> Int#)
-> (a -> Int#)
-> (ByteArray# -> Int# -> a)
-> (forall s.
    MutableByteArray# s -> Int# -> State# s -> (# State# s, a #))
-> (forall s.
    MutableByteArray# s -> Int# -> a -> State# s -> State# s)
-> (forall s.
    MutableByteArray# s -> Int# -> Int# -> a -> State# s -> State# s)
-> (Addr# -> Int# -> a)
-> (forall s. Addr# -> Int# -> State# s -> (# State# s, a #))
-> (forall s. Addr# -> Int# -> a -> State# s -> State# s)
-> (forall s. Addr# -> Int# -> Int# -> a -> State# s -> State# s)
-> Prim a
$csizeOfType# :: Proxy Index -> Int#
sizeOfType# :: Proxy Index -> Int#
$csizeOf# :: Index -> Int#
sizeOf# :: Index -> Int#
$calignmentOfType# :: Proxy Index -> Int#
alignmentOfType# :: Proxy Index -> Int#
$calignment# :: Index -> Int#
alignment# :: Index -> Int#
$cindexByteArray# :: ByteArray# -> Int# -> Index
indexByteArray# :: ByteArray# -> Int# -> Index
$creadByteArray# :: forall s.
MutableByteArray# s -> Int# -> State# s -> (# State# s, Index #)
readByteArray# :: forall s.
MutableByteArray# s -> Int# -> State# s -> (# State# s, Index #)
$cwriteByteArray# :: forall s.
MutableByteArray# s -> Int# -> Index -> State# s -> State# s
writeByteArray# :: forall s.
MutableByteArray# s -> Int# -> Index -> State# s -> State# s
$csetByteArray# :: forall s.
MutableByteArray# s
-> Int# -> Int# -> Index -> State# s -> State# s
setByteArray# :: forall s.
MutableByteArray# s
-> Int# -> Int# -> Index -> State# s -> State# s
$cindexOffAddr# :: Addr# -> Int# -> Index
indexOffAddr# :: Addr# -> Int# -> Index
$creadOffAddr# :: forall s. Addr# -> Int# -> State# s -> (# State# s, Index #)
readOffAddr# :: forall s. Addr# -> Int# -> State# s -> (# State# s, Index #)
$cwriteOffAddr# :: forall s. Addr# -> Int# -> Index -> State# s -> State# s
writeOffAddr# :: forall s. Addr# -> Int# -> Index -> State# s -> State# s
$csetOffAddr# :: forall s. Addr# -> Int# -> Int# -> Index -> State# s -> State# s
setOffAddr# :: forall s. Addr# -> Int# -> Int# -> Index -> State# s -> State# s
VP.Prim)
  deriving newtype (Eq Index
Eq Index =>
(Index -> Index -> Ordering)
-> (Index -> Index -> Bool)
-> (Index -> Index -> Bool)
-> (Index -> Index -> Bool)
-> (Index -> Index -> Bool)
-> (Index -> Index -> Index)
-> (Index -> Index -> Index)
-> Ord Index
Index -> Index -> Bool
Index -> Index -> Ordering
Index -> Index -> Index
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Index -> Index -> Ordering
compare :: Index -> Index -> Ordering
$c< :: Index -> Index -> Bool
< :: Index -> Index -> Bool
$c<= :: Index -> Index -> Bool
<= :: Index -> Index -> Bool
$c> :: Index -> Index -> Bool
> :: Index -> Index -> Bool
$c>= :: Index -> Index -> Bool
>= :: Index -> Index -> Bool
$cmax :: Index -> Index -> Index
max :: Index -> Index -> Index
$cmin :: Index -> Index -> Index
min :: Index -> Index -> Index
Ord, Int -> Index -> ShowS
[Index] -> ShowS
Index -> String
(Int -> Index -> ShowS)
-> (Index -> String) -> ([Index] -> ShowS) -> Show Index
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Index -> ShowS
showsPrec :: Int -> Index -> ShowS
$cshow :: Index -> String
show :: Index -> String
$cshowList :: [Index] -> ShowS
showList :: [Index] -> ShowS
Show)

newtype instance VU.MVector s Index = MV_Index (VP.MVector s Index)

newtype instance VU.Vector Index = V_Index (VP.Vector Index)

deriving via (VU.UnboxViaPrim Index) instance VGM.MVector VUM.MVector Index

deriving via (VU.UnboxViaPrim Index) instance VG.Vector VU.Vector Index

instance VU.Unbox Index

-- | Invalid, null `Index`.
{-# INLINE undefIndex #-}
undefIndex :: Index
undefIndex :: Index
undefIndex = Int -> Index
Index (-Int
1)

-- | Returns `True` for `undefIndex`.
{-# INLINE nullIndex #-}
nullIndex :: Index -> Bool
nullIndex :: Index -> Bool
nullIndex = (Index -> Index -> Bool
forall a. Eq a => a -> a -> Bool
== Index
undefIndex)

-- | \(O(n)\) Creates a pool with the specified @capacity@.
{-# INLINE new #-}
new :: (VU.Unbox a, PrimMonad m) => Int -> m (Pool (PrimState m) a)
new :: forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Int -> m (Pool (PrimState m) a)
new Int
cap = ST (PrimState m) (Pool (PrimState m) a) -> m (Pool (PrimState m) a)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (Pool (PrimState m) a)
 -> m (Pool (PrimState m) a))
-> ST (PrimState m) (Pool (PrimState m) a)
-> m (Pool (PrimState m) a)
forall a b. (a -> b) -> a -> b
$ Int -> ST (PrimState m) (Pool (PrimState m) a)
forall a s. Unbox a => Int -> ST s (Pool s a)
newST Int
cap

-- | \(O(1)\) Resets the pool to the initial state.
{-# INLINE clear #-}
clear :: (PrimMonad m) => Pool (PrimState m) a -> m ()
clear :: forall (m :: * -> *) a. PrimMonad m => Pool (PrimState m) a -> m ()
clear Pool (PrimState m) a
pool = 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
$ Pool (PrimState m) a -> ST (PrimState m) ()
forall s a. Pool s a -> ST s ()
clearST Pool (PrimState m) a
pool

-- | \(O(1)\) Returns the maximum number of elements the pool can store.
{-# INLINE capacity #-}
capacity :: (VU.Unbox a) => Pool s a -> Int
capacity :: forall a s. Unbox a => Pool s a -> Int
capacity = MVector s a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
VGM.length (MVector s a -> Int)
-> (Pool s a -> MVector s a) -> Pool s a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pool s a -> MVector s a
forall s a. Pool s a -> MVector s a
dataPool

-- | \(O(1)\) Returns the number of elements in the pool.
{-# INLINE size #-}
size :: (PrimMonad m, VU.Unbox a) => Pool (PrimState m) a -> m Int
size :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Pool (PrimState m) a -> m Int
size Pool (PrimState m) a
pool = 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
$ Pool (PrimState m) a -> ST (PrimState m) Int
forall a s. Unbox a => Pool s a -> ST s Int
sizeST Pool (PrimState m) a
pool

-- | \(O(1)\) Allocates a new element.
--
-- ==== Constraints
-- - The number of elements must not exceed the `capacity`.
{-# INLINE alloc #-}
alloc :: (HasCallStack, PrimMonad m, VU.Unbox a) => Pool (PrimState m) a -> a -> m Index
alloc :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Pool (PrimState m) a -> a -> m Index
alloc Pool (PrimState m) a
pool a
x = ST (PrimState m) Index -> m Index
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) Index -> m Index)
-> ST (PrimState m) Index -> m Index
forall a b. (a -> b) -> a -> b
$ Pool (PrimState m) a -> a -> ST (PrimState m) Index
forall a s. (HasCallStack, Unbox a) => Pool s a -> a -> ST s Index
allocST Pool (PrimState m) a
pool a
x

-- | \(O(1)\) Frees an element. Be sure to not free a deleted element.
--
-- ==== Constraints
-- - \(0 \le i \lt n\)
{-# INLINE free #-}
free :: (PrimMonad m) => Pool (PrimState m) a -> Index -> m ()
free :: forall (m :: * -> *) a.
PrimMonad m =>
Pool (PrimState m) a -> Index -> m ()
free Pool {MVector (PrimState m) a
MVector (PrimState m) Index
Buffer (PrimState m) Index
dataPool :: forall s a. Pool s a -> MVector s a
freePool :: forall s a. Pool s a -> Buffer s Index
nextPool :: forall s a. Pool s a -> MVector s Index
dataPool :: MVector (PrimState m) a
freePool :: Buffer (PrimState m) Index
nextPool :: MVector (PrimState m) Index
..} Index
i = do
  Buffer (PrimState m) Index -> Index -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> a -> m ()
B.pushBack Buffer (PrimState m) Index
freePool Index
i

-- | \(O(1)\) Reads the \(k\)-th value.
--
-- ==== Constraints
-- - \(0 \le i \lt n\)
{-# INLINE read #-}
read :: (PrimMonad m, VU.Unbox a) => Pool (PrimState m) a -> Index -> m a
read :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Pool (PrimState m) a -> Index -> m a
read Pool {MVector (PrimState m) a
dataPool :: forall s a. Pool s a -> MVector s a
dataPool :: MVector (PrimState m) a
dataPool} !Index
i = do
  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
dataPool (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
i)

-- | \(O(1)\) Writes to the \(k\)-th value.
--
-- ==== Constraints
-- - \(0 \le i \lt n\)
{-# INLINE write #-}
write :: (PrimMonad m, VU.Unbox a) => Pool (PrimState m) a -> Index -> a -> m ()
write :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Pool (PrimState m) a -> Index -> a -> m ()
write Pool {MVector (PrimState m) a
dataPool :: forall s a. Pool s a -> MVector s a
dataPool :: MVector (PrimState m) a
dataPool} !Index
i !a
x = 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
dataPool (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
i) a
x

-- | \(O(1)\) Modifies the \(k\)-th value.
--
-- ==== Constraints
-- - \(0 \le i \lt n\)
{-# INLINE modify #-}
modify :: (PrimMonad m, VU.Unbox a) => Pool (PrimState m) a -> (a -> a) -> Index -> m ()
modify :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Pool (PrimState m) a -> (a -> a) -> Index -> m ()
modify Pool {MVector (PrimState m) a
dataPool :: forall s a. Pool s a -> MVector s a
dataPool :: MVector (PrimState m) a
dataPool} !a -> a
f !Index
i = do
  MVector (PrimState m) a -> (a -> a) -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector (PrimState m) a
dataPool a -> a
f (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
i)

-- | \(O(1)\) Exchanges the \(k\)-th value.
--
-- ==== Constraints
-- - \(0 \le i \lt n\)
{-# INLINE exchange #-}
exchange :: (PrimMonad m, VU.Unbox a) => Pool (PrimState m) a -> Index -> a -> m a
exchange :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Pool (PrimState m) a -> Index -> a -> m a
exchange Pool {MVector (PrimState m) a
dataPool :: forall s a. Pool s a -> MVector s a
dataPool :: MVector (PrimState m) a
dataPool} !Index
i !a
x = do
  MVector (PrimState m) a -> Int -> a -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m a
VGM.exchange MVector (PrimState m) a
dataPool (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
i) a
x

-- | Mutable `Handle` of an `Index`.
--
-- @since 1.2.0.0
newtype Handle s = Handle
  { -- | @since 1.2.0.0
    forall s. Handle s -> MVector s Index
unHandle :: VUM.MVector s Index
  }

-- | \(O(1)\) Creates a new sequence `Handle` from a root node index.
--
-- @since 1.2.0.0
{-# INLINE newHandle #-}
newHandle :: (PrimMonad m) => Index -> m (Handle (PrimState m))
newHandle :: forall (m :: * -> *).
PrimMonad m =>
Index -> m (Handle (PrimState m))
newHandle Index
x = MVector (PrimState m) Index -> Handle (PrimState m)
forall s. MVector s Index -> Handle s
Handle (MVector (PrimState m) Index -> Handle (PrimState m))
-> m (MVector (PrimState m) Index) -> m (Handle (PrimState m))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Index -> m (MVector (PrimState m) Index)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
1 Index
x

-- | \(O(1)\) Returns whether the sequence is empty.
--
-- @since 1.2.0.0
{-# INLINE nullHandle #-}
nullHandle :: (PrimMonad m) => Handle (PrimState m) -> m Bool
nullHandle :: forall (m :: * -> *). PrimMonad m => Handle (PrimState m) -> m Bool
nullHandle (Handle MVector (PrimState m) Index
h) = Index -> Bool
nullIndex (Index -> Bool) -> m Index -> m Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState m) Index -> Int -> m Index
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector (PrimState m) Index
h Int
0

-- | \(O(1)\) Invalidates a sequence handle. Note that it does not change or `free` the sequence.
--
-- @since 1.2.0.0
{-# INLINE invalidateHandle #-}
invalidateHandle :: (PrimMonad m) => Handle (PrimState m) -> m ()
invalidateHandle :: forall (m :: * -> *). PrimMonad m => Handle (PrimState m) -> m ()
invalidateHandle (Handle MVector (PrimState m) Index
h) = MVector (PrimState m) Index -> Int -> Index -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.unsafeWrite MVector (PrimState m) Index
h Int
0 Index
undefIndex

-- -------------------------------------------------------------------------------------------------
-- Internal
-- -------------------------------------------------------------------------------------------------

{-# INLINEABLE newST #-}
newST :: (VU.Unbox a) => Int -> ST s (Pool s a)
newST :: forall a s. Unbox a => Int -> ST s (Pool s a)
newST Int
cap = do
  MVector s a
dataPool <- Int -> ST s (MVector (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
cap
  Buffer s Index
freePool <- Int -> ST s (Buffer (PrimState (ST s)) Index)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (Buffer (PrimState m) a)
B.new Int
cap
  MVector s Index
nextPool <- Int -> Index -> ST s (MVector (PrimState (ST s)) Index)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
1 (Int -> Index
Index Int
0)
  Pool s a -> ST s (Pool s a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Pool {MVector s a
MVector s Index
Buffer s Index
dataPool :: MVector s a
freePool :: Buffer s Index
nextPool :: MVector s Index
dataPool :: MVector s a
freePool :: Buffer s Index
nextPool :: MVector s Index
..}

{-# INLINEABLE clearST #-}
clearST :: Pool s a -> ST s ()
clearST :: forall s a. Pool s a -> ST s ()
clearST Pool {MVector s a
MVector s Index
Buffer s Index
dataPool :: forall s a. Pool s a -> MVector s a
freePool :: forall s a. Pool s a -> Buffer s Index
nextPool :: forall s a. Pool s a -> MVector s Index
dataPool :: MVector s a
freePool :: Buffer s Index
nextPool :: MVector s Index
..} = do
  Buffer (PrimState (ST s)) Index -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> m ()
B.clear Buffer s Index
Buffer (PrimState (ST s)) Index
freePool
  MVector (PrimState (ST s)) Index -> Int -> Index -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.unsafeWrite MVector s Index
MVector (PrimState (ST s)) Index
nextPool Int
0 (Index -> ST s ()) -> Index -> ST s ()
forall a b. (a -> b) -> a -> b
$ Int -> Index
Index Int
0

{-# INLINEABLE sizeST #-}
sizeST :: (VU.Unbox a) => Pool s a -> ST s Int
sizeST :: forall a s. Unbox a => Pool s a -> ST s Int
sizeST Pool {MVector s a
MVector s Index
Buffer s Index
dataPool :: forall s a. Pool s a -> MVector s a
freePool :: forall s a. Pool s a -> Buffer s Index
nextPool :: forall s a. Pool s a -> MVector s Index
dataPool :: MVector s a
freePool :: Buffer s Index
nextPool :: MVector s Index
..} = do
  !Int
nFree <- Buffer (PrimState (ST s)) Index -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> m Int
B.length Buffer s Index
Buffer (PrimState (ST s)) Index
freePool
  Index !Int
next <- MVector (PrimState (ST s)) Index -> Int -> ST s Index
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector s Index
MVector (PrimState (ST s)) Index
nextPool Int
0
  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
next Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
nFree

{-# INLINEABLE allocST #-}
allocST :: (HasCallStack, VU.Unbox a) => Pool s a -> a -> ST s Index
allocST :: forall a s. (HasCallStack, Unbox a) => Pool s a -> a -> ST s Index
allocST Pool {MVector s a
MVector s Index
Buffer s Index
dataPool :: forall s a. Pool s a -> MVector s a
freePool :: forall s a. Pool s a -> Buffer s Index
nextPool :: forall s a. Pool s a -> MVector s Index
dataPool :: MVector s a
freePool :: Buffer s Index
nextPool :: MVector s Index
..} !a
x = do
  Buffer (PrimState (ST s)) Index -> ST s (Maybe Index)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> m (Maybe a)
B.popBack Buffer s Index
Buffer (PrimState (ST s)) Index
freePool ST s (Maybe Index) -> (Maybe Index -> ST s Index) -> ST s Index
forall a b. ST s a -> (a -> ST s b) -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Just Index
i -> Index -> ST s Index
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Index
i
    Maybe Index
Nothing -> do
      Index Int
i <- MVector (PrimState (ST s)) Index -> Int -> ST s Index
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector s Index
MVector (PrimState (ST s)) Index
nextPool Int
0
      if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< MVector s a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
VGM.length MVector s a
dataPool
        then do
          MVector (PrimState (ST s)) Index -> Int -> Index -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.unsafeWrite MVector s Index
MVector (PrimState (ST s)) Index
nextPool Int
0 (Index -> ST s ()) -> Index -> ST s ()
forall a b. (a -> b) -> a -> b
$ Int -> Index
forall a b. Coercible a b => a -> b
coerce (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
          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
dataPool Int
i a
x
          Index -> ST s Index
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Index -> ST s Index) -> Index -> ST s Index
forall a b. (a -> b) -> a -> b
$ Int -> Index
forall a b. Coercible a b => a -> b
coerce Int
i
        else do
          String -> ST s Index
forall a. HasCallStack => String -> a
error String
"AtCoder.Extra.Pool.allocST: capacity out of bounds"