-- | This is an internal module.
--
-- A 'Buffer' is an array with a fixed capacity, used to build up 'Data.RRBVector.Internal.Array.Array's.
-- It is used in the implementation of 'Data.RRBVector.fromList' and 'Data.RRBVector.><'.

module Data.RRBVector.Internal.Buffer
    ( Buffer
    , new
    , push
    , get
    , size
    ) where

import Control.Monad.ST

import Data.RRBVector.Internal.IntRef
import qualified Data.RRBVector.Internal.Array as A

-- | A mutable array buffer with a fixed capacity.
data Buffer s a = Buffer !(A.MutableArray s a) !(IntRef s)

-- | \(O(n)\). Create a new empty buffer with the given capacity.
new :: Int -> ST s (Buffer s a)
new :: forall s a. Int -> ST s (Buffer s a)
new Int
capacity = do
    MutableArray s a
buffer <- forall s a. Int -> ST s (MutableArray s a)
A.new Int
capacity
    IntRef s
offset <- forall s. Int -> ST s (IntRef s)
newIntRef Int
0
    forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall s a. MutableArray s a -> IntRef s -> Buffer s a
Buffer MutableArray s a
buffer IntRef s
offset)

-- | \(O(1)\). Push a new element onto the buffer.
-- The size of the buffer must not exceed the capacity, but this is not checked.
push :: Buffer s a -> a -> ST s ()
push :: forall s a. Buffer s a -> a -> ST s ()
push (Buffer MutableArray s a
buffer IntRef s
offset) a
x = do
    Int
idx <- forall s. IntRef s -> ST s Int
readIntRef IntRef s
offset
    forall s a. MutableArray s a -> Int -> a -> ST s ()
A.write MutableArray s a
buffer Int
idx a
x
    forall s. IntRef s -> Int -> ST s ()
writeIntRef IntRef s
offset (Int
idx forall a. Num a => a -> a -> a
+ Int
1)

-- | \(O(n)\). Freeze the content of the buffer and return it.
-- This resets the buffer so that it is empty.
get :: Buffer s a -> ST s (A.Array a)
get :: forall s a. Buffer s a -> ST s (Array a)
get (Buffer MutableArray s a
buffer IntRef s
offset) = do
    Int
len <- forall s. IntRef s -> ST s Int
readIntRef IntRef s
offset
    Array a
result <- forall s a. MutableArray s a -> Int -> Int -> ST s (Array a)
A.freeze MutableArray s a
buffer Int
0 Int
len
    forall s. IntRef s -> Int -> ST s ()
writeIntRef IntRef s
offset Int
0
    forall (f :: * -> *) a. Applicative f => a -> f a
pure Array a
result

-- | \(O(1)\). Return the current size of the buffer.
size :: Buffer s a -> ST s Int
size :: forall s a. Buffer s a -> ST s Int
size (Buffer MutableArray s a
_ IntRef s
offset) = forall s. IntRef s -> ST s Int
readIntRef IntRef s
offset
{-# INLInE size #-}