{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}

module Aztecs.ECS.Query
  ( -- * Queries
    Query (..),
    ArrowQueryReader (..),
    ArrowQuery (..),
    ArrowDynamicQueryReader (..),
    ArrowDynamicQuery (..),

    -- ** Running
    all,
    map,

    -- ** Conversion
    fromReader,
    toReader,

    -- * Filters
    QueryFilter (..),
    with,
    without,

    -- * Reads and writes
    ReadsWrites (..),
    disjoint,
  )
where

import Aztecs.ECS.Component
import Aztecs.ECS.Query.Class (ArrowQuery (..))
import Aztecs.ECS.Query.Dynamic (DynamicQuery (..), fromDynReader, mapDyn, toDynReader)
import Aztecs.ECS.Query.Dynamic.Class (ArrowDynamicQuery (..))
import Aztecs.ECS.Query.Dynamic.Reader.Class (ArrowDynamicQueryReader (..))
import Aztecs.ECS.Query.Reader (QueryFilter (..), QueryReader (..), with, without)
import qualified Aztecs.ECS.Query.Reader as QR
import Aztecs.ECS.Query.Reader.Class (ArrowQueryReader (..))
import Aztecs.ECS.World.Components (Components)
import qualified Aztecs.ECS.World.Components as CS
import Aztecs.ECS.World.Entities (Entities (..))
import Control.Arrow (Arrow (..), ArrowChoice (..))
import Control.Category (Category (..))
import Data.Set (Set)
import qualified Data.Set as Set
import Prelude hiding (all, id, map, reads, (.))

-- | Query for matching entities.
--
-- === Do notation:
-- > move :: (ArrowQuery arr) => arr () Position
-- > move = proc () -> do
-- >   Velocity v <- Q.fetch -< ()
-- >   Position p <- Q.fetch -< ()
-- >   Q.set -< Position $ p + v
--
-- === Arrow combinators:
-- > move :: (ArrowQuery arr) => arr () Position
-- > move = Q.fetch &&& Q.fetch >>> arr (\(Position p, Velocity v) -> Position $ p + v) >>> Q.set
--
-- === Applicative combinators:
-- > move :: (ArrowQuery arr) => arr () Position
-- > move = (,) <$> Q.fetch <*> Q.fetch >>> arr (\(Position p, Velocity v) -> Position $ p + v) >>> Q.set
newtype Query i o = Query {forall i o.
Query i o
-> Components -> (ReadsWrites, Components, DynamicQuery i o)
runQuery :: Components -> (ReadsWrites, Components, DynamicQuery i o)}
  deriving ((forall a b. (a -> b) -> Query i a -> Query i b)
-> (forall a b. a -> Query i b -> Query i a) -> Functor (Query i)
forall a b. a -> Query i b -> Query i a
forall a b. (a -> b) -> Query i a -> Query i b
forall i a b. a -> Query i b -> Query i a
forall i a b. (a -> b) -> Query i a -> Query 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) -> Query i a -> Query i b
fmap :: forall a b. (a -> b) -> Query i a -> Query i b
$c<$ :: forall i a b. a -> Query i b -> Query i a
<$ :: forall a b. a -> Query i b -> Query i a
Functor)

instance Applicative (Query i) where
  pure :: forall a. a -> Query i a
pure a
a = (Components -> (ReadsWrites, Components, DynamicQuery i a))
-> Query i a
forall i o.
(Components -> (ReadsWrites, Components, DynamicQuery i o))
-> Query i o
Query (ReadsWrites
forall a. Monoid a => a
mempty,,a -> DynamicQuery i a
forall a. a -> DynamicQuery i a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a)
  (Query Components -> (ReadsWrites, Components, DynamicQuery i (a -> b))
f) <*> :: forall a b. Query i (a -> b) -> Query i a -> Query i b
<*> (Query Components -> (ReadsWrites, Components, DynamicQuery i a)
g) = (Components -> (ReadsWrites, Components, DynamicQuery i b))
-> Query i b
forall i o.
(Components -> (ReadsWrites, Components, DynamicQuery i o))
-> Query i o
Query ((Components -> (ReadsWrites, Components, DynamicQuery i b))
 -> Query i b)
