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 Debug.Trace (traceShow)
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
return newVec