{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}

-- |

-- Module      : Aztecs.ECS.World.Entities

-- Copyright   : (c) Matt Hunzinger, 2025

-- License     : BSD-style (see the LICENSE file in the distribution)

--

-- Maintainer  : matt@hunzinger.me

-- Stability   : provisional

-- Portability : non-portable (GHC extensions)

module Aztecs.ECS.World.Entities
  ( Entities (..),
    empty,
    spawn,
    spawnWithArchetypeId,
    insert,
    insertDyn,
    lookup,
    remove,
    removeWithId,
    despawn,
  )
where

import Aztecs.ECS.Component
import Aztecs.ECS.Entity
import qualified Aztecs.ECS.World.Archetype as A
import Aztecs.ECS.World.Archetypes (ArchetypeID, Archetypes, Node (..))
import qualified Aztecs.ECS.World.Archetypes as AS
import Aztecs.ECS.World.Bundle
import Aztecs.ECS.World.Bundle.Dynamic
import Aztecs.ECS.World.Components (Components (..))
import qualified Aztecs.ECS.World.Components as CS
import Control.DeepSeq
import Data.Dynamic
import Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe
import Data.Set (Set)
import qualified Data.Set as Set
import GHC.Generics
import Prelude hiding (lookup)

-- | World of entities and their components.

--

-- @since 0.9

data Entities = Entities
  { -- | Archetypes.

    --

    -- @since 0.9

    Entities -> Archetypes
archetypes :: !Archetypes,
    -- | Components.

    --

    -- @since 0.9

    Entities -> Components
components :: !Components,
    -- | Entities and their archetype identifiers.

    --

    -- @since 0.9

    Entities -> Map EntityID ArchetypeID
entities :: !(Map EntityID ArchetypeID)
  }
  deriving (Int -> Entities -> ShowS
[Entities] -> ShowS
Entities -> String
(Int -> Entities -> ShowS)
-> (Entities -> String) -> ([Entities] -> ShowS) -> Show Entities
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Entities -> ShowS
showsPrec :: Int -> Entities -> ShowS
$cshow :: Entities -> String
show :: Entities -> String
$cshowList :: [Entities] -> ShowS
showList :: [Entities] -> ShowS
Show, (forall x. Entities -> Rep Entities x)
-> (forall x. Rep Entities x -> Entities) -> Generic Entities
forall x. Rep Entities x -> Entities
forall x. Entities -> Rep Entities x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. Entities -> Rep Entities x
from :: forall x. Entities -> Rep Entities x
$cto :: forall x. Rep Entities x -> Entities
to :: forall x. Rep Entities x -> Entities
Generic, Entities -> ()
(Entities -> ()) -> NFData Entities
forall a. (a -> ()) -> NFData a
$crnf :: Entities -> ()
rnf :: Entities -> ()
NFData)

-- | Empty `World`.

--

-- @since 0.9

empty :: Entities
empty :: Entities
empty =
  Entities
    { archetypes :: Archetypes
archetypes = Archetypes
AS.empty,
      components :: Components
components = Components
CS.empty,
      entities :: Map EntityID ArchetypeID
entities = Map EntityID ArchetypeID
forall a. Monoid a => a
mempty
    }

-- | Spawn a `Bundle`.

--

-- @since 0.9

spawn :: EntityID -> Bundle -> Entities -> Entities
spawn :: EntityID -> Bundle -> Entities -> Entities
spawn EntityID
eId Bundle
b Entities
w =
  let (Set ComponentID
cIds, Components
components', DynamicBundle
dynB) = Bundle
-> Components -> (Set ComponentID, Components, DynamicBundle)
unBundle Bundle
b (Entities -> Components
components Entities
w)
   in case Set ComponentID -> Archetypes -> Maybe ArchetypeID
AS.lookupArchetypeId Set ComponentID
cIds (Entities -> Archetypes
archetypes Entities
w) of
        Just ArchetypeID
aId -> Entities -> Maybe Entities -> Entities
forall a. a -> Maybe a -> a
fromMaybe Entities
w (Maybe Entities -> Entities) -> Maybe Entities -> Entities
forall a b. (a -> b) -> a -> b
$ do
          Node
node <- ArchetypeID -> Archetypes -> Maybe Node
AS.lookup ArchetypeID
aId (Archetypes -> Maybe Node) -> Archetypes -> Maybe Node
forall a b. (a -> b) -> a -> b
$ Entities -> Archetypes
archetypes Entities
w
          let arch' :: Archetype
