{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

module Data.Aztecs.Query
  ( ReadWrites (..),
    Query (..),
    entity,
    read,
    buildQuery,
    all,
    all',
    get,
    get',
    write,
  )
where

import Control.Monad.IO.Class (MonadIO (..))
import Data.Aztecs.World (Component, Entity, World (..))
import Data.Aztecs.World.Archetypes (Archetype, ArchetypeId, ArchetypeState (..), archetype)
import qualified Data.Aztecs.World.Archetypes as A
import Data.Foldable (foldrM)
import qualified Data.Map as Map
import Data.Set (Set)
import Data.Typeable
import Prelude hiding (all, read)

-- | Component IDs to read and write.
data ReadWrites = ReadWrites (Set TypeRep) (Set TypeRep)

instance Semigroup ReadWrites where
  ReadWrites Set TypeRep
rs Set TypeRep
ws <> :: ReadWrites -> ReadWrites -> ReadWrites
<> ReadWrites Set TypeRep
rs' Set TypeRep
ws' = Set TypeRep -> Set TypeRep -> ReadWrites
ReadWrites (Set TypeRep
rs Set TypeRep -> Set TypeRep -> Set TypeRep
forall a. Semigroup a => a -> a -> a
<> Set TypeRep
rs') (Set TypeRep
ws Set TypeRep -> Set TypeRep -> Set TypeRep
forall a. Semigroup a => a -> a -> a
<> Set TypeRep
ws')

instance Monoid ReadWrites where
  mempty :: ReadWrites
mempty = Set TypeRep -> Set TypeRep -> ReadWrites
ReadWrites Set TypeRep
forall a. Monoid a => a
mempty Set TypeRep
forall a. Monoid a => a
mempty

-- | Builder for a `Query`.
data Query m a where
  PureQ :: a -> Query m a
  MapQ :: (a -> b) -> Query m a -> Query m b
  AppQ :: Query m (a -> b) -> Query m a -> Query m b
  BindQ :: Query m a -> (a -> Query m b) -> Query m b
  EntityQ :: Query m Entity
  ReadQ :: (Component c) => Archetype -> Query m c
  WriteQ :: (Component c) => (c -> c) -> Archetype -> Query m c
  LiftQ :: m a -> Query m a

instance Functor (Query m) where
  fmap :: forall a b. (a -> b) -> Query m a -> Query m b
fmap = (a -> b) -> Query m a -> Query m b
forall a b (m :: * -> *). (a -> b) -> Query m a -> Query m b
MapQ

instance Applicative (Query m) where
  pure :: forall a. a -> Query m a
pure = a -> Query m a
forall a (m :: * -> *). a -> Query m a
PureQ
  <*> :: forall a b. Query m (a -> b) -> Query m a -> Query m b
(<*>) = Query m (a -> b) -> Query m a -> Query m b
forall (m :: * -> *) a b.
Query m (a -> b) -> Query m a -> Query m b
AppQ

instance Monad (Query m) where
  >>= :: forall a b. Query m a -> (a -> Query m b) -> Query m b
(>>=) = Query m a -> (a -> Query m b) -> Query m b
forall (m :: * -> *) a b.
Query m a -> (a -> Query m b) -> Query m b
BindQ

entity :: Query m Entity
entity :: forall (m :: * -> *). Query m Entity
entity = Query m Entity
forall (m :: * -> *). Query m Entity
EntityQ

-- | Read a `Component`.
read :: forall m c. (Component c) => Query m c
read :: forall (m :: * -> *) c. Component c => Query m c
read = Archetype -> Query m c
forall c (m :: * -> *). Component c => Archetype -> Query m c
ReadQ (forall c. Component c => Archetype
archetype @c)

-- | Alter a `Component`.
write :: forall m c. (Component c) => (c -> c) -> Query m c
write :: forall (m :: * -> *) c. Component c => (c -> c) -> Query m c
write c -> c
c = (c -> c) -> Archetype -> Query m c
forall c (m :: * -> *).
Component c =>
(c -> c) -> Archetype -> Query m c
WriteQ c -> c
c (forall c. Component c => Archetype
archetype @c)

buildQuery :: Query m a -> Archetype
buildQuery :: forall (m :: * -> *) a. Query m a -> Archetype
buildQuery (PureQ a
_) = Archetype
forall a. Monoid a => a
mempty
buildQuery (MapQ a -> a
_ Query m a
qb) = Query m a -> Archetype
forall (m :: * -> *) a. Query m a -> Archetype
buildQuery Query m a
qb
buildQuery (AppQ Query m (a -> a)
f Query m a
a) = Query m (a -> a) -> Archetype
forall (m :: * -> *) a. Query m a -> Archetype
buildQuery Query m (a -> a)
f Archetype -> Archetype -> Archetype
forall a. Semigroup a => a -> a -> a
<> Query m a -> Archetype
forall (m :: * -> *) a. Query m a -> Archetype
buildQuery Query m a
a
buildQuery Query m a
EntityQ = Archetype
forall a. Monoid a => a
mempty
buildQuery (ReadQ Archetype
a) = Archetype
a
buildQuery (WriteQ a -> a
_ Archetype
a) = Archetype
a
buildQuery (BindQ Query m a
a a -> Query m a
_) = Query m a -> Archetype
forall (m :: * -> *) a. Query m a -> Archetype
buildQuery Query m a
a
buildQuery (LiftQ m a
_) = Archetype
forall a. Monoid a => a
mempty

all :: (MonadIO m) => ArchetypeId -> Query m a -> World -> m [a]
all :: forall (m :: * -> *) a.
MonadIO m =>
ArchetypeId -> Query m a -> World -> m [a]
all ArchetypeId
a Query m a
qb w :: World
w@(World Components
_ Archetypes
as) = case ArchetypeId -> Archetypes -> Maybe ArchetypeState
A.getArchetype ArchetypeId
a Archetypes
as of
  Just ArchetypeState
s -> ArchetypeState -> Query m a -> World -> m [a]
forall (m :: * -> *) a.
MonadIO m =>
ArchetypeState -> Query m a -> World -> m [a]
all' ArchetypeState
s Query m a
qb World
w
  Maybe ArchetypeState
Nothing -> [a] -> m [a]
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return []

all' :: (MonadIO m) => ArchetypeState -> Query m a -> World -> m [a]
all' :: forall (m :: * -> *) a.
MonadIO m =>
ArchetypeState -> Query m a -> World -> m [a]
all' es :: ArchetypeState
es@(ArchetypeState Archetype
_ Map Entity ArchetypeComponents
m [ArchetypeId]
_) Query m a
q World
w =
  (Entity -> [a] -> m [a]) -> [a] -> [Entity] -> m [a]
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> b -> m b) -> b -> t a -> m b
foldrM
    ( \Entity
e [a]
acc -> do
        Maybe a
a <- Entity -> ArchetypeState -> Query m a -> World -> m (Maybe a)
forall (m :: * -> *) a.
MonadIO m =>
Entity -> ArchetypeState -> Query m a -> World -> m (Maybe a)
get' Entity
e ArchetypeState
es Query m a
q World
w
        [a] -> m [a]
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ([a] -> m [a]) -> [a] -> m [a]
forall a b. (a -> b) -> a -> b
$ case Maybe a
a of
          Just a
a' -> (a
a' a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
acc)
          Maybe a
Nothing -> [a]
acc
    )
    []
    (Map Entity ArchetypeComponents -> [Entity]
forall k a. Map k a -> [k]
Map.keys Map Entity ArchetypeComponents
m)

