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
  { forall s a. Buffer s a -> Int
offset :: !Int,
    forall s a. Buffer s a -> SmallMutableArray s a
marr :: !(SmallMutableArray s a)
  }

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

read :: (PrimMonad m, s ~ PrimState m) => Int -> Buffer s a -> m a
read :: forall (m :: * -> *) s a.
(PrimMonad m, s ~ PrimState m) =>
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} = forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> m a
readSmallArray SmallMutableArray s a
marr Int
i
{-# INLINE read #-}

write :: (PrimMonad m, s ~ PrimState m) => Int -> a -> Buffer s a -> m ()
write :: forall (m :: * -> *) s a.
(PrimMonad m, s ~ PrimState m) =>
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} = forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s a
marr Int
i a
a
{-# INLINE write #-}

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

shrink :: Int -> Buffer s a -> Buffer s a
shrink :: forall s a. 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 :: forall (m :: * -> *) s a.
(PrimMonad m, s ~ PrimState m) =>
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 <- forall (m :: * -> *) a.
PrimMonad m =>
MArray (PrimState m) a -> Int -> m (MArray (PrimState m) a)
shrinkSmallMutableArray_ SmallMutableArray s a
marr Int
i
  forall (f :: * -> *) a. Applicative f => a -> f a
pure 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 :: forall s a. 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} = forall s a. SmallMutableArray s a -> Int
sizeofSmallMutableArray SmallMutableArray s a
marr
{-# INLINE capacity #-}

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

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

undefinedElem :: forall a. a
undefinedElem :: forall a. a
undefinedElem = 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 :: forall (m :: * -> *) s a.
(PrimMonad m, s ~ PrimState m) =>
Buffer s a -> m (Buffer s a)
resize Buffer s a
buffer = do
  if forall s a. Buffer s a -> Int
capacity Buffer s a
buffer forall a. Eq a => a -> a -> Bool
== Int
0
    then 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 forall (m :: * -> *) s a.
(PrimMonad m, s ~ PrimState m) =>
Int -> Buffer s a -> m (Buffer s a)
grow (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 :: forall (m :: * -> *) s a.
(PrimMonad m, s ~ PrimState m) =>
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' <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray (forall s a. SmallMutableArray s a -> Int
sizeofSmallMutableArray SmallMutableArray s a
marr forall a. Num a => a -> a -> a
+ Int
more) forall a. a
undefinedElem
  forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallMutableArray (PrimState m) a -> Int -> Int -> m ()
copySmallMutableArray SmallMutableArray s a
marr' Int
0 SmallMutableArray s a
marr Int
0 Int
offset
  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 :: forall (m :: * -> *) s a.
(PrimMonad m, s ~ PrimState m) =>
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} = forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> Int -> m (SmallArray a)
freezeSmallArray SmallMutableArray s a
marr Int
0 Int
offset
{-# INLINE freeze #-}

unsafeFreeze :: (PrimMonad m, s ~ PrimState m) => Buffer s a -> m (SmallArray a)
unsafeFreeze :: forall (m :: * -> *) s a.
(PrimMonad m, s ~ PrimState m) =>
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 <- forall (m :: * -> *) a.
PrimMonad m =>
MArray (PrimState m) a -> Int -> m (MArray (PrimState m) a)
shrinkSmallMutableArray_ SmallMutableArray s a
marr Int
offset
  forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
unsafeFreezeSmallArray SmallMutableArray s a
marr
{-# INLINE unsafeFreeze #-}