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

data Buffer s a = Buffer !(A.MutableArray s a) !(IntRef s)

new :: Int -> ST s (Buffer s a)
new :: Int -> ST s (Buffer s a)
new Int
capacity = do
    MutableArray s a
buffer <- Int -> ST s (MutableArray s a)
forall s a. Int -> ST s (MutableArray s a)
A.new Int
capacity
    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 (f :: * -> *) a. Applicative f => a -> f a
pure (MutableArray s a -> IntRef s -> Buffer s a
forall s a. MutableArray s a -> IntRef s -> Buffer s a
Buffer MutableArray s a
buffer IntRef s
offset)

push :: Buffer s a -> a -> ST s ()
push :: Buffer s a -> a -> ST s ()
push (Buffer MutableArray 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
    MutableArray s a -> Int -> a -> ST s ()
forall s a. MutableArray s a -> Int -> a -> ST s ()
A.write MutableArray s a
buffer Int
idx a
x
    IntRef s -> (Int -> Int) -> ST s ()
forall s. IntRef s -> (Int -> Int) -> ST s ()
modifyIntRef IntRef s
offset (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

get :: Buffer s a -> ST s (A.Array a)
get :: Buffer s a -> ST s (Array a)
get (Buffer MutableArray 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
    Array a
result <- MutableArray s a -> Int -> Int -> ST s (Array a)
forall s a. MutableArray s a -> Int -> Int -> ST s (Array a)
A.freeze MutableArray 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 (f :: * -> *) a. Applicative f => a -> f a
pure Array a
result

size :: Buffer s a -> ST s Int
size :: Buffer s a -> ST s Int
size (Buffer MutableArray s a
_ IntRef s
offset) = IntRef s -> ST s Int
forall s. IntRef s -> ST s Int
readIntRef IntRef s
offset
{-# INLInE size #-}