get :: (MonadIO m) => ArchetypeId -> Query m a -> Entity -> World -> m (Maybe a)
get :: forall (m :: * -> *) a.
MonadIO m =>
ArchetypeId -> Query m a -> Entity -> World -> m (Maybe a)
get ArchetypeId
a Query m a
qb Entity
e w :: World
w@(World Components
_ Archetypes
as) = case ArchetypeId -> Archetypes -> Maybe ArchetypeState
A.getArchetype ArchetypeId
a Archetypes
as of
  Just ArchetypeState
s -> Entity -> ArchetypeState -> Query m a -> World -> m (Maybe a)
forall (m :: * -> *) a.
MonadIO m =>
Entity -> ArchetypeState -> Query m a -> World -> m (Maybe a)
get' Entity
e ArchetypeState
s Query m a
qb World
w
  Maybe ArchetypeState
Nothing -> Maybe a -> m (Maybe a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing

get' :: (MonadIO m) => Entity -> ArchetypeState -> Query m a -> World -> m (Maybe a)
get' :: forall (m :: * -> *) a.
MonadIO m =>
Entity -> ArchetypeState -> Query m a -> World -> m (Maybe a)
get' Entity
_ ArchetypeState
_ (PureQ a
a) World
_ = Maybe a -> m (Maybe a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe a -> m (Maybe a)) -> Maybe a -> m (Maybe a)
forall a b. (a -> b) -> a -> b
$ a -> Maybe a
forall a. a -> Maybe a
Just a
a
get' Entity
e ArchetypeState
es (MapQ a -> a
f Query m a
qb) World
w = do
  Maybe a
a <- Entity -> ArchetypeState -> Query m a -> World -> m (Maybe a)
forall (m :: * -> *) a.
MonadIO m =>
Entity -> ArchetypeState -> Query m a -> World -> m (Maybe a)
get' Entity
e ArchetypeState
es Query m a
qb World
w
  Maybe a -> m (Maybe a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe a -> m (Maybe a)) -> Maybe a -> m (Maybe a)
forall a b. (a -> b) -> a -> b
$ (a -> a) -> Maybe a -> Maybe a
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> a
f Maybe a
a
get' Entity
e ArchetypeState
es (AppQ Query m (a -> a)
fqb Query m a
aqb) World
w = do
  Maybe (a -> a)
f <- Entity
-> ArchetypeState
-> Query m (a -> a)
-> World
-> m (Maybe (a -> a))
forall (m :: * -> *) a.
MonadIO m =>
Entity -> ArchetypeState -> Query m a -> World -> m (Maybe a)
get' Entity
e ArchetypeState
es Query m (a -> a)
fqb World
w
  Maybe a
a <- Entity -> ArchetypeState -> Query m a -> World -> m (Maybe a)
forall (m :: * -> *) a.
MonadIO m =>
Entity -> ArchetypeState -> Query m a -> World -> m (Maybe a)
get' Entity
e ArchetypeState
es Query m a
aqb World
w
  Maybe a -> m (Maybe a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe a -> m (Maybe a)) -> Maybe a -> m (Maybe a)
forall a b. (a -> b) -> a -> b
$ Maybe (a -> a)
f Maybe (a -> a) -> Maybe a -> Maybe a
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe a
a
get' Entity
e ArchetypeState
_ Query m a
EntityQ World
_ = Maybe a -> m (Maybe a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe a -> m (Maybe a)) -> Maybe a -> m (Maybe a)
forall a b. (a -> b) -> a -> b
$ a -> Maybe a
forall a. a -> Maybe a
Just a
Entity
e
get' Entity
e (ArchetypeState Archetype
_ Map Entity ArchetypeComponents
m [ArchetypeId]
_) (ReadQ Archetype
_) World
_ = Maybe a -> m (Maybe a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe a -> m (Maybe a)) -> Maybe a -> m (Maybe a)
forall a b. (a -> b) -> a -> b
$ do
  ArchetypeComponents
cs <- Entity
-> Map Entity ArchetypeComponents -> Maybe ArchetypeComponents
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Entity
e Map Entity ArchetypeComponents
m
  (a
c, a -> IO ()
_) <- ArchetypeComponents -> Maybe (a, a -> IO ())
forall c.
Component c =>
ArchetypeComponents -> Maybe (c, c -> IO ())
A.getArchetypeComponent ArchetypeComponents
cs
  a -> Maybe a
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return a
c
get' Entity
e (ArchetypeState Archetype
_ Map Entity ArchetypeComponents
m [ArchetypeId]
_) (WriteQ a -> a
f Archetype
_) World
_ = do
  let res :: Maybe (a, a -> IO ())
res = do
        ArchetypeComponents
cs <- Entity
-> Map Entity ArchetypeComponents -> Maybe ArchetypeComponents
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Entity
e Map Entity ArchetypeComponents
m
        ArchetypeComponents -> Maybe (a, a -> IO ())
forall c.
Component c =>
ArchetypeComponents -> Maybe (c, c -> IO ())
A.getArchetypeComponent ArchetypeComponents
cs
  case Maybe (a, a -> IO ())
res of
    Just (a
c, a -> IO ()
g) -> do
      let c' :: a
c' = a -> a
f a
c
      IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ a -> IO ()
g a
c'
      Maybe a -> m (Maybe a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe a -> m (Maybe a)) -> Maybe a -> m (Maybe a)
forall a b. (a -> b) -> a -> b
$ a -> Maybe a
forall a. a -> Maybe a
Just a
c'
    Maybe (a, a -> IO ())
Nothing -> Maybe a -> m (Maybe a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing
get' Entity
e ArchetypeState
es (BindQ Query m a
qb a -> Query m a
f) World
w = do
  Maybe a
a <- Entity -> ArchetypeState -> Query m a -> World -> m (Maybe a)
forall (m :: * -> *) a.
MonadIO m =>
Entity -> ArchetypeState -> Query m a -> World -> m (Maybe a)
get' Entity
e ArchetypeState
es Query m a
qb World
w
  case Maybe a
a of
    Just a
a' -> Entity -> ArchetypeState -> Query m a -> World -> m (Maybe a)
forall (m :: * -> *) a.
MonadIO m =>
Entity -> ArchetypeState -> Query m a -> World -> m (Maybe a)
get' Entity
e ArchetypeState
es (a -> Query m a
f a
a') World
w
    Maybe a
Nothing -> Maybe a -> m (Maybe a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing
get' Entity
_ ArchetypeState
_ (LiftQ m a
a) World
_ = (a -> Maybe a) -> m a -> m (Maybe a)
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> Maybe a
forall a. a -> Maybe a
Just m a
a