arch' =
                DynamicBundle -> EntityID -> Archetype -> Archetype
runDynamicBundle
                  DynamicBundle
dynB
                  EntityID
eId
                  ( (Node -> Archetype
nodeArchetype Node
node)
                      { A.entities = Set.insert eId . A.entities $ nodeArchetype node
                      }
                  )
          Entities -> Maybe Entities
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return
            Entities
w
              { archetypes = (archetypes w) {AS.nodes = Map.insert aId node {nodeArchetype = arch'} (AS.nodes $ archetypes w)},
                components = components',
                entities = Map.insert eId aId (entities w)
              }
        Maybe ArchetypeID
Nothing ->
          let arch' :: Archetype
arch' = DynamicBundle -> EntityID -> Archetype -> Archetype
runDynamicBundle DynamicBundle
dynB EntityID
eId (Archetype -> Archetype) -> Archetype -> Archetype
forall a b. (a -> b) -> a -> b
$ EntityID -> Archetype
A.singleton EntityID
eId
              node' :: Node
node' = Node {nodeComponentIds :: Set ComponentID
nodeComponentIds = Set ComponentID
cIds, nodeArchetype :: Archetype
nodeArchetype = Archetype
arch'}
              (ArchetypeID
aId, Archetypes
arches) = Set ComponentID -> Node -> Archetypes -> (ArchetypeID, Archetypes)
AS.insertArchetype Set ComponentID
cIds Node
node' (Archetypes -> (ArchetypeID, Archetypes))
-> Archetypes -> (ArchetypeID, Archetypes)
forall a b. (a -> b) -> a -> b
$ Entities -> Archetypes
archetypes Entities
w
           in Entities
w
                { archetypes = arches,
                  entities = Map.insert eId aId (entities w),
                  components = components'
                }

-- | Spawn a `DynamicBundle` with a specified `ArchetypeID`.

--

-- @since 0.9

spawnWithArchetypeId ::
  EntityID ->
  ArchetypeID ->
  DynamicBundle ->
  Entities ->
  Entities
spawnWithArchetypeId :: EntityID -> ArchetypeID -> DynamicBundle -> Entities -> Entities
spawnWithArchetypeId EntityID
e ArchetypeID
aId DynamicBundle
b Entities
w =
  let f :: Node -> Node
f Node
n = Node
n {nodeArchetype = runDynamicBundle b e ((nodeArchetype n) {A.entities = Set.insert e . A.entities $ nodeArchetype n})}
   in Entities
w
        { archetypes = (archetypes w) {AS.nodes = Map.adjust f aId (AS.nodes $ archetypes w)},
          entities = Map.insert e aId (entities w)
        }

-- | Insert a component into an entity.

--

-- @since 0.9

insert :: EntityID -> Bundle -> Entities -> Entities
insert :: EntityID -> Bundle -> Entities -> Entities
insert EntityID
e Bundle
b Entities
w =
  let !(Set ComponentID
cIds, Components
components', DynamicBundle
dynB) = Bundle
-> Components -> (Set ComponentID, Components, DynamicBundle)
unBundle Bundle
b (Entities -> Components
components Entities
w)
   in EntityID
-> Set ComponentID -> DynamicBundle -> Entities -> Entities
insertDyn EntityID
e Set ComponentID
cIds DynamicBundle
dynB Entities
w {components = components'}

-- | Insert a component into an entity with its `ComponentID`.

--

-- @since 0.9

insertDyn :: EntityID -> Set ComponentID -> DynamicBundle -> Entities -> Entities
insertDyn :: EntityID
-> Set ComponentID -> DynamicBundle -> Entities -> Entities
insertDyn EntityID
e Set ComponentID
cIds DynamicBundle
b Entities
w = case EntityID -> Map EntityID ArchetypeID -> Maybe ArchetypeID
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup EntityID
e (Map EntityID ArchetypeID -> Maybe ArchetypeID)
-> Map EntityID ArchetypeID -> Maybe ArchetypeID
forall a b. (a -> b) -> a -> b
$ Entities -> Map EntityID ArchetypeID
entities Entities
w of
  Just ArchetypeID
aId ->
    let (Maybe ArchetypeID
maybeNextAId, Archetypes
arches) = EntityID
-> ArchetypeID
-> Set ComponentID
-> DynamicBundle
-> Archetypes
-> (Maybe ArchetypeID, Archetypes)
AS.insert EntityID
e ArchetypeID
aId Set ComponentID
cIds DynamicBundle
b (Archetypes -> (Maybe ArchetypeID, Archetypes))
-> Archetypes -> (Maybe ArchetypeID, Archetypes)
forall a b. (a -> b) -> a -> b
$ Entities -> Archetypes
archetypes Entities
w
        es :: Map EntityID ArchetypeID
es = case Maybe ArchetypeID
maybeNextAId of
          Just ArchetypeID
nextAId -> EntityID
-> ArchetypeID
-> Map EntityID ArchetypeID
-> Map EntityID ArchetypeID
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert EntityID
e ArchetypeID
nextAId (Map EntityID ArchetypeID -> Map EntityID ArchetypeID)
-> Map EntityID ArchetypeID -> Map EntityID ArchetypeID
forall a b. (a -> b) -> a -> b
$ Entities -> Map EntityID ArchetypeID
entities Entities
w
          Maybe ArchetypeID
Nothing -> Entities -> Map EntityID ArchetypeID
entities Entities
w
     in Entities
w {archetypes = arches, entities = es}
  Maybe ArchetypeID
Nothing -> case Set ComponentID -> Archetypes -> Maybe ArchetypeID
AS.lookupArchetypeId Set ComponentID
cIds (Archetypes -> Maybe ArchetypeID)
-> Archetypes -> Maybe ArchetypeID
forall a b. (a -> b) -> a -> b
$ Entities -> Archetypes
archetypes Entities
w of
    Just ArchetypeID
aId -> EntityID -> ArchetypeID -> DynamicBundle -> Entities -> Entities
spawnWithArchetypeId EntityID
e ArchetypeID
aId DynamicBundle
b Entities
w
    Maybe ArchetypeID
Nothing ->
      let arch :: Archetype
arch = DynamicBundle -> EntityID -> Archetype -> Archetype
runDynamicBundle DynamicBundle
b EntityID
e (Archetype -> Archetype) -> Archetype -> Archetype
forall a b. (a -> b) -> a -> b
$ EntityID -> Archetype
A.singleton EntityID
e
          node :: Node
node = Node {nodeComponentIds :: Set ComponentID
nodeComponentIds = Set ComponentID
cIds, nodeArchetype :: Archetype
nodeArchetype = Archetype
arch}
          (ArchetypeID
aId, Archetypes
arches) = Set ComponentID -> Node -> Archetypes -> (ArchetypeID, Archetypes)
AS.insertArchetype Set ComponentID
cIds Node
node (Archetypes -> (ArchetypeID, Archetypes))
-> Archetypes -> (ArchetypeID, Archetypes)
forall a b. (a -> b) -> a -> b
$ Entities -> Archetypes
archetypes Entities
w
       in Entities
w {archetypes = arches, entities = Map.insert e aId (entities w)}

-- | Lookup a component in an entity.

--

-- @since 0.9

lookup :: forall a. (Component a) => EntityID -> Entities -> Maybe a
lookup :: forall a. Component a => EntityID -> Entities -> Maybe a
lookup EntityID
e Entities
w = do
  !ComponentID
cId <- forall a. Typeable a => Components -> Maybe ComponentID
CS.lookup @a (Components -> Maybe ComponentID)
-> Components -> Maybe ComponentID
forall a b. (a -> b) -> a -> b
$ Entities -> Components
components Entities
w
  !ArchetypeID
aId <- EntityID -> Map EntityID ArchetypeID -> Maybe ArchetypeID
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup EntityID
e (Map EntityID ArchetypeID -> Maybe ArchetypeID)
-> Map EntityID ArchetypeID -> Maybe ArchetypeID
forall a b. (a -> b) -> a -> b
$ Entities -> Map EntityID ArchetypeID
entities Entities
w
  !Node
node <- ArchetypeID -> Archetypes -> Maybe Node
AS.lookup ArchetypeID
aId (Archetypes -> Maybe Node) -> Archetypes -> Maybe Node
forall a b. (a -> b) -> a -> b
$ Entities -> Archetypes
archetypes Entities
w
  EntityID -> ComponentID -> Archetype -> Maybe a
forall a.
Component a =>
EntityID -> ComponentID -> Archetype -> Maybe a
A.lookupComponent EntityID
e ComponentID
cId (Archetype -> Maybe a) -> Archetype -> Maybe a
forall a b. (a -> b) -> a -> b
$ Node -> Archetype
nodeArchetype Node
node

-- | Insert a component into an entity.

--

-- @since 0.9

remove :: forall a. (Component a) => EntityID -> Entities -> (Maybe a, Entities)
remove :: forall a.
Component a =>
EntityID -> Entities -> (Maybe a, Entities)
remove EntityID
e Entities
w =
  let !(ComponentID
cId, Components
components') = forall a. Component a => Components -> (ComponentID, Components)
CS.insert @a (Entities -> Components
components Entities
w)
   in forall a.
Component a =>
EntityID -> ComponentID -> Entities -> (Maybe a, Entities)
removeWithId @a EntityID
e ComponentID
cId Entities
w {components = components'}

-- | Remove a component from an entity with its `ComponentID`.

--

-- @since 0.9

removeWithId :: forall a. (Component a) => EntityID -> ComponentID -> Entities -> (Maybe a, Entities)
removeWithId :: forall a.
Component a =>
EntityID -> ComponentID -> Entities -> (Maybe a, Entities)
removeWithId EntityID
e ComponentID
cId Entities
w = case EntityID -> Map EntityID ArchetypeID -> Maybe ArchetypeID
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup EntityID
e (Entities -> Map EntityID ArchetypeID
entities Entities
w) of
  Just ArchetypeID
aId ->
    let (Maybe (a, ArchetypeID)
res, Archetypes
as) = forall a.
Component a =>
EntityID
-> ArchetypeID
-> ComponentID
-> Archetypes
-> (Maybe (a, ArchetypeID), Archetypes)
AS.remove @a EntityID
e ArchetypeID
aId ComponentID
cId (Archetypes -> (Maybe (a, ArchetypeID), Archetypes))
-> Archetypes -> (Maybe (a, ArchetypeID), Archetypes)
forall a b. (a -> b) -> a -> b
$ Entities -> Archetypes
archetypes Entities
w
        (Maybe a
maybeA, Map EntityID ArchetypeID
es) = case Maybe (a, ArchetypeID)
res of
          Just (a
a, ArchetypeID
nextAId) -> (a -> Maybe a
forall a. a -> Maybe a
Just a
a, EntityID
-> ArchetypeID
-> Map EntityID ArchetypeID
-> Map EntityID ArchetypeID
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert EntityID
e ArchetypeID
nextAId (Entities -> Map EntityID ArchetypeID
entities Entities
w))
          Maybe (a, ArchetypeID)
Nothing -> (Maybe a
forall a. Maybe a
Nothing, Entities -> Map EntityID ArchetypeID
entities Entities
w)
     in (Maybe a
maybeA, Entities
w {archetypes = as, entities = es})
  Maybe ArchetypeID
Nothing -> (Maybe a
forall a. Maybe a
Nothing, Entities
w)

-- | Despawn an entity, returning its components.

--

-- @since 0.9

despawn :: EntityID -> Entities -> (IntMap Dynamic, Entities)
despawn :: EntityID -> Entities -> (IntMap Dynamic, Entities)
despawn EntityID
e Entities
w =
  let res :: Maybe (ArchetypeID, Node)
res = do
        !ArchetypeID
aId <- EntityID -> Map EntityID ArchetypeID -> Maybe ArchetypeID
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup EntityID
e (Map EntityID ArchetypeID -> Maybe ArchetypeID)
-> Map EntityID ArchetypeID -> Maybe ArchetypeID
forall a b. (a -> b) -> a -> b
$ Entities -> Map EntityID ArchetypeID
entities Entities
w
        !Node
node <- ArchetypeID -> Archetypes -> Maybe Node
AS.lookup ArchetypeID
aId (Archetypes -> Maybe Node) -> Archetypes -> Maybe Node
forall a b. (a -> b) -> a -> b
$ Entities -> Archetypes
archetypes Entities
w
        (ArchetypeID, Node) -> Maybe (ArchetypeID, Node)
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (ArchetypeID
aId, Node
node)
   in case Maybe (ArchetypeID, Node)
res of
        Just (ArchetypeID
aId, Node
node) ->
          let !(IntMap Dynamic
dynAcc, Archetype
arch') = EntityID -> Archetype -> (IntMap Dynamic, Archetype)
A.remove EntityID
e (Node -> Archetype
nodeArchetype Node
node)
           in ( IntMap Dynamic
dynAcc,
                Entities
w
                  { archetypes = (archetypes w) {AS.nodes = Map.insert aId node {nodeArchetype = arch'} (AS.nodes $ archetypes w)},
                    entities = Map.delete e (entities w)
                  }
              )
        Maybe (ArchetypeID, Node)
Nothing -> (IntMap Dynamic
forall a. IntMap a
IntMap.empty, Entities
w)