module PureSAT.Vec (
Vec,
newVec,
sizeofVec,
insertVec,
readVec,
writeVec,
shrinkVec,
) where
import Data.Primitive.PrimVar
import Unsafe.Coerce (unsafeCoerce)
import PureSAT.Base
import PureSAT.Utils
import PureSAT.Prim
data Vec s a = Vec {-# UNPACK #-} !(PrimVar s Int) {-# UNPACK #-} !(MutableArray s a)
newVec
:: Int
-> ST s (Vec s a)
newVec :: forall s a. Int -> ST s (Vec s a)
newVec Int
capacity = do
arr <- Int -> a -> ST s (MutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
newArray (Int -> Int
nextPowerOf2 (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
64 Int
capacity)) a
forall a. a
unused
size <- newPrimVar 0
return (Vec size arr)
unused :: a
unused :: forall a. a
unused = a
forall a. HasCallStack => a
undefined
sizeofVec :: Vec s a -> ST s Int
sizeofVec :: forall s a. Vec s a -> ST s Int
sizeofVec (Vec PrimVar s Int
size MutableArray s a
_) = PrimVar (PrimState (ST s)) Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> m a
readPrimVar PrimVar s Int
PrimVar (PrimState (ST s)) Int
size
{-# INLINE sizeofVec #-}
insertVec :: Vec s a -> a -> ST s (Vec s a)
insertVec :: forall s a. Vec s a -> a -> ST s (Vec s a)
insertVec vec :: Vec s a
vec@(Vec PrimVar s Int
sizeRef MutableArray s a
arr) a
x = do
size <- PrimVar (PrimState (ST s)) Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> m a
readPrimVar PrimVar s Int
PrimVar (PrimState (ST s)) Int
sizeRef
let !capacity = MutableArray s a -> Int
forall s a. MutableArray s a -> Int
sizeofMutableArray MutableArray s a
arr
if size < capacity
then do
writeArray arr size x
writePrimVar sizeRef (size + 1)
return vec
else do
new <- newArray (capacity * 2) unused
copyMutableArray new 0 arr 0 size
writeArray new size x
writePrimVar sizeRef (size + 1)
return (Vec sizeRef new)
readVec :: Vec s a -> Int -> ST s a
readVec :: forall s a. Vec s a -> Int -> ST s a
readVec (Vec PrimVar s Int
_ MutableArray s a
arr) Int
i = MutableArray s a -> Int -> ST s a
forall s a. HasCallStack => MutableArray s a -> Int -> ST s a
readArray MutableArray s a
arr Int
i
writeVec :: Vec s a -> Int -> a -> ST s ()
writeVec :: forall s a. Vec s a -> Int -> a -> ST s ()
writeVec (Vec PrimVar s Int
_ MutableArray s a
arr) Int
i a
x = MutableArray s a -> Int -> a -> ST s ()
forall s a. HasCallStack => MutableArray s a -> Int -> a -> ST s ()
writeArray MutableArray s a
arr Int
i a
x
shrinkVec :: Vec s a -> Int -> ST s ()
shrinkVec :: forall s a. Vec s a -> Int -> ST s ()
shrinkVec (Vec PrimVar s Int
sizeRef MutableArray s a
arr) Int
newSize = do
size <- PrimVar (PrimState (ST s)) Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> m a
readPrimVar PrimVar s Int
PrimVar (PrimState (ST s)) Int
sizeRef
forM_ [newSize .. size - 1] $ \Int
i -> MutableArray s a -> Int -> a -> ST s ()
forall s a. HasCallStack => MutableArray s a -> Int -> a -> ST s ()
writeArray MutableArray s a
arr Int
i a
forall a. a
unused
writePrimVar sizeRef newSize