-- | 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.Primitive.SmallArray
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 !(SmallMutableArray 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
    SmallMutableArray s a
buffer <- Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
capacity a
forall a. a
uninitialized
    IntRef s
offset <- Int -> ST s (IntRef s)
forall s. Int -> ST s (IntRef s)
newIntRef Int
0
    Buffer s a -> ST s (Buffer s a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SmallMutableArray s a -> IntRef s -> Buffer s a
forall s a. SmallMutableArray s a -> IntRef s -> Buffer s a
Buffer SmallMutableArray s a
buffer IntRef s
offset)

uninitialized :: a
uninitialized :: forall a. a
uninitialized = [Char] -> a
forall a. [Char] -> a
errorWithoutStackTrace [Char]
"uninitialized"

-- | \(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 SmallMutableArray s a
buffer IntRef s
offset) a
x = do
    Int
idx <- IntRef s -> ST s Int
forall s. IntRef s -> ST s Int
readIntRef IntRef s
offset
    SmallMutableArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
buffer Int
idx a
x
    IntRef s -> Int -> ST s ()
forall s. IntRef s -> Int -> ST s ()
writeIntRef IntRef s
offset (Int
idx Int -> Int -> Int
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 SmallMutableArray s a
buffer IntRef s
offset) = do
    Int
len <- IntRef s -> ST s Int
forall s. IntRef s -> ST s Int
readIntRef IntRef s
offset
    SmallArray a
result <- SmallMutableArray (PrimState (ST s)) a
-> Int -> Int -> ST s (SmallArray a)
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> Int -> m (SmallArray a)
freezeSmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
buffer Int
0 Int
len
    IntRef s -> Int -> ST s ()
forall s. IntRef s -> Int -> ST s ()
writeIntRef IntRef s
offset Int
0
    Array a -> ST s (Array a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SmallArray a -> Array a
forall a. SmallArray a -> Array a
A.wrap SmallArray 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 SmallMutableArray s a
_ IntRef s
offset) = IntRef s -> ST s Int
forall s. IntRef s -> ST s Int
readIntRef IntRef s
offset
{-# INLInE size #-}