module Data.Vector.Persistent.Internal.Buffer where

import Control.Monad.Primitive
import Data.Primitive.SmallArray
import Data.Vector.Persistent.Internal.Array (shrinkSmallMutableArray_)
import Prelude hiding (length)

data Buffer s a = Buffer
  { Buffer s a -> Int
offset :: !Int,
    Buffer s a -> SmallMutableArray s a
marr :: !(SmallMutableArray s a)
  }

new :: (PrimMonad m, s ~ PrimState m) => m (Buffer s a)
new :: m (Buffer s a)
new = do
  SmallMutableArray s a
marr <- Int -> a -> m (SmallMutableArray (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
0 a
forall a. a
undefinedElem
  Buffer s a -> m (Buffer s a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Buffer :: forall s a. Int -> SmallMutableArray s a -> Buffer s a
Buffer {$sel:offset:Buffer :: Int
offset = Int
0, SmallMutableArray s a
marr :: SmallMutableArray s a
$sel:marr:Buffer :: SmallMutableArray s a
marr}
{-# INLINE new #-}

newWithCapacity :: (PrimMonad m, s ~ PrimState m) => Int -> m (Buffer s a)
newWithCapacity :: Int -> m (Buffer s a)
newWithCapacity Int
cap = do
  SmallMutableArray s a
marr <- Int -> a -> m (SmallMutableArray (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
cap a
forall a. a
undefinedElem
  Buffer s a -> m (Buffer s a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Buffer :: forall s a. Int -> SmallMutableArray s a -> Buffer s a
Buffer {$sel:offset:Buffer :: Int
offset = Int
0, SmallMutableArray s a
marr :: SmallMutableArray s a
$sel:marr:Buffer :: SmallMutableArray s a
marr}
{-# INLINE newWithCapacity #-}

push :: (PrimMonad m, s ~ PrimState m) => a -> Buffer s a -> m (Buffer s a)
push :: a -> Buffer s a -> m (Buffer s a)
push a
a Buffer s a
buffer = do
  Buffer s a
buffer' <-
    if Buffer s a -> Int
forall s a. Buffer s a -> Int
length Buffer s a
buffer Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Buffer s a -> Int
forall s a. Buffer s a -> Int
capacity Buffer s a
buffer
      then Buffer s a -> m (Buffer s a)
forall (m :: * -> *) s a.
(PrimMonad m, s ~ PrimState m) =>
Buffer s a -> m (Buffer s a)
resize Buffer s a
buffer
      else Buffer s a -> m (Buffer s a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Buffer s a
buffer
  SmallMutableArray (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray (Buffer s a -> SmallMutableArray s a
forall s a. Buffer s a -> SmallMutableArray s a
marr Buffer s a
buffer') (Buffer s a -> Int
forall s a. Buffer s a -> Int
length Buffer s a
buffer) a
a
  Buffer s a -> m (Buffer s a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Buffer s a
buffer' {$sel:offset:Buffer :: Int
offset = Buffer s a -> Int
forall s a. Buffer s a -> Int
offset Buffer s a
buffer' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1}
{-# INLINE push #-}

read :: (PrimMonad m, s ~ PrimState m) => Int -> Buffer s a -> m a
read :: Int -> Buffer s a -> m a
read Int
i Buffer {SmallMutableArray s a
marr :: SmallMutableArray s a
$sel:marr:Buffer :: forall s a. Buffer s a -> SmallMutableArray s a
marr} = SmallMutableArray (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> m a
readSmallArray SmallMutableArray s a
SmallMutableArray (PrimState m) a
marr Int
i
{-# INLINE read #-}

write :: (PrimMonad m, s ~ PrimState m) => Int -> a -> Buffer s a -> m ()
write :: Int -> a -> Buffer s a -> m ()
write Int
i a
a Buffer {SmallMutableArray s a
marr :: SmallMutableArray s a
$sel:marr:Buffer :: forall s a. Buffer s a -> SmallMutableArray s a
marr} = SmallMutableArray (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s a
SmallMutableArray (PrimState m) a
marr Int
i a
a
{-# INLINE write #-}

clear :: Buffer s a -> Buffer s a
clear :: Buffer s a -> Buffer s a
clear = Int -> Buffer s a -> Buffer s a
forall s a. Int -> Buffer s a -> Buffer s a
shrink Int
0
{-# INLINE clear #-}

shrink :: Int -> Buffer s a -> Buffer s a
shrink :: Int -> Buffer s a -> Buffer s a
shrink Int
i Buffer s a
buffer = Buffer s a
buffer {$sel:offset:Buffer :: Int
offset = Int
i}
{-# INLINE shrink #-}

unsafeShrink :: (PrimMonad m, s ~ PrimState m) => Int -> Buffer s a -> m (Buffer s a)
unsafeShrink :: Int -> Buffer s a -> m (Buffer s a)
unsafeShrink Int
i Buffer {SmallMutableArray s a
marr :: SmallMutableArray s a
$sel:marr:Buffer :: forall s a. Buffer s a -> SmallMutableArray s a
marr} = do
  SmallMutableArray s a
marr <- MArray (PrimState m) a -> Int -> m (MArray (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
MArray (PrimState m) a -> Int -> m (MArray (PrimState m) a)
shrinkSmallMutableArray_ SmallMutableArray s a
MArray (PrimState m) a
marr Int
i
  Buffer s a -> m (Buffer s a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Buffer :: forall s a. Int -> SmallMutableArray s a -> Buffer s a
Buffer {SmallMutableArray s a
marr :: SmallMutableArray s a
$sel:marr:Buffer :: SmallMutableArray s a
marr, $sel:offset:Buffer :: Int
offset = Int
i}
{-# INLINE unsafeShrink #-}

capacity :: Buffer s a -> Int
capacity :: Buffer s a -> Int
capacity Buffer {SmallMutableArray s a
marr :: SmallMutableArray s a
$sel:marr:Buffer :: forall s a. Buffer s a -> SmallMutableArray s a
marr} = SmallMutableArray s a -> Int
forall s a. SmallMutableArray s a -> Int
sizeofSmallMutableArray SmallMutableArray s a
marr
{-# INLINE capacity #-}

null :: Buffer s a -> Bool
null :: Buffer s a -> Bool
null = (Int
0 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
==) (Int -> Bool) -> (Buffer s a -> Int) -> Buffer s a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Buffer s a -> Int
forall s a. Buffer s a -> Int
length

length :: Buffer s a -> Int
length :: Buffer s a -> Int
length = Buffer s a -> Int
forall s a. Buffer s a -> Int
offset
{-# INLINE length #-}

undefinedElem :: forall a. a
undefinedElem :: a
undefinedElem = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"undefined element"
{-# NOINLINE undefinedElem #-}

resize :: (PrimMonad m, s ~ PrimState m) => Buffer s a -> m (Buffer s a)
resize :: Buffer s a -> m (Buffer s a)
resize Buffer s a
buffer = do
  if Buffer s a -> Int
forall s a. Buffer s a -> Int
capacity Buffer s a
buffer Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
    then Int -> Buffer s a -> m (Buffer s a)
forall (m :: * -> *) s a.
(PrimMonad m, s ~ PrimState m) =>
Int -> Buffer s a -> m (Buffer s a)
grow Int
32 Buffer s a
buffer
    else Int -> Buffer s a -> m (Buffer s a)
forall (m :: * -> *) s a.
(PrimMonad m, s ~ PrimState m) =>
Int -> Buffer s a -> m (Buffer s a)
grow (Buffer s a -> Int
forall s a. Buffer s a -> Int
capacity Buffer s a
buffer) Buffer s a
buffer
{-# INLINE resize #-}

grow :: (PrimMonad m, s ~ PrimState m) => Int -> Buffer s a -> m (Buffer s a)
grow :: Int -> Buffer s a -> m (Buffer s a)
grow Int
more buffer :: Buffer s a
buffer@Buffer {SmallMutableArray s a
marr :: SmallMutableArray s a
$sel:marr:Buffer :: forall s a. Buffer s a -> SmallMutableArray s a
marr, Int
offset :: Int
$sel:offset:Buffer :: forall s a. Buffer s a -> Int
offset} = do
  SmallMutableArray s a
marr' <- Int -> a -> m (SmallMutableArray (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray (SmallMutableArray s a -> Int
forall s a. SmallMutableArray s a -> Int
sizeofSmallMutableArray SmallMutableArray s a
marr Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
more) a
forall a. a
undefinedElem
  SmallMutableArray (PrimState m) a
-> Int -> SmallMutableArray (PrimState m) a -> Int -> Int -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallMutableArray (PrimState m) a -> Int -> Int -> m ()
copySmallMutableArray SmallMutableArray s a
SmallMutableArray (PrimState m) a
marr' Int
0 SmallMutableArray s a
SmallMutableArray (PrimState m) a
marr Int
0 Int
offset
  Buffer s a -> m (Buffer s a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Buffer s a
buffer {$sel:marr:Buffer :: SmallMutableArray s a
marr = SmallMutableArray s a
marr'}
{-# INLINE grow #-}

freeze :: (PrimMonad m, s ~ PrimState m) => Buffer s a -> m (SmallArray a)
freeze :: Buffer s a -> m (SmallArray a)
freeze Buffer {SmallMutableArray s a
marr :: SmallMutableArray s a
$sel:marr:Buffer :: forall s a. Buffer s a -> SmallMutableArray s a
marr, Int
offset :: Int
$sel:offset:Buffer :: forall s a. Buffer s a -> Int
offset} = SmallMutableArray (PrimState m) a -> Int -> Int -> m (SmallArray a)
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> Int -> m (SmallArray a)
freezeSmallArray SmallMutableArray s a
SmallMutableArray (PrimState m) a
marr Int
0 Int
offset
{-# INLINE freeze #-}

unsafeFreeze :: (PrimMonad m, s ~ PrimState m) => Buffer s a -> m (SmallArray a)
unsafeFreeze :: Buffer s a -> m (SmallArray a)
unsafeFreeze Buffer {SmallMutableArray s a
marr :: SmallMutableArray s a
$sel:marr:Buffer :: forall s a. Buffer s a -> SmallMutableArray s a
marr, Int
offset :: Int
$sel:offset:Buffer :: forall s a. Buffer s a -> Int
offset} = do
  SmallMutableArray s a
marr <- MArray (PrimState m) a -> Int -> m (MArray (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
MArray (PrimState m) a -> Int -> m (MArray (PrimState m) a)
shrinkSmallMutableArray_ SmallMutableArray s a
MArray (PrimState m) a
marr Int
offset
  MArray (PrimState m) a -> m (SmallArray a)
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
unsafeFreezeSmallArray SmallMutableArray s a
MArray (PrimState m) a
marr
{-# INLINE unsafeFreeze #-}