module Mischief.ECS.Vec
( Vec (..),
IOVec,
length,
null,
capacity,
new,
newSized,
slice,
thaw,
freeze,
toList,
ensure,
ensureAppend,
read,
write,
unsafeRead,
unsafeWrite,
pushBack,
unsafePushBack,
modify,
modify_,
modifyM,
modifyM_,
swap,
tap,
shrink,
takeSwap,
removeSwap,
clone,
)
where
import Control.Monad
import Control.Monad.Primitive
import Data.Foldable (for_)
import Data.Primitive.MutVar
import Data.Vector (Vector)
import Data.Vector qualified as Vector
import Data.Vector.Mutable (MVector)
import Data.Vector.Mutable qualified as MVector
import GHC.Generics
import GHC.Stack (HasCallStack)
import Prelude hiding (length, null, read)
data Vec s a = Vec
{ forall s a. Vec s a -> MutVar s (MVector s a)
buffer :: !(MutVar s (MVector s a)),
forall s a. Vec s a -> MutVar s Int
len :: !(MutVar s Int)
}
deriving ((forall x. Vec s a -> Rep (Vec s a) x)
-> (forall x. Rep (Vec s a) x -> Vec s a) -> Generic (Vec s a)
forall x. Rep (Vec s a) x -> Vec s a
forall x. Vec s a -> Rep (Vec s a) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall s a x. Rep (Vec s a) x -> Vec s a
forall s a x. Vec s a -> Rep (Vec s a) x
$cfrom :: forall s a x. Vec s a -> Rep (Vec s a) x
from :: forall x. Vec s a -> Rep (Vec s a) x
$cto :: forall s a x. Rep (Vec s a) x -> Vec s a
to :: forall x. Rep (Vec s a) x -> Vec s a
Generic)
type IOVec a = Vec RealWorld a
capacity :: (PrimMonad m) => Vec (PrimState m) a -> m Int
capacity :: forall (m :: * -> *) a. PrimMonad m => Vec (PrimState m) a -> m Int
capacity !Vec (PrimState m) a
v =
MVector (PrimState m) a -> Int
forall s a. MVector s a -> Int
MVector.length (MVector (PrimState m) a -> Int)
-> m (MVector (PrimState m) a) -> m Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MutVar (PrimState m) (MVector (PrimState m) a)
-> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar Vec (PrimState m) a
v.buffer
{-# INLINE capacity #-}
length :: (PrimMonad m) => Vec (PrimState m) a -> m Int
length :: forall (m :: * -> *) a. PrimMonad m => Vec (PrimState m) a -> m Int
length !Vec (PrimState m) a
v = MutVar (PrimState m) Int -> m Int
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar Vec (PrimState m) a
v.len
{-# INLINE length #-}
null :: (PrimMonad m) => Vec (PrimState m) a -> m Bool
null :: forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> m Bool
null !Vec (PrimState m) a
v = (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) (Int -> Bool) -> m Int -> m Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Vec (PrimState m) a -> m Int
forall (m :: * -> *) a. PrimMonad m => Vec (PrimState m) a -> m Int
length Vec (PrimState m) a
v
new :: (PrimMonad m) => Int -> m (Vec (PrimState m) a)
new :: forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (Vec (PrimState m) a)
new = Int -> Int -> m (Vec (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> Int -> m (Vec (PrimState m) a)
newSized Int
0
{-# INLINE new #-}
newSized :: (PrimMonad m) => Int -> Int -> m (Vec (PrimState m) a)
newSized :: forall (m :: * -> *) a.
PrimMonad m =>
Int -> Int -> m (Vec (PrimState m) a)
newSized !Int
n !Int
cap = do
!buffer <- Int -> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (MVector (PrimState m) a)
MVector.new Int
cap
!buffer <- newMutVar buffer
!len <- newMutVar n
pure Vec {buffer, len}
{-# INLINEABLE newSized #-}
slice ::
(PrimMonad m) =>
Int ->
Int ->
Vec (PrimState m) a ->
m (Vec (PrimState m) a)
slice :: forall (m :: * -> *) a.
PrimMonad m =>
Int -> Int -> Vec (PrimState m) a -> m (Vec (PrimState m) a)
slice !Int
i !Int
n !Vec (PrimState m) a
v = do
!len <- Int -> m (MutVar (PrimState m) Int)
forall (m :: * -> *) a.
PrimMonad m =>
a -> m (MutVar (PrimState m) a)
newMutVar Int
n
!mv <- readMutVar v.buffer
!buffer <- newMutVar $! MVector.slice i n mv
pure $! Vec {len, buffer}
{-# INLINEABLE slice #-}
thaw ::
(PrimMonad m) =>
Vector a ->
m (Vec (PrimState m) a)
thaw :: forall (m :: * -> *) a.
PrimMonad m =>
Vector a -> m (Vec (PrimState m) a)
thaw !Vector a
u = do
!buffer <- MVector (PrimState m) a
-> m (MutVar (PrimState m) (MVector (PrimState m) a))
forall (m :: * -> *) a.
PrimMonad m =>
a -> m (MutVar (PrimState m) a)
newMutVar (MVector (PrimState m) a
-> m (MutVar (PrimState m) (MVector (PrimState m) a)))
-> m (MVector (PrimState m) a)
-> m (MutVar (PrimState m) (MVector (PrimState m) a))
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Vector a -> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
Vector a -> m (MVector (PrimState m) a)
Vector.thaw Vector a
u
!len <- newMutVar $! Vector.length u
pure Vec {buffer, len}
{-# INLINEABLE thaw #-}
freeze ::
(PrimMonad m) =>
Vec (PrimState m) a ->
m (Vector a)
freeze :: forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> m (Vector a)
freeze !Vec (PrimState m) a
v = do
!n <- Vec (PrimState m) a -> m Int
forall (m :: * -> *) a. PrimMonad m => Vec (PrimState m) a -> m Int
length Vec (PrimState m) a
v
!mv <- readMutVar v.buffer
Vector.freeze $ MVector.take n mv
{-# INLINEABLE freeze #-}
toList ::
(PrimMonad m) =>
Vec (PrimState m) a ->
m [a]
toList :: forall (m :: * -> *) a. PrimMonad m => Vec (PrimState m) a -> m [a]
toList !Vec (PrimState m) a
vec = Vector a -> [a]
forall a. Vector a -> [a]
Vector.toList (Vector a -> [a]) -> m (Vector a) -> m [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Vec (PrimState m) a -> m (Vector a)
forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> m (Vector a)
freeze Vec (PrimState m) a
vec
ensure_not_oob ::
(HasCallStack, PrimMonad m) =>
String ->
Int ->
Vec (PrimState m) a ->
m ()
ensure_not_oob :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m) =>
String -> Int -> Vec (PrimState m) a -> m ()
ensure_not_oob !String
fname !Int
i !Vec (PrimState m) a
vec = do
!len <- Vec (PrimState m) a -> m Int
forall (m :: * -> *) a. PrimMonad m => Vec (PrimState m) a -> m Int
length Vec (PrimState m) a
vec
when (i < 0 || i >= len) $ do
error $ mconcat [fname, ": index ", show i, " is out bounds ", show len]
ensure ::
(PrimMonad m) =>
Vec (PrimState m) a ->
Int ->
m ()
ensure :: forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> Int -> m ()
ensure !Vec (PrimState m) a
v !Int
cap = do
!current_cap <- Vec (PrimState m) a -> m Int
forall (m :: * -> *) a. PrimMonad m => Vec (PrimState m) a -> m Int
capacity Vec (PrimState m) a
v
unless (current_cap >= cap) $ do
!buffer <- readMutVar v.buffer
!grown <- MVector.grow buffer (cap - current_cap)
writeMutVar v.buffer grown
{-# INLINEABLE ensure #-}
ensureAppend ::
(PrimMonad m) =>
Vec (PrimState m) a ->
Int ->
m ()
ensureAppend :: forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> Int -> m ()
ensureAppend !Vec (PrimState m) a
vec !Int
i = do
!len <- Vec (PrimState m) a -> m Int
forall (m :: * -> *) a. PrimMonad m => Vec (PrimState m) a -> m Int
length Vec (PrimState m) a
vec
!buf <- readMutVar vec.buffer
let !cap = MVector (PrimState m) a -> Int
forall s a. MVector s a -> Int
MVector.length MVector (PrimState m) a
buf
unless (cap >= len + i) $ do
let !newCap = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
ceiling (Double -> Int) -> Double -> Int
forall a b. (a -> b) -> a -> b
$ Double -> Double -> Double
forall a. Ord a => a -> a -> a
max (Double
growFactor Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
cap) (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
cap Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
growFactor Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
cap))
!new_buf <- MVector.grow buf (newCap - cap)
writeMutVar vec.buffer new_buf
where
growFactor :: Double
growFactor :: Double
growFactor = Double
1.5
{-# INLINEABLE ensureAppend #-}
read ::
(HasCallStack, PrimMonad m) =>
Vec (PrimState m) a ->
Int ->
m a
read :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m) =>
Vec (PrimState m) a -> Int -> m a
read !Vec (PrimState m) a
vec !Int
i = do
String -> Int -> Vec (PrimState m) a -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m) =>
String -> Int -> Vec (PrimState m) a -> m ()
ensure_not_oob String
"Vec.read" Int
i Vec (PrimState m) a
vec
!buf <- MutVar (PrimState m) (MVector (PrimState m) a)
-> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar Vec (PrimState m) a
vec.buffer
MVector.unsafeRead buf i
{-# INLINEABLE read #-}
unsafeRead ::
(PrimMonad m) =>
Vec (PrimState m) a ->
Int ->
m a
unsafeRead :: forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> Int -> m a
unsafeRead !Vec (PrimState m) a
vec !Int
i = do
!buf <- MutVar (PrimState m) (MVector (PrimState m) a)
-> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar Vec (PrimState m) a
vec.buffer
MVector.unsafeRead buf i
{-# INLINEABLE unsafeRead #-}
write ::
(PrimMonad m) =>
Vec (PrimState m) a ->
Int ->
a ->
m ()
write :: forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> Int -> a -> m ()
write !Vec (PrimState m) a
vec !Int
i !a
value = do
String -> Int -> Vec (PrimState m) a -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m) =>
String -> Int -> Vec (PrimState m) a -> m ()
ensure_not_oob String
"Vec.write" Int
i Vec (PrimState m) a
vec
!buf <- MutVar (PrimState m) (MVector (PrimState m) a)
-> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar Vec (PrimState m) a
vec.buffer
MVector.unsafeWrite buf i value
{-# INLINEABLE write #-}
unsafeWrite ::
(PrimMonad m) =>
Vec (PrimState m) a ->
Int ->
a ->
m ()
unsafeWrite :: forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> Int -> a -> m ()
unsafeWrite !Vec (PrimState m) a
vec !Int
i !a
value = do
!buf <- MutVar (PrimState m) (MVector (PrimState m) a)
-> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar Vec (PrimState m) a
vec.buffer
MVector.unsafeWrite buf i value
{-# INLINEABLE unsafeWrite #-}
modify ::
(PrimMonad m) =>
Vec (PrimState m) a ->
Int ->
(a -> a) ->
m a
modify :: forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> Int -> (a -> a) -> m a
modify !Vec (PrimState m) a
vec !Int
i !a -> a
f = do
String -> Int -> Vec (PrimState m) a -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m) =>
String -> Int -> Vec (PrimState m) a -> m ()
ensure_not_oob String
"Vec.modify" Int
i Vec (PrimState m) a
vec
!old_val <- Vec (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m) =>
Vec (PrimState m) a -> Int -> m a
read Vec (PrimState m) a
vec Int
i
write vec i (f old_val)
pure old_val
{-# INLINEABLE modify #-}
modify_ ::
(PrimMonad m) =>
Vec (PrimState m) a ->
Int ->
(a -> a) ->
m ()
modify_ :: forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> Int -> (a -> a) -> m ()
modify_ !Vec (PrimState m) a
vec !Int
i !a -> a
f = do
String -> Int -> Vec (PrimState m) a -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m) =>
String -> Int -> Vec (PrimState m) a -> m ()
ensure_not_oob String
"Vec.modify_" Int
i Vec (PrimState m) a
vec
!old_val <- Vec (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m) =>
Vec (PrimState m) a -> Int -> m a
read Vec (PrimState m) a
vec Int
i
write vec i (f old_val)
{-# INLINEABLE modify_ #-}
modifyM ::
(PrimMonad m) =>
Vec (PrimState m) a ->
Int ->
(a -> m a) ->
m a
modifyM :: forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> Int -> (a -> m a) -> m a
modifyM !Vec (PrimState m) a
vec !Int
i !a -> m a
f = do
String -> Int -> Vec (PrimState m) a -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m) =>
String -> Int -> Vec (PrimState m) a -> m ()
ensure_not_oob String
"Vec.modifyM" Int
i Vec (PrimState m) a
vec
!old_val <- Vec (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m) =>
Vec (PrimState m) a -> Int -> m a
read Vec (PrimState m) a
vec Int
i
write vec i =<< f old_val
pure old_val
{-# INLINEABLE modifyM #-}
modifyM_ ::
(PrimMonad m) =>
Vec (PrimState m) a ->
Int ->
(a -> m a) ->
m ()
modifyM_ :: forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> Int -> (a -> m a) -> m ()
modifyM_ !Vec (PrimState m) a
vec !Int
i !a -> m a
f = do
String -> Int -> Vec (PrimState m) a -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m) =>
String -> Int -> Vec (PrimState m) a -> m ()
ensure_not_oob String
"Vec.modifyM_" Int
i Vec (PrimState m) a
vec
!old_val <- Vec (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m) =>
Vec (PrimState m) a -> Int -> m a
read Vec (PrimState m) a
vec Int
i
write vec i =<< f old_val
{-# INLINEABLE modifyM_ #-}
tap ::
(PrimMonad m) =>
Vec (PrimState m) a ->
Int ->
(a -> m ()) ->
m ()
tap :: forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> Int -> (a -> m ()) -> m ()
tap !Vec (PrimState m) a
vec !Int
i !a -> m ()
act = do
String -> Int -> Vec (PrimState m) a -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m) =>
String -> Int -> Vec (PrimState m) a -> m ()
ensure_not_oob String
"Vec.tap" Int
i Vec (PrimState m) a
vec
a -> m ()
act (a -> m ()) -> m a -> m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Vec (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m) =>
Vec (PrimState m) a -> Int -> m a
read Vec (PrimState m) a
vec Int
i
{-# INLINE tap #-}
pushBack ::
(PrimMonad m) =>
Vec (PrimState m) a ->
a ->
m ()
pushBack :: forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> a -> m ()
pushBack !Vec (PrimState m) a
vec !a
value = do
Vec (PrimState m) a -> Int -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> Int -> m ()
ensureAppend Vec (PrimState m) a
vec Int
1
Vec (PrimState m) a -> a -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> a -> m ()
unsafePushBack Vec (PrimState m) a
vec a
value
{-# INLINEABLE pushBack #-}
unsafePushBack ::
(PrimMonad m) =>
Vec (PrimState m) a ->
a ->
m ()
unsafePushBack :: forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> a -> m ()
unsafePushBack !Vec (PrimState m) a
vec !a
a = do
!len <- Vec (PrimState m) a -> m Int
forall (m :: * -> *) a. PrimMonad m => Vec (PrimState m) a -> m Int
length Vec (PrimState m) a
vec
!buf <- readMutVar vec.buffer
MVector.write buf len a
writeMutVar vec.len (len + 1)
{-# INLINEABLE unsafePushBack #-}
swap ::
(PrimMonad m) =>
Vec (PrimState m) a ->
Int ->
Int ->
m ()
swap :: forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> Int -> Int -> m ()
swap !Vec (PrimState m) a
vec !Int
i !Int
j = do
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
j) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
String -> Int -> Vec (PrimState m) a -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m) =>
String -> Int -> Vec (PrimState m) a -> m ()
ensure_not_oob String
"Vec.swap" Int
i Vec (PrimState m) a
vec
String -> Int -> Vec (PrimState m) a -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m) =>
String -> Int -> Vec (PrimState m) a -> m ()
ensure_not_oob String
"Vec.swap" Int
j Vec (PrimState m) a
vec
!buffer <- MutVar (PrimState m) (MVector (PrimState m) a)
-> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar Vec (PrimState m) a
vec.buffer
MVector.swap buffer i j
shrink ::
(PrimMonad m) =>
Vec (PrimState m) a ->
Int ->
m ()
shrink :: forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> Int -> m ()
shrink !Vec (PrimState m) a
vec !Int
amount = do
!old_len <- Vec (PrimState m) a -> m Int
forall (m :: * -> *) a. PrimMonad m => Vec (PrimState m) a -> m Int
length Vec (PrimState m) a
vec
writeMutVar vec.len (old_len - amount)
takeSwap ::
(PrimMonad m) =>
Vec (PrimState m) a ->
Int ->
m a
takeSwap :: forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> Int -> m a
takeSwap !Vec (PrimState m) a
vec !Int
i = do
!value <- Vec (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m) =>
Vec (PrimState m) a -> Int -> m a
read Vec (PrimState m) a
vec Int
i
removeSwap vec i
pure value
removeSwap ::
(PrimMonad m) =>
Vec (PrimState m) a ->
Int ->
m ()
removeSwap :: forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> Int -> m ()
removeSwap !Vec (PrimState m) a
vec !Int
i = do
!len <- Vec (PrimState m) a -> m Int
forall (m :: * -> *) a. PrimMonad m => Vec (PrimState m) a -> m Int
length Vec (PrimState m) a
vec
swap vec i (len - 1)
shrink vec 1
clone :: (PrimMonad m) => Vec (PrimState m) a -> m (Vec (PrimState m) a)
clone :: forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> m (Vec (PrimState m) a)
clone !Vec (PrimState m) a
vec = do
!list <- Vec (PrimState m) a -> m [a]
forall (m :: * -> *) a. PrimMonad m => Vec (PrimState m) a -> m [a]
toList Vec (PrimState m) a
vec
!len <- length vec
!newVec <- new len
for_ list $ pushBack newVec
pure newVec