aztecs-0.11.0: src/Aztecs/ECS/Query/Dynamic.hs
{-# LANGUAGE ApplicativeDo #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeFamilies #-}
-- |
-- Module : Aztecs.ECS.Query.Dynamic
-- Copyright : (c) Matt Hunzinger, 2025
-- License : BSD-style (see the LICENSE file in the distribution)
--
-- Maintainer : matt@hunzinger.me
-- Stability : provisional
-- Portability : non-portable (GHC extensions)
module Aztecs.ECS.Query.Dynamic
( -- * Dynamic queries
DynamicQuery,
DynamicQueryT (..),
-- ** Operations
entityDyn,
fetchDyn,
fetchMaybeDyn,
fetchMapDyn,
fetchMapDynM,
zipFetchMapDyn,
zipFetchMapAccumDyn,
zipFetchMapDynM,
zipFetchMapAccumDynM,
-- ** Filters
withDyn,
withoutDyn,
-- ** Running
queryDyn,
singleDyn,
singleMaybeDyn,
queryEntitiesDyn,
readQueryDyn,
mapSingleDyn,
mapSingleMaybeDyn,
readQueryEntitiesDyn,
-- *** Internal
QueryFilter (..),
queryFilter,
runDynQuery,
runDynQueryEntities,
readDynQuery,
readDynQueryEntities,
)
where
import Aztecs.ECS.Component
import Aztecs.ECS.Entity
import Aztecs.ECS.World.Archetype (Archetype)
import qualified Aztecs.ECS.World.Archetype as A
import Aztecs.ECS.World.Archetypes (ArchetypeID, Node (..))
import qualified Aztecs.ECS.World.Archetypes as AS
import Aztecs.ECS.World.Entities
import Control.Applicative
import Control.Monad.Identity
import Data.Bifunctor
import Data.Foldable
import Data.Map (Map)
import qualified Data.Map.Strict as Map
import Data.Maybe
import Data.Set (Set)
import qualified Data.Set as Set
import GHC.Stack
import Prelude hiding (reads)
data Operation f a where
Entity :: Operation f EntityID
Fetch :: (Component a) => !ComponentID -> Operation f a
FetchMaybe :: (Component a) => !ComponentID -> Operation f (Maybe a)
FetchMap :: (Component a) => !(a -> a) -> !ComponentID -> Operation f a
FetchMapM :: (Monad f, Component a) => !(a -> f a) -> !ComponentID -> Operation f a
Adjust :: (Component a) => !(b -> a -> (c, a)) -> !ComponentID -> !(DynamicQueryT f b) -> Operation f (c, a)
AdjustM :: (Monad f, Component a) => !(b -> a -> f (c, a)) -> !ComponentID -> !(DynamicQueryT f b) -> Operation f (c, a)
With :: !ComponentID -> Operation f ()
Without :: !ComponentID -> Operation f ()
-- @since 0.9
type DynamicQuery = DynamicQueryT Identity
-- | Dynamic query for components by ID.
--
-- @since 0.11
data DynamicQueryT f a where
Pure :: !a -> DynamicQueryT f a
Map :: !(a -> b) -> !(DynamicQueryT f a) -> DynamicQueryT f b
Ap :: !(DynamicQueryT f (a -> b)) -> !(DynamicQueryT f a) -> DynamicQueryT f b
Op :: !(Operation f a) -> DynamicQueryT f a
instance Functor (DynamicQueryT f) where
{-# INLINE fmap #-}
fmap = Map
-- | @since 0.11
instance Applicative (DynamicQueryT f) where
{-# INLINE pure #-}
pure = Pure
{-# INLINE (<*>) #-}
(<*>) = Ap
{-# INLINE entityDyn #-}
entityDyn :: DynamicQueryT f EntityID
entityDyn = Op Entity
{-# INLINE fetchDyn #-}
fetchDyn :: (Component a) => ComponentID -> DynamicQueryT f a
fetchDyn = Op . Fetch
{-# INLINE fetchMaybeDyn #-}
fetchMaybeDyn :: (Component a) => ComponentID -> DynamicQueryT f (Maybe a)
fetchMaybeDyn = Op . FetchMaybe
{-# INLINE fetchMapDyn #-}
fetchMapDyn :: (Component a) => (a -> a) -> ComponentID -> DynamicQueryT f a
fetchMapDyn f = Op . FetchMap f
{-# INLINE fetchMapDynM #-}
fetchMapDynM :: (Monad f, Component a) => (a -> f a) -> ComponentID -> DynamicQueryT f a
fetchMapDynM f = Op . FetchMapM f
{-# INLINE zipFetchMapDyn #-}
zipFetchMapDyn :: (Component a) => (b -> a -> a) -> ComponentID -> DynamicQueryT f b -> DynamicQueryT f a
zipFetchMapDyn f cId q = snd <$> Op (Adjust (\b a -> ((), f b a)) cId q)
{-# INLINE zipFetchMapAccumDyn #-}
zipFetchMapAccumDyn :: (Component a) => (b -> a -> (c, a)) -> ComponentID -> DynamicQueryT f b -> DynamicQueryT f (c, a)
zipFetchMapAccumDyn f cId q = Op $ Adjust f cId q
{-# INLINE zipFetchMapDynM #-}
zipFetchMapDynM :: (Monad f, Component a) => (b -> a -> f a) -> ComponentID -> DynamicQueryT f b -> DynamicQueryT f a
zipFetchMapDynM f cId q = snd <$> zipFetchMapAccumDynM (\b a -> ((),) <$> f b a) cId q
{-# INLINE zipFetchMapAccumDynM #-}
zipFetchMapAccumDynM :: (Monad f, Component a) => (b -> a -> f (c, a)) -> ComponentID -> DynamicQueryT f b -> DynamicQueryT f (c, a)
zipFetchMapAccumDynM f cId q = Op $ AdjustM f cId q
{-# INLINE withDyn #-}
withDyn :: ComponentID -> DynamicQueryT f ()
withDyn = Op . With
{-# INLINE withoutDyn #-}
withoutDyn :: ComponentID -> DynamicQueryT f ()
withoutDyn = Op . Without
{-# INLINE opFilter #-}
opFilter :: Operation f a -> QueryFilter
opFilter Entity = mempty
opFilter (Fetch cId) = mempty {filterWith = Set.singleton cId}
opFilter (FetchMaybe cId) = mempty {filterWith = Set.singleton cId}
opFilter (FetchMap _ cId) = mempty {filterWith = Set.singleton cId}
opFilter (FetchMapM _ cId) = mempty {filterWith = Set.singleton cId}
opFilter (Adjust _ cId q) = queryFilter q <> mempty {filterWith = Set.singleton cId}
opFilter (AdjustM _ cId q) = queryFilter q <> mempty {filterWith = Set.singleton cId}
opFilter (With cId) = mempty {filterWith = Set.singleton cId}
opFilter (Without cId) = mempty {filterWithout = Set.singleton cId}
{-# INLINE queryFilter #-}
queryFilter :: DynamicQueryT f a -> QueryFilter
queryFilter (Pure _) = mempty
queryFilter (Map _ q) = queryFilter q
queryFilter (Ap f g) = queryFilter f <> queryFilter g
queryFilter (Op op) = opFilter op
{-# INLINE runOp #-}
runOp :: (Applicative f) => Operation f a -> Archetype -> f ([a], Archetype)
runOp Entity arch = pure (Set.toList $ A.entities arch, mempty)
runOp (Fetch cId) arch = pure (A.lookupComponentsAsc cId arch, mempty)
runOp (FetchMaybe cId) arch =
pure
( case A.lookupComponentsAscMaybe cId arch of
Just as -> fmap Just as
Nothing -> replicate (length $ A.entities arch) Nothing,
mempty
)
runOp (FetchMap f cId) arch = pure $ A.map f cId arch
runOp (FetchMapM f cId) arch = do
(as, arch') <- A.mapM f cId arch
return (as, arch')
runOp (Adjust f cId q) arch = do
res <- runDynQuery q arch
return $
let !(bs, arch') = res
!(as, arch'') = A.zipWith bs f cId arch
in (as, arch'' <> arch')
runOp (AdjustM f cId q) arch = do
(as, arch') <- runDynQuery q arch
(bs, arch'') <- A.zipWithM as f cId arch
return (bs, arch'' <> arch')
runOp (With _) _ = pure ([], mempty)
runOp (Without _) _ = pure ([], mempty)
{-# INLINE readOp #-}
readOp :: (Applicative f) => Operation f a -> Archetype -> f [a]
readOp Entity arch = pure $ Set.toList $ A.entities arch
readOp (Fetch cId) arch = pure $ A.lookupComponentsAsc cId arch
readOp (FetchMaybe cId) arch =
pure $
case A.lookupComponentsAscMaybe cId arch of
Just as -> fmap Just as
Nothing -> replicate (length $ A.entities arch) Nothing
readOp (FetchMap f cId) arch = do
bs <- readOp (Fetch cId) arch
return $ map f bs
readOp (FetchMapM f cId) arch = do
bs <- readOp (Fetch cId) arch
mapM f bs
readOp (Adjust f cId q) arch = do
as <- readDynQuery q arch
bs <- readOp (Fetch cId) arch
return $ zipWith f as bs
readOp (AdjustM f cId q) arch = do
as <- readDynQuery q arch
bs <- readOp (Fetch cId) arch
zipWithM f as bs
readOp (With _) _ = pure []
readOp (Without _) _ = pure []
{-# INLINE queryEntitiesDyn #-}
queryEntitiesDyn ::
(Monad m) =>
[EntityID] ->
DynamicQueryT m a ->
Entities ->
m ([a], Entities)
queryEntitiesDyn eIds q es =
let qf = queryFilter q
go = runDynQueryEntities eIds q
in if Set.null $ filterWith qf
then do
(as, _) <- go A.empty {A.entities = Map.keysSet $ entities es}
return (as, es)
else
let go' (acc, esAcc) (aId, n) = do
(as', arch') <- go $ nodeArchetype n
let n' = n {nodeArchetype = arch' <> nodeArchetype n}
!nodes = Map.insert aId n' . AS.nodes $ archetypes esAcc
return (as' ++ acc, esAcc {archetypes = (archetypes esAcc) {AS.nodes = nodes}})
in foldlM go' ([], es) $ Map.toList . AS.find (filterWith qf) (filterWithout qf) $ archetypes es
readQueryEntitiesDyn :: (Monad m) => [EntityID] -> DynamicQueryT m a -> Entities -> m [a]
readQueryEntitiesDyn eIds q es =
let qf = queryFilter q
in if Set.null $ filterWith qf
then readDynQueryEntities eIds q A.empty {A.entities = Map.keysSet $ entities es}
else
let go n = readDynQuery q $ AS.nodeArchetype n
in concat <$> mapM go (AS.find (filterWith qf) (filterWithout qf) $ archetypes es)
runOpEntities :: (Applicative f) => Operation f a -> [EntityID] -> Archetype -> f ([a], Archetype)
runOpEntities Entity es _ = pure (es, mempty)
runOpEntities (Fetch cId) es arch =
pure
( map snd
. filter (\(e, _) -> e `elem` es)
. Map.toList
$ A.lookupComponents cId arch,
mempty
)
runOpEntities (FetchMaybe cId) es arch =
pure
( map (\(e, a) -> if e `elem` es then Just a else Nothing)
. Map.toList
$ A.lookupComponents cId arch,
mempty
)
runOpEntities (FetchMap f cId) es arch =
pure $
let go e a =
if e `elem` es
then let a' = f a in (Just a', a')
else (Nothing, a)
!(as, arch') = A.zipWith es go cId arch
in (mapMaybe fst as, arch')
runOpEntities (FetchMapM f cId) es arch = do
(as, arch') <- runOpEntities (AdjustM (\() a -> (,a) <$> f a) cId (pure ())) es arch
return (map snd as, arch')
runOpEntities (Adjust f cId q) es arch = do
res <- runDynQuery q arch
return $
let go (e, b) a =
if e `elem` es
then let (x, y) = f b a in (Just x, y)
else (Nothing, a)
!(bs, arch') = res
!(as, arch'') = A.zipWith (zip es bs) go cId arch
in (mapMaybe (\(m, b) -> fmap (,b) m) as, arch'' <> arch')
runOpEntities (AdjustM f cId q) es arch = do
(bs, arch') <- runDynQuery q arch
let go (e, b) a =
if e `elem` es
then do
(x, y) <- f b a
return (Just x, y)
else return (Nothing, a)
(as, arch'') <- A.zipWithM (zip es bs) go cId arch
return (mapMaybe (\(m, b) -> fmap (,b) m) as, arch'' <> arch')
runOpEntities (With _) _ arch = pure ([], arch)
runOpEntities (Without _) _ arch = pure ([], arch)
runDynQueryEntities :: (Applicative f) => [EntityID] -> DynamicQueryT f a -> Archetype -> f ([a], Archetype)
runDynQueryEntities es (Pure a) _ = pure (replicate (length es) a, mempty)
runDynQueryEntities es (Map f q) arch = first (fmap f) <$> runDynQueryEntities es q arch
runDynQueryEntities es (Ap f g) arch = do
res <- runDynQueryEntities es g arch
res' <- runDynQueryEntities es f arch
return $
let (as, arch') = res
(bs, arch'') = res'
in (zipWith ($) bs as, arch'' <> arch')
runDynQueryEntities es (Op op) arch = runOpEntities op es arch
{-# INLINE readOpEntities #-}
readOpEntities :: (Applicative f) => Operation f a -> [EntityID] -> Archetype -> f [a]
readOpEntities Entity es _ = pure es
readOpEntities (Fetch cId) es arch =
pure
. map snd
. filter (\(e, _) -> e `elem` es)
. Map.toList
$ A.lookupComponents cId arch
readOpEntities (FetchMaybe cId) es arch =
pure
. map (\(e, a) -> if e `elem` es then Just a else Nothing)
. Map.toList
$ A.lookupComponents cId arch
readOpEntities (FetchMap f cId) es arch = do
b <- readOpEntities (Fetch cId) es arch
pure $ map f b
readOpEntities (FetchMapM f cId) es arch = do
b <- readOpEntities (Fetch cId) es arch
mapM f b
readOpEntities (Adjust f cId q) es arch = do
a <- readDynQueryEntities es q arch
b <- readOpEntities (Fetch cId) es arch
pure $ zipWith f a b
readOpEntities (AdjustM f cId q) es arch = do
a <- readDynQueryEntities es q arch
b <- readOpEntities (Fetch cId) es arch
zipWithM f a b
readOpEntities (With _) _ _ = pure []
readOpEntities (Without _) _ _ = pure []
{-# INLINE readDynQueryEntities #-}
readDynQueryEntities :: (Applicative f) => [EntityID] -> DynamicQueryT f a -> Archetype -> f [a]
readDynQueryEntities es (Pure a) _ = pure $ replicate (length es) a
readDynQueryEntities es (Map f q) arch = fmap f <$> readDynQueryEntities es q arch
readDynQueryEntities es (Ap f g) arch = do
a <- readDynQueryEntities es g arch
b <- readDynQueryEntities es f arch
pure $ b <*> a
readDynQueryEntities es (Op op) arch = readOpEntities op es arch
{-# INLINE runDynQuery #-}
runDynQuery :: (Applicative f) => DynamicQueryT f a -> Archetype -> f ([a], Archetype)
runDynQuery (Pure a) arch = pure (replicate (length $ A.entities arch) a, mempty)
runDynQuery (Map f q) arch = do
res <- runDynQuery q arch
return $ first (fmap f) res
runDynQuery (Ap f g) arch = do
res <- runDynQuery g arch
res' <- runDynQuery f arch
return $
let (as, arch') = res
(bs, arch'') = res'
in (zipWith ($) bs as, arch'' <> arch')
runDynQuery (Op op) arch = runOp op arch
{-# INLINE readDynQuery #-}
readDynQuery :: (Applicative f) => DynamicQueryT f a -> Archetype -> f [a]
readDynQuery (Pure a) arch = pure $ replicate (length $ A.entities arch) a
readDynQuery (Map f q) arch = fmap f <$> readDynQuery q arch
readDynQuery (Ap f g) arch = do
as <- readDynQuery g arch
bs <- readDynQuery f arch
pure $ zipWith ($) bs as
readDynQuery (Op op) arch = readOp op arch
-- | Match all entities.
--
-- @since 0.11
readQueryDyn :: (Monad m) => DynamicQueryT m a -> Entities -> m [a]
readQueryDyn q es =
let qf = queryFilter q
in if Set.null $ filterWith qf
then readDynQuery q $ A.empty {A.entities = Map.keysSet $ entities es}
else
let go n = readDynQuery q $ AS.nodeArchetype n
in concat <$> mapM go (AS.find (filterWith qf) (filterWithout qf) $ archetypes es)
-- | Match a single entity.
--
-- @since 0.11
singleDyn :: (HasCallStack, Monad m) => DynamicQueryT m a -> Entities -> m a
singleDyn q es = do
res <- singleMaybeDyn q es
return $ case res of
Just a -> a
_ -> error "singleDyn: expected a single entity"
-- | Match a single entity, or `Nothing`.
--
-- @since 0.11
singleMaybeDyn :: (Monad m) => DynamicQueryT m a -> Entities -> m (Maybe a)
singleMaybeDyn q es =
let qf = queryFilter q
in if Set.null $ filterWith qf
then case Map.keys $ entities es of
[eId] -> do
res <- readDynQuery q $ A.singleton eId
return $ case res of
[a] -> Just a
_ -> Nothing
_ -> return Nothing
else case Map.elems $ AS.find (filterWith qf) (filterWithout qf) $ archetypes es of
[n] -> do
res <- readDynQuery q $ AS.nodeArchetype n
return $ case res of
[a] -> Just a
_ -> Nothing
_ -> return Nothing
-- | Match and update all matched entities.
--
-- @since 0.11
{-# INLINE queryDyn #-}
queryDyn :: (Monad m) => DynamicQueryT m a -> Entities -> m ([a], Entities)
queryDyn = mapDyn' id
{-# INLINE mapDyn' #-}
mapDyn' ::
(Monad m) =>
(Map ArchetypeID Node -> Map ArchetypeID Node) ->
DynamicQueryT m a ->
Entities ->
m ([a], Entities)
mapDyn' f q es =
let !qf = queryFilter q
go = runDynQuery q
in if Set.null $ filterWith qf
then do
(as, _) <- go A.empty {A.entities = Map.keysSet $ entities es}
return (as, es)
else
let go' (acc, esAcc) (aId, n) = do
(as', arch') <- go $ nodeArchetype n
let n' = n {nodeArchetype = arch' <> nodeArchetype n}
!nodes = Map.insert aId n' . AS.nodes $ archetypes esAcc
return (as' ++ acc, esAcc {archetypes = (archetypes esAcc) {AS.nodes = nodes}})
in foldlM go' ([], es) $ Map.toList . f . AS.find (filterWith qf) (filterWithout qf) $ archetypes es
-- | Map a single matched entity.
--
-- @since 0.11
mapSingleDyn :: (HasCallStack, Monad m) => DynamicQueryT m a -> Entities -> m (a, Entities)
mapSingleDyn q es = do
(res, es') <- mapSingleMaybeDyn q es
return $ case res of
Just a -> (a, es')
_ -> error "mapSingleDyn: expected single matching entity"
-- | Map a single matched entity, or @Nothing@.
--
-- @since 0.11
{-# INLINE mapSingleMaybeDyn #-}
mapSingleMaybeDyn :: (Monad m) => DynamicQueryT m a -> Entities -> m (Maybe a, Entities)
mapSingleMaybeDyn q es =
let qf = queryFilter q
in if Set.null $ filterWith qf
then case Map.keys $ entities es of
[eId] -> do
res <- runDynQuery q $ A.singleton eId
return $ case res of
([a], _) -> (Just a, es)
_ -> (Nothing, es)
_ -> pure (Nothing, es)
else case Map.toList $ AS.find (filterWith qf) (filterWithout qf) $ archetypes es of
[(aId, n)] -> do
(as, arch') <- runDynQuery q $ AS.nodeArchetype n
return $ case as of
[a] ->
let nodes = Map.insert aId n {nodeArchetype = arch' <> nodeArchetype n} . AS.nodes $ archetypes es
in (Just a, es {archetypes = (archetypes es) {AS.nodes = nodes}})
_ -> (Nothing, es)
_ -> pure (Nothing, es)
-- | `Query` filter.
--
-- @since 0.11
data QueryFilter = QueryFilter
{ filterWith :: !(Set ComponentID),
filterWithout :: !(Set ComponentID)
}
deriving (Show)
-- | @since 0.9
instance Semigroup QueryFilter where
QueryFilter r1 w1 <> QueryFilter r2 w2 = QueryFilter (r1 <> r2) (w1 <> w2)
-- | @since 0.9
instance Monoid QueryFilter where
mempty = QueryFilter mempty mempty