{-# 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
  { forall (m :: * -> *) (cs :: [*]).
World m cs -> HSet (WorldComponents m cs)
worldComponents :: !(HSet (WorldComponents m cs)),
    forall (m :: * -> *) (cs :: [*]). World m cs -> Entities
worldEntities :: {-# UNPACK #-} !Entities,
    forall (m :: * -> *) (cs :: [*]).
World m cs -> WorldEntityComponents m cs
worldEntityComponents :: {-# UNPACK #-} !(WorldEntityComponents m cs)
  }

empty :: (Monad m, Empty m (HSet (WorldComponents m cs))) => m (World m cs)
empty :: forall (m :: * -> *) (cs :: [*]).
(Monad m, Empty m (HSet (WorldComponents m cs))) =>
m (World m cs)
empty = do
  HSet (WorldComponents m cs)
cs <- m (HSet (WorldComponents m cs))
forall {k} (m :: k -> *) (a :: k). Empty m a => m a
Storage.empty
  World m cs -> m (World m cs)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (World m cs -> m (World m cs)) -> World m cs -> m (World m cs)
forall a b. (a -> b) -> a -> b
$ HSet (WorldComponents m cs)
-> Entities -> WorldEntityComponents m cs -> World m cs
forall (m :: * -> *) (cs :: [*]).
HSet (WorldComponents m cs)
-> Entities -> WorldEntityComponents m cs -> World m cs
World HSet (WorldComponents m cs)
cs Entities
emptyEntities WorldEntityComponents m cs
forall a. IntMap a
IntMap.empty
{-# INLINE empty #-}

lookupStorage ::
  (Lookup (ComponentStorage m c c) (WorldComponents m cs)) =>
  World m cs ->
  ComponentStorage m c c
lookupStorage :: forall (m :: * -> *) c (cs :: [*]).
Lookup (ComponentStorage m c c) (WorldComponents m cs) =>
World m cs -> ComponentStorage m c c
lookupStorage = HSet (WorldComponents m cs) -> ComponentStorage m c c
forall t (ts :: [*]). Lookup t ts => HSet ts -> t
HS.lookup (HSet (WorldComponents m cs) -> ComponentStorage m c c)
-> (World m cs -> HSet (WorldComponents m cs))
-> World m cs
-> ComponentStorage m c c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. World m cs -> HSet (WorldComponents m cs)
forall (m :: * -> *) (cs :: [*]).
World m cs -> HSet (WorldComponents m cs)
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 :: 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 ComponentStorage m c c -> m (ComponentStorage m c c)
f World m cs
w = do
  HSet (WorldComponents m cs)
cs <- forall (m :: * -> *) t (ts :: [*]).
AdjustM m t ts =>
(t -> m t) -> HSet ts -> m (HSet ts)
HS.adjustM @m @(ComponentStorage m c c) ComponentStorage m c c -> m (ComponentStorage m c c)
f (World m cs -> HSet (WorldComponents m cs)
forall (m :: * -> *) (cs :: [*]).
World m cs -> HSet (WorldComponents m cs)
worldComponents World m cs
w)
  World m cs -> m (World m cs)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (World m cs -> m (World m cs)) -> World m cs -> m (World m cs)
forall a b. (a -> b) -> a -> b
$ World m cs
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 :: forall (m :: * -> *) (cs :: [*]) c.
(AdjustM m (ComponentStorage m c c) (WorldComponents m cs),
 PrimMonad m, Component m c, Entity m ~ Entity, Typeable c,
 Storage m (ComponentStorage m c)) =>
Entity -> World m cs -> m (World m cs)
removeComponent Entity
entity World m cs
w = do
  let entityIdx :: Key
entityIdx = Word32 -> Key
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Entity -> Word32
entityIndex Entity
entity)
      componentType :: TypeRep
componentType = Proxy c -> TypeRep
forall {k} (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep (Proxy c
forall {k} (t :: k). Proxy t
Proxy :: Proxy c)
      hooks :: Hooks m
hooks = Proxy c -> Hooks m
forall (proxy :: * -> *). proxy c -> Hooks m
forall (m :: * -> *) a (proxy :: * -> *).
Component m a =>
proxy a -> Hooks m
componentHooks (Proxy c
forall {k} (t :: k). Proxy t
Proxy :: Proxy c)
  -- Run the onRemove hook first

  Hooks m -> Entity m -> m ()
forall (m :: * -> *). Hooks m -> Entity m -> m ()
onRemove Hooks m
hooks Entity m
Entity
entity
  World m cs
w' <- 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 @_ @_ @c (Entity -> ComponentStorage m c c -> m (ComponentStorage m c c)
forall a.
Entity -> ComponentStorage m c a -> m (ComponentStorage m c a)
forall (m :: * -> *) (s :: * -> *) a.
Storage m s =>
Entity -> s a -> m (s a)
removeStorage Entity
entity) World m cs
w
  let entityComponents' :: IntMap
  (Map
     TypeRep
     (HSet (WorldComponents m cs) -> m (HSet (WorldComponents m cs))))
entityComponents' = (Map
   TypeRep
   (HSet (WorldComponents m cs) -> m (HSet (WorldComponents m cs)))
 -> Map
      TypeRep
      (HSet (WorldComponents m cs) -> m (HSet (WorldComponents m cs))))
-> Key
-> IntMap
     (Map
        TypeRep
        (HSet (WorldComponents m cs) -> m (HSet (WorldComponents m cs))))
-> IntMap
     (Map
        TypeRep
        (HSet (WorldComponents m cs) -> m (HSet (WorldComponents m cs))))
forall a. (a -> a) -> Key -> IntMap a -> IntMap a
IntMap.adjust (TypeRep
-> Map
     TypeRep
     (HSet (WorldComponents m cs) -> m (HSet (WorldComponents m cs)))
-> Map
     TypeRep
     (HSet (WorldComponents m cs) -> m (HSet (WorldComponents m cs)))
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete TypeRep
componentType) Key
entityIdx (World m cs
-> IntMap
     (Map
        TypeRep
        (HSet (WorldComponents m cs) -> m (HSet (WorldComponents m cs))))
forall (m :: * -> *) (cs :: [*]).
World m cs -> WorldEntityComponents m cs
worldEntityComponents World m cs
w)
  World m cs -> m (World m cs)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (World m cs -> m (World m cs)) -> World m cs -> m (World m cs)
forall a b. (a -> b) -> a -> b
$ World m cs
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' :: forall (m :: * -> *) c (cs :: [*]).
(AdjustM m (SparseStorage m c) cs, PrimMonad m) =>
Entity -> HSet cs -> m (HSet cs)
removeComponent' Entity
e HSet cs
components = forall (m :: * -> *) t (ts :: [*]).
AdjustM m t ts =>
(t -> m t) -> HSet ts -> m (HSet ts)
HS.adjustM @m @(SparseStorage m c) SparseStorage m c -> m (SparseStorage m c)
go HSet cs
components
  where
    go :: SparseStorage m c -> m (SparseStorage m c)
go SparseStorage m c
s = do
      SparseSet Word32 c
s' <- SparseStorage m c -> m (SparseSet Word32 c)
forall (m :: * -> *) i a.
PrimMonad m =>
MSparseSet (PrimState m) i a -> m (SparseSet i a)
S.freeze SparseStorage m c
s
      SparseSet Word32 c -> m (SparseStorage m c)
forall (m :: * -> *) i a.
PrimMonad m =>
SparseSet i a -> m (MSparseSet (PrimState m) i a)
S.thaw (Word32 -> SparseSet Word32 c -> SparseSet Word32 c
forall i a. Integral i => i -> SparseSet i a -> SparseSet i a
S.delete (Entity -> Word32
entityIndex Entity
e) SparseSet Word32 c
s')
{-# INLINE removeComponent' #-}

remove :: (Monad m) => Entity -> World m cs -> m (World m cs)
remove :: forall (m :: * -> *) (cs :: [*]).
Monad m =>
Entity -> World m cs -> m (World m cs)
remove Entity
entity World m cs
w = do
  let entityIdx :: Key
entityIdx = Word32 -> Key
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Entity -> Word32
entityIndex Entity
entity)
  case Key
-> IntMap
     (Map
        TypeRep
        (HSet (WorldComponents m cs) -> m (HSet (WorldComponents m cs))))
-> Maybe
     (Map
        TypeRep
        (HSet (WorldComponents m cs) -> m (HSet (WorldComponents m cs))))
forall a. Key -> IntMap a -> Maybe a
IntMap.lookup Key
entityIdx (World m cs
-> IntMap
     (Map
        TypeRep
        (HSet (WorldComponents m cs) -> m (HSet (WorldComponents m cs))))
forall (m :: * -> *) (cs :: [*]).
World m cs -> WorldEntityComponents m cs
worldEntityComponents World m cs
w) of
    Maybe
  (Map
     TypeRep
     (HSet (WorldComponents m cs) -> m (HSet (WorldComponents m cs))))
Nothing -> World m cs -> m (World m cs)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return World m cs
w
    Just Map
  TypeRep
  (HSet (WorldComponents m cs) -> m (HSet (WorldComponents m cs)))
componentRemovalFunctions -> do
      HSet (WorldComponents m cs)
cs' <- ((HSet (WorldComponents m cs) -> m (HSet (WorldComponents m cs)))
 -> (HSet (WorldComponents m cs) -> m (HSet (WorldComponents m cs)))
 -> HSet (WorldComponents m cs)
 -> m (HSet (WorldComponents m cs)))
-> (HSet (WorldComponents m cs) -> m (HSet (WorldComponents m cs)))
-> [HSet (WorldComponents m cs) -> m (HSet (WorldComponents m cs))]
-> HSet (WorldComponents m cs)
-> m (HSet (WorldComponents m cs))
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (HSet (WorldComponents m cs) -> m (HSet (WorldComponents m cs)))
-> (HSet (WorldComponents m cs) -> m (HSet (WorldComponents m cs)))
-> HSet (WorldComponents m cs)
-> m (HSet (WorldComponents m cs))
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
(<=<) HSet (WorldComponents m cs) -> m (HSet (WorldComponents m cs))
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Map
  TypeRep
  (HSet (WorldComponents m cs) -> m (HSet (WorldComponents m cs)))
-> [HSet (WorldComponents m cs) -> m (HSet (WorldComponents m cs))]
forall k a. Map k a -> [a]
Map.elems Map
  TypeRep
  (HSet (WorldComponents m cs) -> m (HSet (WorldComponents m cs)))
componentRemovalFunctions) (World m cs -> HSet (WorldComponents m cs)
forall (m :: * -> *) (cs :: [*]).
World m cs -> HSet (WorldComponents m cs)
worldComponents World m cs
w)
      let entityComponents' :: IntMap
  (Map
     TypeRep
     (HSet (WorldComponents m cs) -> m (HSet (WorldComponents m cs))))
entityComponents' = Key
-> IntMap
     (Map
        TypeRep
        (HSet (WorldComponents m cs) -> m (HSet (WorldComponents m cs))))
-> IntMap
     (Map
        TypeRep
        (HSet (WorldComponents m cs) -> m (HSet (WorldComponents m cs))))
forall a. Key -> IntMap a -> IntMap a
IntMap.delete Key
entityIdx (World m cs
-> IntMap
     (Map
        TypeRep
        (HSet (WorldComponents m cs) -> m (HSet (WorldComponents m cs))))
forall (m :: * -> *) (cs :: [*]).
World m cs -> WorldEntityComponents m cs
worldEntityComponents World m cs
w)
      World m cs -> m (World m cs)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (World m cs -> m (World m cs)) -> World m cs -> m (World m cs)
forall a b. (a -> b) -> a -> b
$ World m cs
w {worldComponents = cs', worldEntityComponents = entityComponents'}
{-# INLINE remove #-}