-> (Components -> (ReadsWrites, Components, DynamicQuery i b))
-> Query i b
forall a b. (a -> b) -> a -> b
$ \Components
cs ->
    let (ReadsWrites
cIdsG, Components
cs', DynamicQuery i a
aQS) = Components -> (ReadsWrites, Components, DynamicQuery i a)
g Components
cs
        (ReadsWrites
cIdsF, Components
cs'', DynamicQuery i (a -> b)
bQS) = Components -> (ReadsWrites, Components, DynamicQuery i (a -> b))
f Components
cs'
     in (ReadsWrites
cIdsG ReadsWrites -> ReadsWrites -> ReadsWrites
forall a. Semigroup a => a -> a -> a
<> ReadsWrites
cIdsF, Components
cs'', DynamicQuery i (a -> b)
bQS DynamicQuery i (a -> b) -> DynamicQuery i a -> DynamicQuery i b
forall a b.
DynamicQuery i (a -> b) -> DynamicQuery i a -> DynamicQuery i b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> DynamicQuery i a
aQS)

instance Category Query where
  id :: forall a. Query a a
id = (Components -> (ReadsWrites, Components, DynamicQuery a a))
-> Query a a
forall i o.
(Components -> (ReadsWrites, Components, DynamicQuery i o))
-> Query i o
Query (ReadsWrites
forall a. Monoid a => a
mempty,,DynamicQuery a a
forall a. DynamicQuery a a
forall {k} (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id)
  (Query Components -> (ReadsWrites, Components, DynamicQuery b c)
f) . :: forall b c a. Query b c -> Query a b -> Query a c
. (Query Components -> (ReadsWrites, Components, DynamicQuery a b)
g) = (Components -> (ReadsWrites, Components, DynamicQuery a c))
-> Query a c
forall i o.
(Components -> (ReadsWrites, Components, DynamicQuery i o))
-> Query i o
Query ((Components -> (ReadsWrites, Components, DynamicQuery a c))
 -> Query a c)
-> (Components -> (ReadsWrites, Components, DynamicQuery a c))
-> Query a c
forall a b. (a -> b) -> a -> b
$ \Components
cs ->
    let (ReadsWrites
cIdsG, Components
cs', DynamicQuery a b
aQS) = Components -> (ReadsWrites, Components, DynamicQuery a b)
g Components
cs
        (ReadsWrites
cIdsF, Components
cs'', DynamicQuery b c
bQS) = Components -> (ReadsWrites, Components, DynamicQuery b c)
f Components
cs'
     in (ReadsWrites
cIdsG ReadsWrites -> ReadsWrites -> ReadsWrites
forall a. Semigroup a => a -> a -> a
<> ReadsWrites
cIdsF, Components
cs'', DynamicQuery b c
bQS DynamicQuery b c -> DynamicQuery a b -> DynamicQuery a c
forall b c a.
DynamicQuery b c -> DynamicQuery a b -> DynamicQuery a c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. DynamicQuery a b
aQS)

instance Arrow Query where
  arr :: forall b c. (b -> c) -> Query b c
arr b -> c
f = (Components -> (ReadsWrites, Components, DynamicQuery b c))
-> Query b c
forall i o.
(Components -> (ReadsWrites, Components, DynamicQuery i o))
-> Query i o
Query (ReadsWrites
forall a. Monoid a => a
mempty,,(b -> c) -> DynamicQuery b c
forall b c. (b -> c) -> DynamicQuery b c
forall (a :: * -> * -> *) b c. Arrow a => (b -> c) -> a b c
arr b -> c
f)
  first :: forall b c d. Query b c -> Query (b, d) (c, d)
first (Query Components -> (ReadsWrites, Components, DynamicQuery b c)
f) = (Components
 -> (ReadsWrites, Components, DynamicQuery (b, d) (c, d)))
-> Query (b, d) (c, d)
forall i o.
(Components -> (ReadsWrites, Components, DynamicQuery i o))
-> Query i o
Query ((Components
  -> (ReadsWrites, Components, DynamicQuery (b, d) (c, d)))
 -> Query (b, d) (c, d))
-> (Components
    -> (ReadsWrites, Components, DynamicQuery (b, d) (c, d)))
-> Query (b, d) (c, d)
forall a b. (a -> b) -> a -> b
$ \Components
comps -> let (ReadsWrites
cIds, Components
comps', DynamicQuery b c
qS) = Components -> (ReadsWrites, Components, DynamicQuery b c)
f Components
comps in (ReadsWrites
cIds, Components
comps', DynamicQuery b c -> DynamicQuery (b, d) (c, d)
forall b c d. DynamicQuery b c -> DynamicQuery (b, d) (c, d)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first DynamicQuery b c
qS)

instance ArrowChoice Query where
  left :: forall b c d. Query b c -> Query (Either b d) (Either c d)
left (Query Components -> (ReadsWrites, Components, DynamicQuery b c)
f) = (Components
 -> (ReadsWrites, Components,
     DynamicQuery (Either b d) (Either c d)))
-> Query (Either b d) (Either c d)
forall i o.
(Components -> (ReadsWrites, Components, DynamicQuery i o))
-> Query i o
Query ((Components
  -> (ReadsWrites, Components,
      DynamicQuery (Either b d) (Either c d)))
 -> Query (Either b d) (Either c d))
-> (Components
    -> (ReadsWrites, Components,
        DynamicQuery (Either b d) (Either c d)))
-> Query (Either b d) (Either c d)
forall a b. (a -> b) -> a -> b
$ \Components
comps -> let (ReadsWrites
cIds, Components
comps', DynamicQuery b c
qS) = Components -> (ReadsWrites, Components, DynamicQuery b c)
f Components
comps in (ReadsWrites
cIds, Components
comps', DynamicQuery b c -> DynamicQuery (Either b d) (Either c d)
forall b c d.
DynamicQuery b c -> DynamicQuery (Either b d) (Either c d)
forall (a :: * -> * -> *) b c d.
ArrowChoice a =>
a b c -> a (Either b d) (Either c d)
left DynamicQuery b c
qS)

instance ArrowQueryReader Query where
  fetch :: forall a. Component a => Query () a
fetch = QueryReader () a -> Query () a
forall i o. QueryReader i o -> Query i o
fromReader QueryReader () a
forall a. Component a => QueryReader () a
forall (arr :: * -> * -> *) a.
(ArrowQueryReader arr, Component a) =>
arr () a
fetch
  fetchMaybe :: forall a. Component a => Query () (Maybe a)
fetchMaybe = QueryReader () (Maybe a) -> Query () (Maybe a)
forall i o. QueryReader i o -> Query i o
fromReader QueryReader () (Maybe a)
forall a. Component a => QueryReader () (Maybe a)
forall (arr :: * -> * -> *) a.
(ArrowQueryReader arr, Component a) =>
arr () (Maybe a)
fetchMaybe

instance ArrowDynamicQueryReader Query where
  entity :: Query () EntityID
entity = QueryReader () EntityID -> Query () EntityID
forall i o. QueryReader i o -> Query i o
fromReader QueryReader () EntityID
forall (arr :: * -> * -> *).
ArrowDynamicQueryReader arr =>
arr () EntityID
entity
  fetchDyn :: forall a. Component a => ComponentID -> Query () a
fetchDyn = QueryReader () a -> Query () a
forall i o. QueryReader i o -> Query i o
fromReader (QueryReader () a -> Query () a)
-> (ComponentID -> QueryReader () a) -> ComponentID -> Query () 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 -> QueryReader () a
forall a. Component a => ComponentID -> QueryReader () a
forall (arr :: * -> * -> *) a.
(ArrowDynamicQueryReader arr, Component a) =>
ComponentID -> arr () a
fetchDyn
  fetchMaybeDyn :: forall a. Component a => ComponentID -> Query () (Maybe a)
fetchMaybeDyn = QueryReader () (Maybe a) -> Query () (Maybe a)
forall i o. QueryReader i o -> Query i o
fromReader (QueryReader () (Maybe a) -> Query () (Maybe a))
-> (ComponentID -> QueryReader () (Maybe a))
-> ComponentID
-> Query () (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 -> QueryReader () (Maybe a)
forall a. Component a => ComponentID -> QueryReader () (Maybe a)
forall (arr :: * -> * -> *) a.
(ArrowDynamicQueryReader arr, Component a) =>
ComponentID -> arr () (Maybe a)
fetchMaybeDyn

instance ArrowDynamicQuery Query where
  setDyn :: forall a. Component a => ComponentID -> Query a a
setDyn ComponentID
cId = (Components -> (ReadsWrites, Components, DynamicQuery a a))
-> Query a a
forall i o.
(Components -> (ReadsWrites, Components, DynamicQuery i o))
-> Query i o
Query (Set ComponentID -> Set ComponentID -> ReadsWrites
ReadsWrites Set ComponentID
forall a. Set a
Set.empty (ComponentID -> Set ComponentID
forall a. a -> Set a
Set.singleton ComponentID
cId),,ComponentID -> DynamicQuery a a
forall a. Component a => ComponentID -> DynamicQuery a a
forall (arr :: * -> * -> *) a.
(ArrowDynamicQuery arr, Component a) =>
ComponentID -> arr a a
setDyn ComponentID
cId)

instance ArrowQuery Query where
  set :: forall a. (Component a) => Query a a
  set :: forall a. Component a => Query a a
set = (Components -> (ReadsWrites, Components, DynamicQuery a a))
-> Query a a
forall i o.
(Components -> (ReadsWrites, Components, DynamicQuery i o))
-> Query i o
Query ((Components -> (ReadsWrites, Components, DynamicQuery a a))
 -> Query a a)
-> (Components -> (ReadsWrites, Components, DynamicQuery a a))
-> Query a a
forall a b. (a -> b) -> a -> b
$ \Components
cs ->
    let (ComponentID
cId, Components
cs') = forall a. Component a => Components -> (ComponentID, Components)
CS.insert @a Components
cs
     in (Set ComponentID -> Set ComponentID -> ReadsWrites
ReadsWrites Set ComponentID
forall a. Set a
Set.empty (ComponentID -> Set ComponentID
forall a. a -> Set a
Set.singleton ComponentID
cId), Components
cs', ComponentID -> DynamicQuery a a
forall a. Component a => ComponentID -> DynamicQuery a a
forall (arr :: * -> * -> *) a.
(ArrowDynamicQuery arr, Component a) =>
ComponentID -> arr a a
setDyn ComponentID
cId)

fromReader :: QueryReader i o -> Query i o
fromReader :: forall i o. QueryReader i o -> Query i o
fromReader (QueryReader Components -> (Set ComponentID, Components, DynamicQueryReader i o)
f) = (Components -> (ReadsWrites, Components, DynamicQuery i o))
-> Query i o
forall i o.
(Components -> (ReadsWrites, Components, DynamicQuery i o))
-> Query i o
Query ((Components -> (ReadsWrites, Components, DynamicQuery i o))
 -> Query i o)
-> (Components -> (ReadsWrites, Components, DynamicQuery i o))
-> Query i o
forall a b. (a -> b) -> a -> b
$ \Components
cs ->
  let (Set ComponentID
cIds, Components
cs', DynamicQueryReader i o
dynQ) = Components -> (Set ComponentID, Components, DynamicQueryReader i o)
f Components
cs in (Set ComponentID -> Set ComponentID -> ReadsWrites
ReadsWrites Set ComponentID
cIds Set ComponentID
forall a. Set a
Set.empty, Components
cs', DynamicQueryReader i o -> DynamicQuery i o
forall i o. DynamicQueryReader i o -> DynamicQuery i o
fromDynReader DynamicQueryReader i o
dynQ)

toReader :: Query i o -> QueryReader i o
toReader :: forall i o. Query i o -> QueryReader i o
toReader (Query Components -> (ReadsWrites, Components, DynamicQuery i o)
f) = (Components
 -> (Set ComponentID, Components, DynamicQueryReader i o))
-> QueryReader i o
forall i o.
(Components
 -> (Set ComponentID, Components, DynamicQueryReader i o))
-> QueryReader i o
QueryReader ((Components
  -> (Set ComponentID, Components, DynamicQueryReader i o))
 -> QueryReader i o)
-> (Components
    -> (Set ComponentID, Components, DynamicQueryReader i o))
-> QueryReader i o
forall a b. (a -> b) -> a -> b
$ \Components
cs ->
  let (ReadsWrites
rws, Components
cs', DynamicQuery i o
dynQ) = Components -> (ReadsWrites, Components, DynamicQuery i o)
f Components
cs in (ReadsWrites -> Set ComponentID
reads ReadsWrites
rws, Components
cs', DynamicQuery i o -> DynamicQueryReader i o
forall i o. DynamicQuery i o -> DynamicQueryReader i o
toDynReader DynamicQuery i o
dynQ)

-- | Reads and writes of a `Query`.
data ReadsWrites = ReadsWrites
  { ReadsWrites -> Set ComponentID
reads :: !(Set ComponentID),
    ReadsWrites -> Set ComponentID
writes :: !(Set ComponentID)
  }
  deriving (Int -> ReadsWrites -> ShowS
[ReadsWrites] -> ShowS
ReadsWrites -> String
(Int -> ReadsWrites -> ShowS)
-> (ReadsWrites -> String)
-> ([ReadsWrites] -> ShowS)
-> Show ReadsWrites
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ReadsWrites -> ShowS
showsPrec :: Int -> ReadsWrites -> ShowS
$cshow :: ReadsWrites -> String
show :: ReadsWrites -> String
$cshowList :: [ReadsWrites] -> ShowS
showList :: [ReadsWrites] -> ShowS
Show)

instance Semigroup ReadsWrites where
  ReadsWrites Set ComponentID
r1 Set ComponentID
w1 <> :: ReadsWrites -> ReadsWrites -> ReadsWrites
<> ReadsWrites Set ComponentID
r2 Set ComponentID
w2 = Set ComponentID -> Set ComponentID -> ReadsWrites
ReadsWrites (Set ComponentID
r1 Set ComponentID -> Set ComponentID -> Set ComponentID
forall a. Semigroup a => a -> a -> a
<> Set ComponentID
r2) (Set ComponentID
w1 Set ComponentID -> Set ComponentID -> Set ComponentID
forall a. Semigroup a => a -> a -> a
<> Set ComponentID
w2)

instance Monoid ReadsWrites where
  mempty :: ReadsWrites
mempty = Set ComponentID -> Set ComponentID -> ReadsWrites
ReadsWrites Set ComponentID
forall a. Monoid a => a
mempty Set ComponentID
forall a. Monoid a => a
mempty

-- | `True` if the reads and writes of two `Query`s overlap.
disjoint :: ReadsWrites -> ReadsWrites -> Bool
disjoint :: ReadsWrites -> ReadsWrites -> Bool
disjoint ReadsWrites
a ReadsWrites
b =
  Set ComponentID -> Set ComponentID -> Bool
forall a. Ord a => Set a -> Set a -> Bool
Set.disjoint (ReadsWrites -> Set ComponentID
reads ReadsWrites
a) (ReadsWrites -> Set ComponentID
writes ReadsWrites
b)
    Bool -> Bool -> Bool
|| Set ComponentID -> Set ComponentID -> Bool
forall a. Ord a => Set a -> Set a -> Bool
Set.disjoint (ReadsWrites -> Set ComponentID
reads ReadsWrites
b) (ReadsWrites -> Set ComponentID
writes ReadsWrites
a)
    Bool -> Bool -> Bool
|| Set ComponentID -> Set ComponentID -> Bool
forall a. Ord a => Set a -> Set a -> Bool
Set.disjoint (ReadsWrites -> Set ComponentID
writes ReadsWrites
b) (ReadsWrites -> Set ComponentID
writes ReadsWrites
a)

-- | Match all entities.
all :: i -> Query i a -> Entities -> ([a], Entities)
all :: forall i a. i -> Query i a -> Entities -> ([a], Entities)
all i
i = i -> QueryReader i a -> Entities -> ([a], Entities)
forall i a. i -> QueryReader i a -> Entities -> ([a], Entities)
QR.all i
i (QueryReader i a -> Entities -> ([a], Entities))
-> (Query i a -> QueryReader i a)
-> Query i a
-> Entities
-> ([a], Entities)
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
. Query i a -> QueryReader i a
forall i o. Query i o -> QueryReader i o
toReader

-- | Map all matched entities.
map :: i -> Query i a -> Entities -> ([a], Entities)
map :: forall i a. i -> Query i a -> Entities -> ([a], Entities)
map i
i Query i a
q Entities
es =
  let (ReadsWrites
rws, Components
cs', DynamicQuery i a
dynQ) = Query i a
-> Components -> (ReadsWrites, Components, DynamicQuery i a)
forall i o.
Query i o
-> Components -> (ReadsWrites, Components, DynamicQuery i o)
runQuery Query i a
q (Entities -> Components
components Entities
es)
      cIds :: Set ComponentID
cIds = ReadsWrites -> Set ComponentID
reads ReadsWrites
rws Set ComponentID -> Set ComponentID -> Set ComponentID
forall a. Semigroup a => a -> a -> a
<> ReadsWrites -> Set ComponentID
writes ReadsWrites
rws
      ([a]
as, Entities
es') = Set ComponentID
-> i -> DynamicQuery i a -> Entities -> ([a], Entities)
forall i a.
Set ComponentID
-> i -> DynamicQuery i a -> Entities -> ([a], Entities)
mapDyn Set ComponentID
cIds i
i DynamicQuery i a
dynQ Entities
es
   in ([a]
as, Entities
es' {components = cs'})