Internal.hs

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE UndecidableInstances #-}

module Aztecs.Internal
  ( AztecsT (..),
    runAztecsT,
    runAztecsT_,
  )
where

import Aztecs.ECS.Bundle
import Aztecs.ECS.Bundle.Class
import Aztecs.ECS.Class
import Aztecs.ECS.Commands
import Aztecs.ECS.Component (Component (ComponentStorage, componentHooks), Hooks (..))
import Aztecs.ECS.HSet (AdjustM, HSet (..), Lookup (..))
import qualified Aztecs.ECS.HSet as HS
import Aztecs.ECS.Query
import Aztecs.ECS.Query.Internal
import Aztecs.ECS.R
import qualified Aztecs.ECS.Scheduler as Scheduler
import Aztecs.ECS.W
import qualified Aztecs.Entity as E
import Aztecs.Storage
import qualified Aztecs.Storage as S
import Aztecs.World (SparseStorage, WorldComponents)
import qualified Aztecs.World as W
import qualified Aztecs.World.Entities as E
import Control.Monad.Primitive
import Control.Monad.State.Strict
import qualified Data.IntMap.Strict as IntMap
import qualified Data.Map.Strict as Map
import Data.Maybe
import Data.Proxy
import qualified Data.Set as Set
import Data.Typeable
import Prelude hiding (Read, lookup)

newtype AztecsT cs m a = AztecsT {unAztecsT :: StateT (W.World m cs) m a}
  deriving (Functor, Applicative, Monad, MonadIO, PrimMonad)

instance MonadTrans (AztecsT cs) where
  lift = AztecsT . lift
  {-# INLINE lift #-}

instance (PrimMonad m) => ECS (AztecsT cs m) where
  type Entity (AztecsT cs m) = E.Entity
  type Task (AztecsT cs m) = (Commands (AztecsT cs) m)

  spawn b = do
    w <- AztecsT get
    let (e, counter) = E.mkEntityWithCounter (W.worldEntities w)
    AztecsT $ put w {W.worldEntities = counter}
    runBundle b e
    return e
  {-# INLINE spawn #-}
  insert e b = runBundle b e
  {-# INLINE insert #-}
  remove e = AztecsT $ do
    w <- get
    w' <- lift $ W.remove e w
    put w'
  {-# INLINE remove #-}
  task = runCommands
  {-# INLINE task #-}

instance
  ( PrimMonad m,
    Typeable c,
    Component (AztecsT cs m) c,
    AdjustM m (SparseStorage m c) (WorldComponents m cs)
  ) =>
  Bundleable c (AztecsT cs m)
  where
  bundle c = Bundle $ \entity -> do
    w <- AztecsT get
    let entityIdx = fromIntegral $ E.entityIndex entity
        componentType = typeOf c
        go = S.insertStorage entity c
        hooks = componentHooks (Proxy :: Proxy c)
    cs <- lift . HS.adjustM @_ @(SparseStorage m c) go $ W.worldComponents w
    let entityComponents' =
          IntMap.insertWith
            Map.union
            entityIdx
            (Map.singleton componentType (W.removeComponent' @m @c entity))
            (W.worldEntityComponents w)
    AztecsT $ put w {W.worldComponents = cs, W.worldEntityComponents = entityComponents'}
    onInsert hooks entity
  {-# INLINE bundle #-}

runAztecsT :: (Monad m) => AztecsT cs m a -> W.World m cs -> m (a, W.World m cs)
runAztecsT (AztecsT m) = runStateT m
{-# INLINE runAztecsT #-}

runAztecsT_ :: (Monad m) => AztecsT cs m a -> W.World m cs -> m a
runAztecsT_ (AztecsT m) = evalStateT m
{-# INLINE runAztecsT_ #-}

instance (PrimMonad m) => Queryable (AztecsT cs m) E.Entity where
  type QueryableAccess E.Entity = '[]
  queryable = AztecsT $ do
    w <- get
    return . Query . map pure . E.entities $ W.worldEntities w
  {-# INLINE queryable #-}

instance
  ( PrimMonad m,
    Lookup (ComponentStorage m a a) (WorldComponents m cs),
    Storage m (ComponentStorage m a)
  ) =>
  Queryable (AztecsT cs m) (With a)
  where
  type QueryableAccess (With a) = '[With a]
  queryable = AztecsT $ do
    w <- get
    withComponent <-
      lift
        . S.queryStorageR
        . HS.lookup @(ComponentStorage m a a)
        $ W.worldComponents w
    return . fmap (const With) $ withComponent
  {-# INLINE queryable #-}

instance
  ( PrimMonad m,
    Lookup (ComponentStorage m a a) (WorldComponents m cs),
    Storage m (ComponentStorage m a)
  ) =>
  Queryable (AztecsT cs m) (Without a)
  where
  type QueryableAccess (Without a) = '[Without a]
  queryable = AztecsT $ do
    w <- get
    (Query cs) <-
      lift
        . S.queryStorageR
        . HS.lookup @(ComponentStorage m a a)
        $ W.worldComponents w
    let go m = case m of
          Just v -> Nothing
          Nothing -> Just Without
    return . Query $ fmap go cs
  {-# INLINE queryable #-}

instance
  ( PrimMonad m,
    Lookup (ComponentStorage m a a) (WorldComponents m cs),
    Storage (AztecsT cs m) (ComponentStorage m a)
  ) =>
  Queryable (AztecsT cs m) (R a)
  where
  type QueryableAccess (R a) = '[Read a]
  queryable = do
    w <- AztecsT get
    S.queryStorageR . HS.lookup @(ComponentStorage m a a) $ W.worldComponents w
  {-# INLINE queryable #-}

instance
  ( PrimMonad m,
    PrimState m ~ s,
    Lookup (ComponentStorage m a a) (WorldComponents m cs),
    Storage m (ComponentStorage m a)
  ) =>
  Queryable (AztecsT cs m) (W (Commands (AztecsT cs) m) a)
  where
  type QueryableAccess (W (Commands (AztecsT cs) m) a) = '[Write a]
  queryable = AztecsT $ do
    w <- get
    Query results <-
      lift
        . S.queryStorageW
        . HS.lookup @(ComponentStorage m a a)
        $ W.worldComponents w
    let liftToCommands m = Commands $ (,pure ()) <$> m
        go (W r wf mf) =
          W
            (liftToCommands r)
            (liftToCommands . wf)
            (liftToCommands . mf)
    return . Query $ map (fmap go) results
  {-# INLINE queryable #-}

-- Additional instance for direct AztecsT usage in scheduler
instance
  ( PrimMonad m,
    PrimState m ~ s,
    Lookup (ComponentStorage m a a) (WorldComponents m cs),
    Storage m (ComponentStorage m a)
  ) =>
  Queryable (AztecsT cs m) (W (AztecsT cs m) a)
  where
  type QueryableAccess (W (AztecsT cs m) a) = '[Write a]
  queryable = AztecsT $ do
    w <- get
    Query results <-
      lift
        . S.queryStorageW
        . HS.lookup @(ComponentStorage m a a)
        $ W.worldComponents w
    let liftToAztecs (W r wf mf) =
          W
            (AztecsT $ lift r)
            (AztecsT . lift . wf)
            (AztecsT . lift . mf)
    return . Query $ map (fmap liftToAztecs) results
  {-# INLINE queryable #-}