{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeSynonymInstances #-}

-- |
-- Module      : Aztecs.ECS.System
-- 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.System
  ( -- * Systems
    System (..),

    -- * Dynamic systems
    DynamicSystem (..),
    runDynamicSystem,

    -- ** Queries
    readQuery,
    readQueryFiltered,
    readQuerySingle,
    readQuerySingleMaybe,
    runQuery,
    runQueryFiltered,
    runQuerySingle,
    runQuerySingleMaybe,

    -- ** Dynamic queries
    readQueryDyn,
    readQueryFilteredDyn,
    readQuerySingleMaybeDyn,
    runQueryDyn,
    runQueryFilteredDyn,
    runQuerySingleMaybeDyn,
  )
where

import Aztecs.ECS.Component
import Aztecs.ECS.Query (Query (..), QueryFilter (..))
import qualified Aztecs.ECS.Query as Q
import Aztecs.ECS.Query.Dynamic (DynamicQuery, DynamicQueryFilter (..))
import Aztecs.ECS.System.Dynamic (DynamicSystem (..), runDynamicSystem)
import qualified Aztecs.ECS.System.Dynamic as DS
import qualified Aztecs.ECS.World.Archetype as A
import Aztecs.ECS.World.Archetypes (Node (..))
import Aztecs.ECS.World.Components (Components)
import qualified Data.Foldable as F
import Data.Set (Set)
import Data.Vector (Vector)
import GHC.Stack
import Prelude hiding (all, filter, map, mapM)

-- | System for querying entities.
newtype System m a = System {forall (m :: * -> *) a.
System m a -> Components -> (Components, DynamicSystem m a)
runSystem :: Components -> (Components, DynamicSystem m a)}

instance Functor (System m) where
  fmap :: forall a b. (a -> b) -> System m a -> System m b
fmap a -> b
f (System Components -> (Components, DynamicSystem m a)
g) = (Components -> (Components, DynamicSystem m b)) -> System m b
forall (m :: * -> *) a.
(Components -> (Components, DynamicSystem m a)) -> System m a
System ((Components -> (Components, DynamicSystem m b)) -> System m b)
-> (Components -> (Components, DynamicSystem m b)) -> System m b
forall a b. (a -> b) -> a -> b
$ \Components
cs ->
    let !(Components
cs', DynamicSystem m a
dynS) = Components -> (Components, DynamicSystem m a)
g Components
cs in (Components
cs', (a -> b) -> DynamicSystem m a -> DynamicSystem m b
forall a b. (a -> b) -> DynamicSystem m a -> DynamicSystem m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f DynamicSystem m a
dynS)
  {-# INLINE fmap #-}

instance Applicative (System m) where
  pure :: forall a. a -> System m a
pure a
a = (Components -> (Components, DynamicSystem m a)) -> System m a
forall (m :: * -> *) a.
(Components -> (Components, DynamicSystem m a)) -> System m a
System (,a -> DynamicSystem m a
forall a (m :: * -> *). a -> DynamicSystem m a
Pure a
a)
  {-# INLINE pure #-}

  (System Components -> (Components, DynamicSystem m (a -> b))
f) <*> :: forall a b. System m (a -> b) -> System m a -> System m b
<*> (System Components -> (Components, DynamicSystem m a)
g) = (Components -> (Components, DynamicSystem m b)) -> System m b
forall (m :: * -> *) a.
(Components -> (Components, DynamicSystem m a)) -> System m a
System ((Components -> (Components, DynamicSystem m b)) -> System m b)
-> (Components -> (Components, DynamicSystem m b)) -> System m b
forall a b. (a -> b) -> a -> b
$ \Components
cs ->
    let !(Components
cs', DynamicSystem m (a -> b)
dynF) = Components -> (Components, DynamicSystem m (a -> b))
f Components
cs
        !(Components
cs'', DynamicSystem m a
dynG) = Components -> (Components, DynamicSystem m a)
g Components
cs'
     in (Components
cs'', DynamicSystem m (a -> b)
dynF DynamicSystem m (a -> b) -> DynamicSystem m a -> DynamicSystem m b
forall a b.
DynamicSystem m (a -> b) -> DynamicSystem m a -> DynamicSystem m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> DynamicSystem m a
dynG)
  {-# INLINE (<*>) #-}

runner :: (Set ComponentID -> DynamicQuery m a -> DynamicSystem m b) -> Query m a -> System m b
runner :: forall (m :: * -> *) a b.
(Set ComponentID -> DynamicQuery m a -> DynamicSystem m b)
-> Query m a -> System m b
runner Set ComponentID -> DynamicQuery m a -> DynamicSystem m b
f Query m a
q = (Components -> (Components, DynamicSystem m b)) -> System m b
forall (m :: * -> *) a.
(Components -> (Components, DynamicSystem m a)) -> System m a
System ((Components -> (Components, DynamicSystem m b)) -> System m b)
-> (Components -> (Components, DynamicSystem m b)) -> System m b
forall a b. (a -> b) -> a -> b
$ \Components
cs ->
  let (ReadsWrites
rws, Components
cs', DynamicQuery m a
dynQ) = Query m a
-> Components -> (ReadsWrites, Components, DynamicQuery m a)
forall (m :: * -> *) a.
Query m a
-> Components -> (ReadsWrites, Components, DynamicQuery m a)
runQuery' Query m a
q Components
cs
   in (Components
cs', Set ComponentID -> DynamicQuery m a -> DynamicSystem m b
f (ReadsWrites -> Set ComponentID
Q.reads ReadsWrites
rws Set ComponentID -> Set ComponentID -> Set ComponentID
forall a. Semigroup a => a -> a -> a
<> ReadsWrites -> Set ComponentID
Q.writes ReadsWrites
rws) DynamicQuery m a
dynQ)

-- | Match all entities.
readQuery :: (Monad m) => Query m a -> System m (Vector a)
readQuery :: forall (m :: * -> *) a. Monad m => Query m a -> System m (Vector a)
readQuery = (Set ComponentID -> DynamicQuery m a -> DynamicSystem m (Vector a))
-> Query m a -> System m (Vector a)
forall (m :: * -> *) a b.
(Set ComponentID -> DynamicQuery m a -> DynamicSystem m b)
-> Query m a -> System m b
runner Set ComponentID -> DynamicQuery m a -> DynamicSystem m (Vector a)
forall (m :: * -> *) a.
Set ComponentID -> DynamicQuery m a -> DynamicSystem m (Vector a)
DS.readQuery

readQuerySingle :: (HasCallStack, Monad m) => Query m a -> System m a
readQuerySingle :: forall (m :: * -> *) a.
(HasCallStack, Monad m) =>
Query m a -> System m a
readQuerySingle = (Set ComponentID -> DynamicQuery m a -> DynamicSystem m a)
-> Query m a -> System m a
forall (m :: * -> *) a b.
(Set ComponentID -> DynamicQuery m a -> DynamicSystem m b)
-> Query m a -> System m b
runner Set ComponentID -> DynamicQuery m a -> DynamicSystem m a
forall (m :: * -> *) a.
Set ComponentID -> DynamicQuery m a -> DynamicSystem m a
DS.readQuerySingle

readQuerySingleMaybe :: (Monad m) => Query m a -> System m (Maybe a)
readQuerySingleMaybe :: forall (m :: * -> *) a. Monad m => Query m a -> System m (Maybe a)
readQuerySingleMaybe = (Set ComponentID -> DynamicQuery m a -> DynamicSystem m (Maybe a))
-> Query m a -> System m (Maybe a)
forall (m :: * -> *) a b.
(Set ComponentID -> DynamicQuery m a -> DynamicSystem m b)
-> Query m a -> System m b
runner Set ComponentID -> DynamicQuery m a -> DynamicSystem m (Maybe a)
forall (m :: * -> *) a.
Set ComponentID -> DynamicQuery m a -> DynamicSystem m (Maybe a)
DS.readQuerySingleMaybe

-- | Match all entities with a filter.
readQueryFiltered :: (Monad m) => Query m a -> QueryFilter -> System m (Vector a)
readQueryFiltered :: forall (m :: * -> *) a.
Monad m =>
Query m a -> QueryFilter -> System m (Vector a)
readQueryFiltered Query m a
q QueryFilter
f = (Components -> (Components, DynamicSystem m (Vector a)))
-> System m (Vector a)
forall (m :: * -> *) a.
(Components -> (Components, DynamicSystem m a)) -> System m a
System ((Components -> (Components, DynamicSystem m (Vector a)))
 -> System m (Vector a))
-> (Components -> (Components, DynamicSystem m (Vector a)))
-> System m (Vector a)
forall a b. (a -> b) -> a -> b
$ \Components
cs ->
  let (ReadsWrites
rws, Components
cs', DynamicQuery m a
dynQ) = Query m a
-> Components -> (ReadsWrites, Components, DynamicQuery m a)
forall (m :: * -> *) a.
Query m a
-> Components -> (ReadsWrites, Components, DynamicQuery m a)
runQuery' Query m a
q Components
cs
      (DynamicQueryFilter
dynF, Components
cs'') = QueryFilter -> Components -> (DynamicQueryFilter, Components)
runQueryFilter QueryFilter
f Components
cs'
      flt :: Node m -> Bool
flt Node m
n =
        (ComponentID -> Bool) -> Set ComponentID -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
F.all (\ComponentID
cId -> ComponentID -> Archetype m -> Bool
forall (m :: * -> *). ComponentID -> Archetype m -> Bool
A.member ComponentID
cId (Archetype m -> Bool) -> Archetype m -> Bool
forall a b. (a -> b) -> a -> b
$ Node m -> Archetype m
forall (m :: * -> *). Node m -> Archetype m
nodeArchetype Node m
n) (DynamicQueryFilter -> Set ComponentID
filterWith DynamicQueryFilter
dynF)
          Bool -> Bool -> Bool
&& (ComponentID -> Bool) -> Set ComponentID -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
F.all (\ComponentID
cId -> Bool -> Bool
not (ComponentID -> Archetype m -> Bool
forall (m :: * -> *). ComponentID -> Archetype m -> Bool
A.member ComponentID
cId (Archetype m -> Bool) -> Archetype m -> Bool
forall a b. (a -> b) -> a -> b
$ Node m -> Archetype m
forall (m :: * -> *). Node m -> Archetype m
nodeArchetype Node m
n)) (DynamicQueryFilter -> Set ComponentID
filterWithout DynamicQueryFilter
dynF)
   in (Components
cs'', Set ComponentID
-> (Node m -> Bool)
-> DynamicQuery m a
-> DynamicSystem m (Vector a)
forall (m :: * -> *) a.
Set ComponentID
-> (Node m -> Bool)
-> DynamicQuery m a
-> DynamicSystem m (Vector a)
DS.readQueryFiltered (ReadsWrites -> Set ComponentID
Q.reads ReadsWrites
rws Set ComponentID -> Set ComponentID -> Set ComponentID
forall a. Semigroup a => a -> a -> a
<> ReadsWrites -> Set ComponentID
Q.writes ReadsWrites
rws) Node m -> Bool
flt DynamicQuery m a
dynQ)

-- | Map all matching entities.
runQuery :: (Monad m) => Query m a -> System m (Vector a)
runQuery :: forall (m :: * -> *) a. Monad m => Query m a -> System m (Vector a)
runQuery = (Set ComponentID -> DynamicQuery m a -> DynamicSystem m (Vector a))
-> Query m a -> System m (Vector a)
forall (m :: * -> *) a b.
(Set ComponentID -> DynamicQuery m a -> DynamicSystem m b)
-> Query m a -> System m b
runner Set ComponentID -> DynamicQuery m a -> DynamicSystem m (Vector a)
forall (m :: * -> *) a.
Set ComponentID -> DynamicQuery m a -> DynamicSystem m (Vector a)
DS.runQuery

runQuerySingle :: (HasCallStack, Monad m) => Query m a -> System m a
runQuerySingle :: forall (m :: * -> *) a.
(HasCallStack, Monad m) =>
Query m a -> System m a
runQuerySingle = (Set ComponentID -> DynamicQuery m a -> DynamicSystem m a)
-> Query m a -> System m a
forall (m :: * -> *) a b.
(Set ComponentID -> DynamicQuery m a -> DynamicSystem m b)
-> Query m a -> System m b
runner Set ComponentID -> DynamicQuery m a -> DynamicSystem m a
forall (m :: * -> *) a.
Set ComponentID -> DynamicQuery m a -> DynamicSystem m a
DS.runQuerySingle

-- | Map a single matching entity, or @Nothing@.
runQuerySingleMaybe :: (Monad m) => Query m a -> System m (Maybe a)
runQuerySingleMaybe :: forall (m :: * -> *) a. Monad m => Query m a -> System m (Maybe a)
runQuerySingleMaybe = (Set ComponentID -> DynamicQuery m a -> DynamicSystem m (Maybe a))
-> Query m a -> System m (Maybe a)
forall (m :: * -> *) a b.
(Set ComponentID -> DynamicQuery m a -> DynamicSystem m b)
-> Query m a -> System m b
runner Set ComponentID -> DynamicQuery m a -> DynamicSystem m (Maybe a)
forall (m :: * -> *) a.
Set ComponentID -> DynamicQuery m a -> DynamicSystem m (Maybe a)
DS.runQuerySingleMaybe

-- | Filter and map all matching entities.
runQueryFiltered :: (Monad m) => Query m a -> QueryFilter -> System m (Vector a)
runQueryFiltered :: forall (m :: * -> *) a.
Monad m =>
Query m a -> QueryFilter -> System m (Vector a)
runQueryFiltered Query m a
q QueryFilter
f = (Components -> (Components, DynamicSystem m (Vector a)))
-> System m (Vector a)
forall (m :: * -> *) a.
(Components -> (Components, DynamicSystem m a)) -> System m a
System ((Components -> (Components, DynamicSystem m (Vector a)))
 -> System m (Vector a))
-> (Components -> (Components, DynamicSystem m (Vector a)))
-> System m (Vector a)
forall a b. (a -> b) -> a -> b
$ \Components
cs ->
  let (ReadsWrites
rws, Components
cs', DynamicQuery m a
dynQ) = Query m a
-> Components -> (ReadsWrites, Components, DynamicQuery m a)
forall (m :: * -> *) a.
Query m a
-> Components -> (ReadsWrites, Components, DynamicQuery m a)
runQuery' Query m a
q Components
cs
      (DynamicQueryFilter
dynF, Components
cs'') = QueryFilter -> Components -> (DynamicQueryFilter, Components)
runQueryFilter QueryFilter
f Components
cs'
      flt :: Node m -> Bool
flt Node m
n =
        (ComponentID -> Bool) -> Set ComponentID -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
F.all (\ComponentID
cId -> ComponentID -> Archetype m -> Bool
forall (m :: * -> *). ComponentID -> Archetype m -> Bool
A.member ComponentID
cId (Archetype m -> Bool) -> Archetype m -> Bool
forall a b. (a -> b) -> a -> b
$ Node m -> Archetype m
forall (m :: * -> *). Node m -> Archetype m
nodeArchetype Node m
n) (DynamicQueryFilter -> Set ComponentID
filterWith DynamicQueryFilter
dynF)
          Bool -> Bool -> Bool
&& (ComponentID -> Bool) -> Set ComponentID -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
F.all (\ComponentID
cId -> Bool -> Bool
not (ComponentID -> Archetype m -> Bool
forall (m :: * -> *). ComponentID -> Archetype m -> Bool
A.member ComponentID
cId (Archetype m -> Bool) -> Archetype m -> Bool
forall a b. (a -> b) -> a -> b
$ Node m -> Archetype m
forall (m :: * -> *). Node m -> Archetype m
nodeArchetype Node m
n)) (DynamicQueryFilter -> Set ComponentID
filterWithout DynamicQueryFilter
dynF)
   in (Components
cs'', Set ComponentID
-> DynamicQuery m a
-> (Node m -> Bool)
-> DynamicSystem m (Vector a)
forall (m :: * -> *) a.
Set ComponentID
-> DynamicQuery m a
-> (Node m -> Bool)
-> DynamicSystem m (Vector a)
DS.runQueryFiltered (ReadsWrites -> Set ComponentID
Q.reads ReadsWrites
rws Set ComponentID -> Set ComponentID -> Set ComponentID
forall a. Semigroup a => a -> a -> a
<> ReadsWrites -> Set ComponentID
Q.writes ReadsWrites
rws) DynamicQuery m a
dynQ Node m -> Bool
flt)

-- | Match all entities with a dynamic query.
readQueryDyn :: Set ComponentID -> DynamicQuery m a -> System m (Vector a)
readQueryDyn :: forall (m :: * -> *) a.
Set ComponentID -> DynamicQuery m a -> System m (Vector a)
readQueryDyn Set ComponentID
cIds DynamicQuery m a
q = (Components -> (Components, DynamicSystem m (Vector a)))
-> System m (Vector a)
forall (m :: * -> *) a.
(Components -> (Components, DynamicSystem m a)) -> System m a
System (,Set ComponentID -> DynamicQuery m a -> DynamicSystem m (Vector a)
forall (m :: * -> *) a.
Set ComponentID -> DynamicQuery m a -> DynamicSystem m (Vector a)
DS.readQuery Set ComponentID
cIds DynamicQuery m a
q)

readQuerySingleMaybeDyn :: Set ComponentID -> DynamicQuery m a -> System m (Maybe a)
readQuerySingleMaybeDyn :: forall (m :: * -> *) a.
Set ComponentID -> DynamicQuery m a -> System m (Maybe a)
readQuerySingleMaybeDyn Set ComponentID
cIds DynamicQuery m a
q = (Components -> (Components, DynamicSystem m (Maybe a)))
-> System m (Maybe a)
forall (m :: * -> *) a.
(Components -> (Components, DynamicSystem m a)) -> System m a
System (,Set ComponentID -> DynamicQuery m a -> DynamicSystem m (Maybe a)
forall (m :: * -> *) a.
Set ComponentID -> DynamicQuery m a -> DynamicSystem m (Maybe a)
DS.readQuerySingleMaybe Set ComponentID
cIds DynamicQuery m a
q)

-- | Match all entities with a dynamic query and filter.
readQueryFilteredDyn :: Set ComponentID -> DynamicQuery m a -> (Node m -> Bool) -> System m (Vector a)
readQueryFilteredDyn :: forall (m :: * -> *) a.
Set ComponentID
-> DynamicQuery m a -> (Node m -> Bool) -> System m (Vector a)
readQueryFilteredDyn Set ComponentID
cIds DynamicQuery m a
q Node m -> Bool
f = (Components -> (Components, DynamicSystem m (Vector a)))
-> System m (Vector a)
forall (m :: * -> *) a.
(Components -> (Components, DynamicSystem m a)) -> System m a
System (,Set ComponentID
-> (Node m -> Bool)
-> DynamicQuery m a
-> DynamicSystem m (Vector a)
forall (m :: * -> *) a.
Set ComponentID
-> (Node m -> Bool)
-> DynamicQuery m a
-> DynamicSystem m (Vector a)
DS.readQueryFiltered Set ComponentID
cIds Node m -> Bool
f DynamicQuery m a
q)

-- | Map all entities with a dynamic query.
runQueryDyn :: Set ComponentID -> DynamicQuery m a -> System m (Vector a)
runQueryDyn :: forall (m :: * -> *) a.
Set ComponentID -> DynamicQuery m a -> System m (Vector a)
runQueryDyn Set ComponentID
cIds DynamicQuery m a
q = (Components -> (Components, DynamicSystem m (Vector a)))
-> System m (Vector a)
forall (m :: * -> *) a.
(Components -> (Components, DynamicSystem m a)) -> System m a
System (,Set ComponentID -> DynamicQuery m a -> DynamicSystem m (Vector a)
forall (m :: * -> *) a.
Set ComponentID -> DynamicQuery m a -> DynamicSystem m (Vector a)
DS.runQuery Set ComponentID
cIds DynamicQuery m a
q)

-- | Map a single entity with a dynamic query.
runQuerySingleMaybeDyn :: Set ComponentID -> DynamicQuery m a -> System m (Maybe a)
runQuerySingleMaybeDyn :: forall (m :: * -> *) a.
Set ComponentID -> DynamicQuery m a -> System m (Maybe a)
runQuerySingleMaybeDyn Set ComponentID
cIds DynamicQuery m a
q = (Components -> (Components, DynamicSystem m (Maybe a)))
-> System m (Maybe a)
forall (m :: * -> *) a.
(Components -> (Components, DynamicSystem m a)) -> System m a
System (,Set ComponentID -> DynamicQuery m a -> DynamicSystem m (Maybe a)
forall (m :: * -> *) a.
Set ComponentID -> DynamicQuery m a -> DynamicSystem m (Maybe a)
DS.runQuerySingleMaybe Set ComponentID
cIds DynamicQuery m a
q)

-- | Filter and map all entities with a dynamic query.
runQueryFilteredDyn :: Set ComponentID -> (Node m -> Bool) -> DynamicQuery m a -> System m (Vector a)
runQueryFilteredDyn :: forall (m :: * -> *) a.
Set ComponentID
-> (Node m -> Bool) -> DynamicQuery m a -> System m (Vector a)
runQueryFilteredDyn Set ComponentID
cIds Node m -> Bool
f DynamicQuery m a
q = (Components -> (Components, DynamicSystem m (Vector a)))
-> System m (Vector a)
forall (m :: * -> *) a.
(Components -> (Components, DynamicSystem m a)) -> System m a
System (,Set ComponentID
-> DynamicQuery m a
-> (Node m -> Bool)
-> DynamicSystem m (Vector a)
forall (m :: * -> *) a.
Set ComponentID
-> DynamicQuery m a
-> (Node m -> Bool)
-> DynamicSystem m (Vector a)
DS.runQueryFiltered Set ComponentID
cIds DynamicQuery m a
q Node m -> Bool
f)