World.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}

module Aztecs.World
  ( World (..),
    WorldComponents (..),
    empty,
    removeComponent,
    removeComponent',
    remove,
    SparseStorage,
  )
where

import qualified Aztecs.ECS.Class as ECS
import Aztecs.ECS.Component
import Aztecs.ECS.HSet
import qualified Aztecs.ECS.HSet as HS
import Aztecs.Entity
import Aztecs.Storage hiding (empty)
import qualified Aztecs.Storage as Storage
import Aztecs.World.Entities
import Control.Monad
import Control.Monad.Primitive
import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as IntMap
import Data.Kind
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Proxy
import qualified Data.SparseSet.Strict as S
import Data.SparseSet.Strict.Mutable (MSparseSet)
import Data.Typeable
import Data.Word

type SparseStorage m = MSparseSet (PrimState m) Word32

type family WorldComponents m (cs :: [Type]) :: [Type] where
  WorldComponents m '[] = '[]
  WorldComponents m (c ': cs) = ComponentStorage m c c ': WorldComponents m cs

type WorldComponentSet m cs =
  HSet (WorldComponents m cs)

type WorldEntityComponents m cs =
  IntMap (Map TypeRep (WorldComponentSet m cs -> m (WorldComponentSet m cs)))

data World m cs = World
  { worldComponents :: !(HSet (WorldComponents m cs)),
    worldEntities :: {-# UNPACK #-} !Entities,
    worldEntityComponents :: {-# UNPACK #-} !(WorldEntityComponents m cs)
  }

empty :: (Monad m, Empty m (HSet (WorldComponents m cs))) => m (World m cs)
empty = do
  cs <- Storage.empty
  return $ World cs emptyEntities IntMap.empty
{-# INLINE empty #-}

lookupStorage ::
  (Lookup (ComponentStorage m c c) (WorldComponents m cs)) =>
  World m cs ->
  ComponentStorage m c c
lookupStorage = HS.lookup . worldComponents
{-# INLINE lookupStorage #-}

adjustStorage ::
  forall m cs c.
  ( PrimMonad m,
    Typeable c,
    Component m c,
    AdjustM m (ComponentStorage m c c) (WorldComponents m cs),
    Storage m (ComponentStorage m c)
  ) =>
  (ComponentStorage m c c -> m (ComponentStorage m c c)) ->
  World m cs ->
  m (World m cs)
adjustStorage f w = do
  cs <- HS.adjustM @m @(ComponentStorage m c c) f (worldComponents w)
  return $ w {worldComponents = cs}
{-# INLINE adjustStorage #-}

removeComponent ::
  forall m cs c.
  ( AdjustM m (ComponentStorage m c c) (WorldComponents m cs),
    PrimMonad m,
    Component m c,
    ECS.Entity m ~ Entity,
    Typeable c,
    Storage m (ComponentStorage m c)
  ) =>
  Entity ->
  World m cs ->
  m (World m cs)
removeComponent entity w = do
  let entityIdx = fromIntegral (entityIndex entity)
      componentType = typeRep (Proxy :: Proxy c)
      hooks = componentHooks (Proxy :: Proxy c)
  -- Run the onRemove hook first
  onRemove hooks entity
  w' <- adjustStorage @_ @_ @c (removeStorage entity) w
  let entityComponents' = IntMap.adjust (Map.delete componentType) entityIdx (worldEntityComponents w)
  return $ w' {worldEntityComponents = entityComponents'}
{-# INLINE removeComponent #-}

removeComponent' ::
  forall m (c :: Type) cs.
  (AdjustM m (SparseStorage m c) cs, PrimMonad m) =>
  Entity ->
  HSet cs ->
  m (HSet cs)
removeComponent' e components = HS.adjustM @m @(SparseStorage m c) go components
  where
    go s = do
      s' <- S.freeze s
      S.thaw (S.delete (entityIndex e) s')
{-# INLINE removeComponent' #-}

remove :: (Monad m) => Entity -> World m cs -> m (World m cs)
remove entity w = do
  let entityIdx = fromIntegral (entityIndex entity)
  case IntMap.lookup entityIdx (worldEntityComponents w) of
    Nothing -> return w
    Just componentRemovalFunctions -> do
      cs' <- foldr (<=<) return (Map.elems componentRemovalFunctions) (worldComponents w)
      let entityComponents' = IntMap.delete entityIdx (worldEntityComponents w)
      return $ w {worldComponents = cs', worldEntityComponents = entityComponents'}
{-# INLINE remove #-}