{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE TypeFamilies #-}

module Aztecs.ECS.Query.Dynamic
  ( -- * Dynamic queries
    DynamicQuery (..),
    ArrowDynamicQueryReader (..),
    ArrowDynamicQuery (..),

    -- ** Conversion
    fromDynReader,
    toDynReader,

    -- ** Running
    mapDyn,

    -- * Dynamic query filters
    DynamicQueryFilter (..),
  )
where

import Aztecs.ECS.Component (ComponentID)
import Aztecs.ECS.Entity (EntityID)
import Aztecs.ECS.Query.Dynamic.Class (ArrowDynamicQuery (..))
import Aztecs.ECS.Query.Dynamic.Reader (DynamicQueryFilter (..), DynamicQueryReader (..))
import Aztecs.ECS.Query.Dynamic.Reader.Class (ArrowDynamicQueryReader (..))
import Aztecs.ECS.World.Archetype (Archetype)
import qualified Aztecs.ECS.World.Archetype as A
import Aztecs.ECS.World.Archetypes (Node (..))
import qualified Aztecs.ECS.World.Archetypes as AS
import Aztecs.ECS.World.Entities (Entities (..))
import Control.Arrow (Arrow (..), ArrowChoice (..))
import Control.Category (Category (..))
import Data.Either (partitionEithers)
import Data.Foldable
import qualified Data.Map as Map
import Data.Set (Set)
import qualified Data.Set as Set
import Prelude hiding ((.))

-- | Dynamic query for components by ID.
newtype DynamicQuery i o
  = DynamicQuery {forall i o.
DynamicQuery i o
-> [i] -> [EntityID] -> Archetype -> ([o], Archetype)
runDynQuery :: [i] -> [EntityID] -> Archetype -> ([o], Archetype)}
  deriving ((forall a b. (a -> b) -> DynamicQuery i a -> DynamicQuery i b)
-> (forall a b. a -> DynamicQuery i b -> DynamicQuery i a)
-> Functor (DynamicQuery i)
forall a b. a -> DynamicQuery i b -> DynamicQuery i a
forall a b. (a -> b) -> DynamicQuery i a -> DynamicQuery i b
forall i a b. a -> DynamicQuery i b -> DynamicQuery i a
forall i a b. (a -> b) -> DynamicQuery i a -> DynamicQuery i b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall i a b. (a -> b) -> DynamicQuery i a -> DynamicQuery i b
fmap :: forall a b. (a -> b) -> DynamicQuery i a -> DynamicQuery i b
$c<$ :: forall i a b. a -> DynamicQuery i b -> DynamicQuery i a
<$ :: forall a b. a -> DynamicQuery i b -> DynamicQuery i a
Functor)

instance Applicative (DynamicQuery i) where
  pure :: forall a. a -> DynamicQuery i a
pure a
a = ([i] -> [EntityID] -> Archetype -> ([a], Archetype))
-> DynamicQuery i a
forall i o.
([i] -> [EntityID] -> Archetype -> ([o], Archetype))
-> DynamicQuery i o
DynamicQuery (([i] -> [EntityID] -> Archetype -> ([a], Archetype))
 -> DynamicQuery i a)
-> ([i] -> [EntityID] -> Archetype -> ([a], Archetype))
-> DynamicQuery i a
forall a b. (a -> b) -> a -> b
$ \[i]
_ [EntityID]
es Archetype
arch -> (Int -> a -> [a]
forall a. Int -> a -> [a]
replicate ([EntityID] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [EntityID]
es) a
a, Archetype
arch)

  DynamicQuery i (a -> b)
f <*> :: forall a b.
DynamicQuery i (a -> b) -> DynamicQuery i a -> DynamicQuery i b
<*> DynamicQuery i a
g = ([i] -> [EntityID] -> Archetype -> ([b], Archetype))
-> DynamicQuery i b
forall i o.
([i] -> [EntityID] -> Archetype -> ([o], Archetype))
-> DynamicQuery i o
DynamicQuery (([i] -> [EntityID] -> Archetype -> ([b], Archetype))
 -> DynamicQuery i b)
-> ([i] -> [EntityID] -> Archetype -> ([b], Archetype))
-> DynamicQuery i b
forall a b. (a -> b) -> a -> b
$ \[i]
i [EntityID]
es Archetype
arch ->
    let ([a]
as, Archetype
arch') = DynamicQuery i a
-> [i] -> [EntityID] -> Archetype -> ([a], Archetype)
forall i o.
DynamicQuery i o
-> [i] -> [EntityID] -> Archetype -> ([o], Archetype)
runDynQuery DynamicQuery i a
g [i]
i [EntityID]
es Archetype
arch
        ([a -> b]
fs, Archetype
arch'') = DynamicQuery i (a -> b)
-> [i] -> [EntityID] -> Archetype -> ([a -> b], Archetype)
forall i o.
DynamicQuery i o
-> [i] -> [EntityID] -> Archetype -> ([o], Archetype)
runDynQuery DynamicQuery i (a -> b)
f [i]
i [EntityID]
es Archetype
arch'
     in (((a -> b) -> a -> b) -> [a -> b] -> [a] -> [b]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (a -> b) -> a -> b
forall a b. (a -> b) -> a -> b
($) [a -> b]
fs [a]
as, Archetype
arch'')

instance Category DynamicQuery where
  id :: forall a. DynamicQuery a a
id = ([a] -> [EntityID] -> Archetype -> ([a], Archetype))
-> DynamicQuery a a
forall i o.
([i] -> [EntityID] -> Archetype -> ([o], Archetype))
-> DynamicQuery i o
DynamicQuery (([a] -> [EntityID] -> Archetype -> ([a], Archetype))
 -> DynamicQuery a a)
-> ([a] -> [EntityID] -> Archetype -> ([a], Archetype))
-> DynamicQuery a a
forall a b. (a -> b) -> a -> b
$ \[a]
as [EntityID]
_ Archetype
arch -> ([a]
as, Archetype
arch)

  DynamicQuery b c
f . :: forall b c a.
DynamicQuery b c -> DynamicQuery a b -> DynamicQuery a c
. DynamicQuery a b
g = ([a] -> [EntityID] -> Archetype -> ([c], Archetype))
-> DynamicQuery a c
forall i o.
([i] -> [EntityID] -> Archetype -> ([o], Archetype))
-> DynamicQuery i o
DynamicQuery (([a] -> [EntityID] -> Archetype -> ([c], Archetype))
 -> DynamicQuery a c)
-> ([a] -> [EntityID] -> Archetype -> ([c], Archetype))
-> DynamicQuery a c
forall a b. (a -> b) -> a -> b
$ \[a]
i [EntityID]
es Archetype
arch ->
    let ([b]
as, Archetype
arch') = DynamicQuery a b
-> [a] -> [EntityID] -> Archetype -> ([b], Archetype)
forall i o.
DynamicQuery i o
-> [i] -> [EntityID] -> Archetype -> ([o], Archetype)
runDynQuery DynamicQuery a b
g [a]
i [EntityID]
es Archetype
arch in DynamicQuery b c
-> [b] -> [EntityID] -> Archetype -> ([c], Archetype)
forall i o.
DynamicQuery i o
-> [i] -> [EntityID] -> Archetype -> ([o], Archetype)
runDynQuery DynamicQuery b c
f [b]
as [EntityID]
es Archetype
arch'

instance Arrow DynamicQuery where
  arr :: forall b c. (b -> c) -> DynamicQuery b c
arr b -> c
f = ([b] -> [EntityID] -> Archetype -> ([c], Archetype))
-> DynamicQuery b c
forall i o.
([i] -> [EntityID] -> Archetype -> ([o], Archetype))
-> DynamicQuery i o
DynamicQuery (([b] -> [EntityID] -> Archetype -> ([c], Archetype))
 -> DynamicQuery b c)
-> ([b] -> [EntityID] -> Archetype -> ([c], Archetype))
-> DynamicQuery b c
forall a b. (a -> b) -> a -> b
$ \[b]
bs [EntityID]
_ Archetype
arch -> ((b -> c) -> [b] -> [c]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> c
f [b]
bs, Archetype
arch)
  first :: forall b c d. DynamicQuery b c -> DynamicQuery (b, d) (c, d)
first DynamicQuery b c
f = ([(b, d)] -> [EntityID] -> Archetype -> ([(c, d)], Archetype))
-> DynamicQuery (b, d) (c, d)
forall i o.
([i] -> [EntityID] -> Archetype -> ([o], Archetype))
-> DynamicQuery i o
DynamicQuery (([(b, d)] -> [EntityID] -> Archetype -> ([(c, d)], Archetype))
 -> DynamicQuery (b, d) (c, d))
-> ([(b, d)] -> [EntityID] -> Archetype -> ([(c, d)], Archetype))
-> DynamicQuery (b, d) (c, d)
forall a b. (a -> b) -> a -> b
$ \[(b, d)]
bds [EntityID]
es Archetype
arch ->
    let ([b]
bs, [d]
ds) = [(b, d)] -> ([b], [d])
forall a b. [(a, b)] -> ([a], [b])
unzip [(b, d)]
bds
        ([c]
cs, Archetype
arch') = DynamicQuery b c
-> [b] -> [EntityID] -> Archetype -> ([c], Archetype)
forall i o.
DynamicQuery i o
-> [i] -> [EntityID] -> Archetype -> ([o], Archetype)
runDynQuery DynamicQuery b c
f [b]
bs [EntityID]
es Archetype
arch
     in ([c] -> [d] -> [(c, d)]
forall a b. [a] -> [b] -> [(a, b)]
zip [c]
cs [d]
ds, Archetype
arch')

instance ArrowChoice DynamicQuery where
  left :: forall b c d.
DynamicQuery b c -> DynamicQuery (Either b d) (Either c d)
left DynamicQuery b c
f = ([Either b d]
 -> [EntityID] -> Archetype -> ([Either c d], Archetype))
-> DynamicQuery (Either b d) (Either c d)
forall i o.
([i] -> [EntityID] -> Archetype -> ([o], Archetype))
-> DynamicQuery i o
DynamicQuery (([Either b d]
  -> [EntityID] -> Archetype -> ([Either c d], Archetype))
 -> DynamicQuery (Either b d) (Either c d))
-> ([Either b d]
    -> [EntityID] -> Archetype -> ([Either c d], Archetype))
-> DynamicQuery (Either b d) (Either c d)
forall a b. (a -> b) -> a -> b
$ \[Either b d]
eds [EntityID]
es Archetype
arch ->
    let ([b]
es', [d]
ds) = [Either b d] -> ([b], [d])
forall a b. [Either a b] -> ([a], [b])
partitionEithers [Either b d]
eds
        ([c]
cs, Archetype
arch') = DynamicQuery b c
-> [b] -> [EntityID] -> Archetype -> ([c], Archetype)
forall i o.
DynamicQuery i o
-> [i] -> [EntityID] -> Archetype -> ([o], Archetype)
runDynQuery DynamicQuery b c
f [b]
es' [EntityID]
es Archetype
arch
     in ((c -> Either c d) -> [c] -> [Either c d]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap c -> Either c d
forall a b. a -> Either a b
Left [c]
cs [Either c d] -> [Either c d] -> [Either c d]
forall a. [a] -> [a] -> [a]
++ (d -> Either c d) -> [d] -> [Either c d]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap d -> Either c d
forall a b. b -> Either a b
Right [d]
ds, Archetype
arch')

instance ArrowDynamicQueryReader DynamicQuery where
  entity :: DynamicQuery () EntityID
entity = DynamicQueryReader () EntityID -> DynamicQuery () EntityID
forall i o. DynamicQueryReader i o -> DynamicQuery i o
fromDynReader DynamicQueryReader () EntityID
forall (arr :: * -> * -> *).
ArrowDynamicQueryReader arr =>
arr () EntityID
entity
  fetchDyn :: forall a. Component a => ComponentID -> DynamicQuery () a
fetchDyn = DynamicQueryReader () a -> DynamicQuery () a
forall i o. DynamicQueryReader i o -> DynamicQuery i o
fromDynReader (DynamicQueryReader () a -> DynamicQuery () a)
-> (ComponentID -> DynamicQueryReader () a)
-> ComponentID
-> DynamicQuery () a
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. ComponentID -> DynamicQueryReader () a
forall a. Component a => ComponentID -> DynamicQueryReader () a
forall (arr :: * -> * -> *) a.
(ArrowDynamicQueryReader arr, Component a) =>
ComponentID -> arr () a
fetchDyn
  fetchMaybeDyn :: forall a. Component a => ComponentID -> DynamicQuery () (Maybe a)
fetchMaybeDyn = DynamicQueryReader () (Maybe a) -> DynamicQuery () (Maybe a)
forall i o. DynamicQueryReader i o -> DynamicQuery i o
fromDynReader (DynamicQueryReader () (Maybe a) -> DynamicQuery () (Maybe a))
-> (ComponentID -> DynamicQueryReader () (Maybe a))
-> ComponentID
-> DynamicQuery () (Maybe a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. ComponentID -> DynamicQueryReader () (Maybe a)
forall a.
Component a =>
ComponentID -> DynamicQueryReader () (Maybe a)
forall (arr :: * -> * -> *) a.
(ArrowDynamicQueryReader arr, Component a) =>
ComponentID -> arr () (Maybe a)
fetchMaybeDyn

instance ArrowDynamicQuery DynamicQuery where
  setDyn :: forall a. Component a => ComponentID -> DynamicQuery a a
setDyn ComponentID
cId = ([a] -> [EntityID] -> Archetype -> ([a], Archetype))
-> DynamicQuery a a
forall i o.
([i] -> [EntityID] -> Archetype -> ([o], Archetype))
-> DynamicQuery i o
DynamicQuery (([a] -> [EntityID] -> Archetype -> ([a], Archetype))
 -> DynamicQuery a a)
-> ([a] -> [EntityID] -> Archetype -> ([a], Archetype))
-> DynamicQuery a a
forall a b. (a -> b) -> a -> b
$ \[a]
is [EntityID]
_ Archetype
arch ->
    let !arch' :: Archetype
arch' = ComponentID -> [a] -> Archetype -> Archetype
forall a.
Component a =>
ComponentID -> [a] -> Archetype -> Archetype
A.insertAscList ComponentID
cId [a]
is Archetype
arch in ([a]
is, Archetype
arch')

fromDynReader :: DynamicQueryReader i o -> DynamicQuery i o
fromDynReader :: forall i o. DynamicQueryReader i o -> DynamicQuery i o
fromDynReader DynamicQueryReader i o
q = ([i] -> [EntityID] -> Archetype -> ([o], Archetype))
-> DynamicQuery i o
forall i o.
([i] -> [EntityID] -> Archetype -> ([o], Archetype))
-> DynamicQuery i o
DynamicQuery (([i] -> [EntityID] -> Archetype -> ([o], Archetype))
 -> DynamicQuery i o)
-> ([i] -> [EntityID] -> Archetype -> ([o], Archetype))
-> DynamicQuery i o
forall a b. (a -> b) -> a -> b
$ \[i]
is [EntityID]
es Archetype
arch ->
  let os :: [o]
os = DynamicQueryReader i o -> [i] -> [EntityID] -> Archetype -> [o]
forall i o.
DynamicQueryReader i o -> [i] -> [EntityID] -> Archetype -> [o]
runDynQueryReader' DynamicQueryReader i o
q [i]
is [EntityID]
es Archetype
arch in ([o]
os, Archetype
arch)

toDynReader :: DynamicQuery i o -> DynamicQueryReader i o
toDynReader :: forall i o. DynamicQuery i o -> DynamicQueryReader i o
toDynReader DynamicQuery i o
q = ([i] -> [EntityID] -> Archetype -> [o]) -> DynamicQueryReader i o
forall i o.
([i] -> [EntityID] -> Archetype -> [o]) -> DynamicQueryReader i o
DynamicQueryReader (([i] -> [EntityID] -> Archetype -> [o]) -> DynamicQueryReader i o)
-> ([i] -> [EntityID] -> Archetype -> [o])
-> DynamicQueryReader i o
forall a b. (a -> b) -> a -> b
$ \[i]
is [EntityID]
es Archetype
arch -> ([o], Archetype) -> [o]
forall a b. (a, b) -> a
fst (([o], Archetype) -> [o]) -> ([o], Archetype) -> [o]
forall a b. (a -> b) -> a -> b
$ DynamicQuery i o
-> [i] -> [EntityID] -> Archetype -> ([o], Archetype)
forall i o.
DynamicQuery i o
-> [i] -> [EntityID] -> Archetype -> ([o], Archetype)
runDynQuery DynamicQuery i o
q [i]
is [EntityID]
es Archetype
arch

-- | Map all matched entities.
mapDyn :: Set ComponentID -> i -> DynamicQuery i a -> Entities -> ([a], Entities)
mapDyn :: forall i a.
Set ComponentID
-> i -> DynamicQuery i a -> Entities -> ([a], Entities)
mapDyn Set ComponentID
cIds i
i DynamicQuery i a
q Entities
es =
  let ([a]
as, Entities
es') =
        if Set ComponentID -> Bool
forall a. Set a -> Bool
Set.null Set ComponentID
cIds
          then (([a], Archetype) -> [a]
forall a b. (a, b) -> a
fst (([a], Archetype) -> [a]) -> ([a], Archetype) -> [a]
forall a b. (a -> b) -> a -> b
$ [EntityID] -> Archetype -> ([a], Archetype)
go (Map EntityID ArchetypeID -> [EntityID]
forall k a. Map k a -> [k]
Map.keys (Map EntityID ArchetypeID -> [EntityID])
-> Map EntityID ArchetypeID -> [EntityID]
forall a b. (a -> b) -> a -> b
$ Entities -> Map EntityID ArchetypeID
entities Entities
es) Archetype
A.empty, Entities
es)
          else
            (([a], Entities) -> (ArchetypeID, Node) -> ([a], Entities))
-> ([a], Entities) -> [(ArchetypeID, Node)] -> ([a], Entities)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
              ( \([a]
acc, Entities
esAcc) (ArchetypeID
aId, Node
n) ->
                  let ([a]
as', Archetype
arch') = [EntityID] -> Archetype -> ([a], Archetype)
go (Set EntityID -> [EntityID]
forall a. Set a -> [a]
Set.toList (Set EntityID -> [EntityID])
-> (Archetype -> Set EntityID) -> Archetype -> [EntityID]
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Archetype -> Set EntityID
A.entities (Archetype -> [EntityID]) -> Archetype -> [EntityID]
forall a b. (a -> b) -> a -> b
$ Node -> Archetype
nodeArchetype Node
n) (Node -> Archetype
nodeArchetype Node
n)
                      nodes :: Map ArchetypeID Node
nodes = ArchetypeID -> Node -> Map ArchetypeID Node -> Map ArchetypeID Node
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert ArchetypeID
aId Node
n {nodeArchetype = arch'} (Archetypes -> Map ArchetypeID Node
AS.nodes (Archetypes -> Map ArchetypeID Node)
-> Archetypes -> Map ArchetypeID Node
forall a b. (a -> b) -> a -> b
$ Entities -> Archetypes
archetypes Entities
esAcc)
                   in ([a]
as' [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
acc, Entities
esAcc {archetypes = (archetypes esAcc) {AS.nodes = nodes}})
              )
              ([], Entities
es)
              (Map ArchetypeID Node -> [(ArchetypeID, Node)]
forall k a. Map k a -> [(k, a)]
Map.toList (Map ArchetypeID Node -> [(ArchetypeID, Node)])
-> Map ArchetypeID Node -> [(ArchetypeID, Node)]
forall a b. (a -> b) -> a -> b
$ Set ComponentID -> Archetypes -> Map ArchetypeID Node
AS.find Set ComponentID
cIds (Entities -> Archetypes
archetypes Entities
es))
      go :: [EntityID] -> Archetype -> ([a], Archetype)
go = DynamicQuery i a
-> [i] -> [EntityID] -> Archetype -> ([a], Archetype)
forall i o.
DynamicQuery i o
-> [i] -> [EntityID] -> Archetype -> ([o], Archetype)
runDynQuery DynamicQuery i a
q (i -> [i]
forall a. a -> [a]
repeat i
i)
   in ([a]
as, Entities
es')