{-# LANGUAGE RecordWildCards #-}

-- | A pushable vector with fixed capacity \(n\). Internally, it tracks the number of elements.
--
-- ==== __Example__
-- Create a buffer with capacity @4@:
--
-- >>> import AtCoder.Internal.Buffer qualified as B
-- >>> buf <- B.new @_ @Int 4
-- >>> B.capacity buf
-- 4
--
-- >>> B.null buf        -- [_   _  _  _]
-- True
--
-- Append elements with `pushBack`:
--
-- >>> B.pushBack buf 10 -- [10  _  _ _]
-- >>> B.pushBack buf 11 -- [10, 11  _  _]
-- >>> length buf
-- 2
--
-- Access each elements with `read`, `write`, `modify` or `modifyM`:
--
-- >>> B.read buf 0
-- 10
--
-- >>> B.write buf 1 0   -- [10, 0,  _  _]
--
-- Remove elements with `pushBack`:
--
-- >>> B.popBack buf     -- [10  _  _  _]
-- Just 0
--
-- Inspect the internal vector with `freeze`:
--
-- >>> B.freeze buf
-- [10]
--
-- >>> B.clear buf       -- []
-- >>> B.null buf
-- True
--
-- >>> B.unsafeFreeze buf
-- []
--
-- @since 1.0.0.0
module AtCoder.Internal.Buffer
  ( -- * Buffer
    Buffer,

    -- * Constructors
    new,
    build,

    -- * Metadata
    capacity,
    length,
    null,

    -- * Reading
    back,
    read,

    -- * Modifications
    pushBack,
    popBack,
    write,
    modify,
    modifyM,
    clear,

    -- * Conversions
    freeze,
    unsafeFreeze,
  )
where

import AtCoder.Internal.Assert qualified as ACIA
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, read)

-- | A pushable vector with fixed capacity \(n\). Internally, it tracks the number of elements.
--
-- @since 1.0.0.0
data Buffer s a = Buffer
  { forall s a. Buffer s a -> MVector s Int
lenB :: !(VUM.MVector s Int),
    forall s a. Buffer s a -> MVector s a
vecB :: !(VUM.MVector s a)
  }

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

-- | \(O(n)\) Creates a buffer with capacity \(n\) with initial values.
--
-- @since 1.0.0.0
{-# INLINE build #-}
build :: (PrimMonad m, VU.Unbox a) => VU.Vector a -> m (Buffer (PrimState m) a)
build :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Vector a -> m (Buffer (PrimState m) a)
build Vector a
xs = do
  MVector (PrimState m) Int
lenB <- Int -> Int -> m (MVector (PrimState m) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
1 (Int -> m (MVector (PrimState m) Int))
-> Int -> m (MVector (PrimState m) Int)
forall a b. (a -> b) -> a -> b
$ Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
xs
  MVector (PrimState m) a
vecB <- Vector a -> m (MVector (PrimState m) a)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
VU.thaw Vector a
xs
  Buffer (PrimState m) a -> m (Buffer (PrimState m) a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Buffer {MVector (PrimState m) a
MVector (PrimState m) Int
lenB :: MVector (PrimState m) Int
vecB :: MVector (PrimState m) a
lenB :: MVector (PrimState m) Int
vecB :: MVector (PrimState m) a
..}

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

-- | \(O(1)\) Returns the number of elements in the buffer.
--
-- @since 1.0.0.0
{-# INLINE length #-}
length :: (PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> m Int
length :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> m Int
length Buffer {MVector (PrimState m) a
MVector (PrimState m) Int
lenB :: forall s a. Buffer s a -> MVector s Int
vecB :: forall s a. Buffer s a -> MVector s a
lenB :: MVector (PrimState m) Int
vecB :: MVector (PrimState m) a
..} = do
  MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) Int
lenB Int
0

-- | \(O(1)\) Returns `True` if the buffer is empty.
--
-- @since 1.0.0.0
{-# INLINE null #-}
null :: (PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> m Bool
null :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Buffer (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)
-> (Buffer (PrimState m) a -> m Int)
-> Buffer (PrimState m) a
-> m Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Buffer (PrimState m) a -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> m Int
length

-- | \(O(1)\) Returns the last value in the buffer, or `Nothing` if it is empty.
--
-- @since 1.0.0.0
{-# INLINE back #-}
back :: (PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> m (Maybe a)
back :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> m (Maybe a)
back Buffer {MVector (PrimState m) a
MVector (PrimState m) Int
lenB :: forall s a. Buffer s a -> MVector s Int
vecB :: forall s a. Buffer s a -> MVector s a
lenB :: MVector (PrimState m) Int
vecB :: MVector (PrimState m) a
..} = do
  Int
len <- MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) Int
lenB Int
0
  if Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
    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
vecB (Int
len 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)\) Yields the element at the given position. Will throw an exception if the index is out
-- of range.
--
-- @since 1.0.0.0
{-# INLINE read #-}
read :: (HasCallStack, PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> Int -> m a
read :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> Int -> m a
read Buffer {MVector (PrimState m) a
MVector (PrimState m) Int
lenB :: forall s a. Buffer s a -> MVector s Int
vecB :: forall s a. Buffer s a -> MVector s a
lenB :: MVector (PrimState m) Int
vecB :: MVector (PrimState m) a
..} Int
i = do
  Int
len <- MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) Int
lenB Int
0
  let !()
_ = HasCallStack => String -> Int -> Int -> ()
String -> Int -> Int -> ()
ACIA.checkIndex String
"AtCoder.Internal.Buffer.read" Int
i Int
len
  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
vecB Int
i

-- | \(O(1)\) Appends an element to the back.
--
-- @since 1.0.0.0
{-# INLINE pushBack #-}
pushBack :: (HasCallStack, PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> a -> m ()
pushBack :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> a -> m ()
pushBack Buffer {MVector (PrimState m) a
MVector (PrimState m) Int
lenB :: forall s a. Buffer s a -> MVector s Int
vecB :: forall s a. Buffer s a -> MVector s a
lenB :: MVector (PrimState m) Int
vecB :: MVector (PrimState m) a
..} a
e = do
  Int
len <- MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) Int
lenB Int
0
  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
vecB Int
len a
e
  MVector (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState m) Int
lenB Int
0 (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

-- | \(O(1)\) Removes the last element from the buffer and returns it, or `Nothing` if it is empty.
--
-- @since 1.0.0.0
{-# INLINE popBack #-}
popBack :: (PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> m (Maybe a)
popBack :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> m (Maybe a)
popBack Buffer {MVector (PrimState m) a
MVector (PrimState m) Int
lenB :: forall s a. Buffer s a -> MVector s Int
vecB :: forall s a. Buffer s a -> MVector s a
lenB :: MVector (PrimState m) Int
vecB :: MVector (PrimState m) a
..} = do
  Int
len <- MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) Int
lenB Int
0
  if Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
    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
vecB (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
      MVector (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState m) Int
lenB Int
0 (Int
len 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)\) Writes to the element at the given position. Will throw an exception if the index is
-- out of bounds.
--
-- @since 1.0.0.0
{-# INLINE write #-}
write :: (HasCallStack, PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> Int -> a -> m ()
write :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> Int -> a -> m ()
write Buffer {MVector (PrimState m) a
MVector (PrimState m) Int
lenB :: forall s a. Buffer s a -> MVector s Int
vecB :: forall s a. Buffer s a -> MVector s a
lenB :: MVector (PrimState m) Int
vecB :: MVector (PrimState m) a
..} Int
i a
e = do
  Int
len <- MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) Int
lenB Int
0
  let !()
_ = HasCallStack => String -> Int -> Int -> ()
String -> Int -> Int -> ()
ACIA.checkIndex String
"AtCoder.Internal.Buffer.write" Int
i Int
len
  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
vecB Int
i a
e

-- | \(O(1)\) Writes to the element at the given position. Will throw an exception if the index is
-- out of bounds.
--
-- @since 1.0.0.0
{-# INLINE modify #-}
modify :: (HasCallStack, PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> (a -> a) -> Int -> m ()
modify :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> (a -> a) -> Int -> m ()
modify Buffer {MVector (PrimState m) a
MVector (PrimState m) Int
lenB :: forall s a. Buffer s a -> MVector s Int
vecB :: forall s a. Buffer s a -> MVector s a
lenB :: MVector (PrimState m) Int
vecB :: MVector (PrimState m) a
..} a -> a
f Int
i = do
  Int
len <- MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) Int
lenB Int
0
  let !()
_ = HasCallStack => String -> Int -> Int -> ()
String -> Int -> Int -> ()
ACIA.checkIndex String
"AtCoder.Internal.Buffer.modify" Int
i Int
len
  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
vecB a -> a
f Int
i

-- | \(O(1)\) Writes to the element at the given position. Will throw an exception if the index is
-- out of bounds.
--
-- @since 1.0.0.0
{-# INLINE modifyM #-}
modifyM :: (HasCallStack, PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> (a -> m a) -> Int -> m ()
modifyM :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> (a -> m a) -> Int -> m ()
modifyM Buffer {MVector (PrimState m) a
MVector (PrimState m) Int
lenB :: forall s a. Buffer s a -> MVector s Int
vecB :: forall s a. Buffer s a -> MVector s a
lenB :: MVector (PrimState m) Int
vecB :: MVector (PrimState m) a
..} a -> m a
f Int
i = do
  Int
len <- MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) Int
lenB Int
0
  let !()
_ = HasCallStack => String -> Int -> Int -> ()
String -> Int -> Int -> ()
ACIA.checkIndex String
"AtCoder.Internal.Buffer.modifyM" Int
i Int
len
  MVector (PrimState m) a -> (a -> m a) -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> m a) -> Int -> m ()
VGM.modifyM MVector (PrimState m) a
vecB a -> m a
f Int
i

-- | \(O(1)\) Sets the `length` to zero.
--
-- @since 1.0.0.0
{-# INLINE clear #-}
clear :: (PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> m ()
clear :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> m ()
clear Buffer {MVector (PrimState m) a
MVector (PrimState m) Int
lenB :: forall s a. Buffer s a -> MVector s Int
vecB :: forall s a. Buffer s a -> MVector s a
lenB :: MVector (PrimState m) Int
vecB :: MVector (PrimState m) a
..} = do
  MVector (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState m) Int
lenB Int
0 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) => Buffer (PrimState m) a -> m (VU.Vector a)
freeze :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> m (Vector a)
freeze Buffer {MVector (PrimState m) a
MVector (PrimState m) Int
lenB :: forall s a. Buffer s a -> MVector s Int
vecB :: forall s a. Buffer s a -> MVector s a
lenB :: MVector (PrimState m) Int
vecB :: MVector (PrimState m) a
..} = do
  Int
len <- MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) Int
lenB Int
0
  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
len MVector (PrimState m) a
vecB

-- | \(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) => Buffer (PrimState m) a -> m (VU.Vector a)
unsafeFreeze :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Buffer (PrimState m) a -> m (Vector a)
unsafeFreeze Buffer {MVector (PrimState m) a
MVector (PrimState m) Int
lenB :: forall s a. Buffer s a -> MVector s Int
vecB :: forall s a. Buffer s a -> MVector s a
lenB :: MVector (PrimState m) Int
vecB :: MVector (PrimState m) a
..} = do
  Int
len <- MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) Int
lenB Int
0
  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
len MVector (PrimState m) a
vecB