{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}

module Aztecs.ECS.World.Archetype
  ( Archetype (..),
    empty,
    lookupComponent,
    lookupStorage,
    member,
    remove,
    removeStorages,
    insertComponent,
    insertAscList,
  )
where

import Aztecs.ECS.Component
import Aztecs.ECS.Entity
import qualified Aztecs.ECS.World.Storage as S
import Aztecs.ECS.World.Storage.Dynamic
import qualified Aztecs.ECS.World.Storage.Dynamic as S
import Control.DeepSeq
import Data.Dynamic
import Data.Foldable
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Set (Set)
import qualified Data.Set as Set
import GHC.Generics

data Archetype = Archetype
  { Archetype -> Map ComponentID DynamicStorage
storages :: !(Map ComponentID DynamicStorage),
    Archetype -> Set EntityID
entities :: !(Set EntityID)
  }
  deriving (Int -> Archetype -> ShowS
[Archetype] -> ShowS
Archetype -> String
(Int -> Archetype -> ShowS)
-> (Archetype -> String)
-> ([Archetype] -> ShowS)
-> Show Archetype
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Archetype -> ShowS
showsPrec :: Int -> Archetype -> ShowS
$cshow :: Archetype -> String
show :: Archetype -> String
$cshowList :: [Archetype] -> ShowS
showList :: [Archetype] -> ShowS
Show, (forall x. Archetype -> Rep Archetype x)
-> (forall x. Rep Archetype x -> Archetype) -> Generic Archetype
forall x. Rep Archetype x -> Archetype
forall x. Archetype -> Rep Archetype x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. Archetype -> Rep Archetype x
from :: forall x. Archetype -> Rep Archetype x
$cto :: forall x. Rep Archetype x -> Archetype
to :: forall x. Rep Archetype x -> Archetype
Generic, Archetype -> ()
(Archetype -> ()) -> NFData Archetype
forall a. (a -> ()) -> NFData a
$crnf :: Archetype -> ()
rnf :: Archetype -> ()
NFData)

empty :: Archetype
empty :: Archetype
empty = Archetype {storages :: Map ComponentID DynamicStorage
storages = Map ComponentID DynamicStorage
forall k a. Map k a
Map.empty, entities :: Set EntityID
entities = Set EntityID
forall a. Set a
Set.empty}

lookupStorage :: (Component a) => ComponentID -> Archetype -> Maybe (StorageT a)
lookupStorage :: forall a.
Component a =>
ComponentID -> Archetype -> Maybe (StorageT a)
lookupStorage ComponentID
cId Archetype
w = do
  DynamicStorage
dynS <- ComponentID
-> Map ComponentID DynamicStorage -> Maybe DynamicStorage
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup ComponentID
cId (Archetype -> Map ComponentID DynamicStorage
storages Archetype
w)
  Dynamic -> Maybe (StorageT a)
forall a. Typeable a => Dynamic -> Maybe a
fromDynamic (DynamicStorage -> Dynamic
storageDyn DynamicStorage
dynS)

lookupComponents :: forall a. (Component a) => ComponentID -> Archetype -> Map EntityID a
lookupComponents :: forall a. Component a => ComponentID -> Archetype -> Map EntityID a
lookupComponents ComponentID
cId Archetype
arch = case forall a.
Component a =>
ComponentID -> Archetype -> Maybe (StorageT a)
lookupStorage @a ComponentID
cId Archetype
arch of
  Just StorageT a
s -> [(EntityID, a)] -> Map EntityID a
forall k a. Eq k => [(k, a)] -> Map k a
Map.fromAscList ([(EntityID, a)] -> Map EntityID a)
-> ([a] -> [(EntityID, a)]) -> [a] -> Map EntityID a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [EntityID] -> [a] -> [(EntityID, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Set EntityID -> [EntityID]
forall a. Set a -> [a]
Set.toList (Set EntityID -> [EntityID]) -> Set EntityID -> [EntityID]
forall a b. (a -> b) -> a -> b
$ Archetype -> Set EntityID
entities Archetype
arch) ([a] -> Map EntityID a) -> [a] -> Map EntityID a
forall a b. (a -> b) -> a -> b
$ StorageT a -> [a]
forall a s. Storage a s => s -> [a]
S.toAscList StorageT a
s
  Maybe (StorageT a)
Nothing -> Map EntityID a
forall k a. Map k a
Map.empty

insertComponent :: forall a. (Component a) => EntityID -> ComponentID -> a -> Archetype -> Archetype
insertComponent :: forall a.
Component a =>
EntityID -> ComponentID -> a -> Archetype -> Archetype
insertComponent EntityID
e ComponentID
cId a
c Archetype
arch =
  let !storage :: StorageT a
storage =
        forall a s. Storage a s => [a] -> s
S.fromAscList @a @(StorageT a) ([a] -> StorageT a)
-> (Map EntityID a -> [a]) -> Map EntityID a -> StorageT a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map EntityID a -> [a]
forall k a. Map k a -> [a]
Map.elems (Map EntityID a -> [a])
-> (Map EntityID a -> Map EntityID a) -> Map EntityID a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. EntityID -> a -> Map EntityID a -> Map EntityID a
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert EntityID
e a
c (Map EntityID a -> StorageT a) -> Map EntityID a -> StorageT a
forall a b. (a -> b) -> a -> b
$ ComponentID -> Archetype -> Map EntityID a
forall a. Component a => ComponentID -> Archetype -> Map EntityID a
lookupComponents ComponentID
cId Archetype
arch
   in Archetype
arch {storages = Map.insert cId (dynStorage @a storage) (storages arch)}

member :: ComponentID -> Archetype -> Bool
member :: ComponentID -> Archetype -> Bool
member ComponentID
cId Archetype
arch = ComponentID -> Map ComponentID DynamicStorage -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.member ComponentID
cId (Archetype -> Map ComponentID DynamicStorage
storages Archetype
arch)

lookupComponent :: forall a. (Component a) => EntityID -> ComponentID -> Archetype -> Maybe a
lookupComponent :: forall a.
Component a =>
EntityID -> ComponentID -> Archetype -> Maybe a
lookupComponent EntityID
e ComponentID
cId Archetype
w = ComponentID -> Archetype -> Map EntityID a
forall a. Component a => ComponentID -> Archetype -> Map EntityID a
lookupComponents ComponentID
cId Archetype
w Map EntityID a -> EntityID -> Maybe a
forall k a. Ord k => Map k a -> k -> Maybe a
Map.!? EntityID
e

-- | Insert a list of components into the archetype, sorted in ascending order by their `EntityID`.
insertAscList :: forall a. (Component a) => ComponentID -> [a] -> Archetype -> Archetype
insertAscList :: forall a.
Component a =>
ComponentID -> [a] -> Archetype -> Archetype
insertAscList ComponentID
cId [a]
as Archetype
arch =
  let !storage :: DynamicStorage
storage = forall a s. Storage a s => s -> DynamicStorage
dynStorage @a (StorageT a -> DynamicStorage) -> StorageT a -> DynamicStorage
forall a b. (a -> b) -> a -> b
$ forall a s. Storage a s => [a] -> s
S.fromAscList @a @(StorageT a) [a]
as
   in Archetype
arch {storages = Map.insert cId storage (storages arch)}

remove :: EntityID -> Archetype -> (Map ComponentID Dynamic, Archetype)
remove :: EntityID -> Archetype -> (Map ComponentID Dynamic, Archetype)
remove EntityID
e Archetype
arch =
  ((Map ComponentID Dynamic, Archetype)
 -> (ComponentID, DynamicStorage)
 -> (Map ComponentID Dynamic, Archetype))
-> (Map ComponentID Dynamic, Archetype)
-> [(ComponentID, DynamicStorage)]
-> (Map ComponentID Dynamic, Archetype)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
    ( \(Map ComponentID Dynamic
dynAcc, Archetype
archAcc) (ComponentID
cId, DynamicStorage
dynS) ->
        let cs :: Map EntityID Dynamic
cs = [(EntityID, Dynamic)] -> Map EntityID Dynamic
forall k a. Eq k => [(k, a)] -> Map k a
Map.fromAscList ([(EntityID, Dynamic)] -> Map EntityID Dynamic)
-> ([Dynamic] -> [(EntityID, Dynamic)])
-> [Dynamic]
-> Map EntityID Dynamic
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [EntityID] -> [Dynamic] -> [(EntityID, Dynamic)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Set EntityID -> [EntityID]
forall a. Set a -> [a]
Set.toList (Set EntityID -> [EntityID]) -> Set EntityID -> [EntityID]
forall a b. (a -> b) -> a -> b
$ Archetype -> Set EntityID
entities Archetype
arch) ([Dynamic] -> Map EntityID Dynamic)
-> [Dynamic] -> Map EntityID Dynamic
forall a b. (a -> b) -> a -> b
$ DynamicStorage -> [Dynamic]
toAscListDyn DynamicStorage
dynS
            !(Maybe Dynamic
dynA, Map EntityID Dynamic
cs') = (EntityID -> Dynamic -> Maybe Dynamic)
-> EntityID
-> Map EntityID Dynamic
-> (Maybe Dynamic, Map EntityID Dynamic)
forall k a.
Ord k =>
(k -> a -> Maybe a) -> k -> Map k a -> (Maybe a, Map k a)
Map.updateLookupWithKey (\EntityID
_ Dynamic
_ -> Maybe Dynamic
forall a. Maybe a
Nothing) EntityID
e Map EntityID Dynamic
cs
            dynS' :: DynamicStorage
dynS' = [Dynamic] -> DynamicStorage -> DynamicStorage
S.fromAscListDyn (Map EntityID Dynamic -> [Dynamic]
forall k a. Map k a -> [a]
Map.elems Map EntityID Dynamic
cs') DynamicStorage
dynS
            !dynAcc' :: Map ComponentID Dynamic
dynAcc' = case Maybe Dynamic
dynA of
              Just Dynamic
d -> ComponentID
-> Dynamic -> Map ComponentID Dynamic -> Map ComponentID Dynamic
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert ComponentID
cId Dynamic
d Map ComponentID Dynamic
dynAcc
              Maybe Dynamic
Nothing -> Map ComponentID Dynamic
dynAcc
         in (Map ComponentID Dynamic
dynAcc', Archetype
archAcc {storages = Map.insert cId dynS' $ storages archAcc})
    )
    (Map ComponentID Dynamic
forall k a. Map k a
Map.empty, Archetype
arch)
    (Map ComponentID DynamicStorage -> [(ComponentID, DynamicStorage)]
forall k a. Map k a -> [(k, a)]
Map.toList (Map ComponentID DynamicStorage -> [(ComponentID, DynamicStorage)])
-> Map ComponentID DynamicStorage
-> [(ComponentID, DynamicStorage)]
forall a b. (a -> b) -> a -> b
$ Archetype -> Map ComponentID DynamicStorage
storages Archetype
arch)

removeStorages :: EntityID -> Archetype -> (Map ComponentID DynamicStorage, Archetype)
removeStorages :: EntityID
-> Archetype -> (Map ComponentID DynamicStorage, Archetype)
removeStorages EntityID
e Archetype
arch =
  ((Map ComponentID DynamicStorage, Archetype)
 -> (ComponentID, DynamicStorage)
 -> (Map ComponentID DynamicStorage, Archetype))
-> (Map ComponentID DynamicStorage, Archetype)
-> [(ComponentID, DynamicStorage)]
-> (Map ComponentID DynamicStorage, Archetype)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
    ( \(Map ComponentID DynamicStorage
dynAcc, Archetype
archAcc) (ComponentID
cId, DynamicStorage
dynS) ->
        let cs :: Map EntityID Dynamic
cs = [(EntityID, Dynamic)] -> Map EntityID Dynamic
forall k a. Eq k => [(k, a)] -> Map k a
Map.fromAscList ([(EntityID, Dynamic)] -> Map EntityID Dynamic)
-> ([Dynamic] -> [(EntityID, Dynamic)])
-> [Dynamic]
-> Map EntityID Dynamic
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [EntityID] -> [Dynamic] -> [(EntityID, Dynamic)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Set EntityID -> [EntityID]
forall a. Set a -> [a]
Set.toList (Set EntityID -> [EntityID]) -> Set EntityID -> [EntityID]
forall a b. (a -> b) -> a -> b
$ Archetype -> Set EntityID
entities Archetype
arch) ([Dynamic] -> Map EntityID Dynamic)
-> [Dynamic] -> Map EntityID Dynamic
forall a b. (a -> b) -> a -> b
$ DynamicStorage -> [Dynamic]
toAscListDyn DynamicStorage
dynS
            !(Maybe Dynamic
dynA, Map EntityID Dynamic
cs') = (EntityID -> Dynamic -> Maybe Dynamic)
-> EntityID
-> Map EntityID Dynamic
-> (Maybe Dynamic, Map EntityID Dynamic)
forall k a.
Ord k =>
(k -> a -> Maybe a) -> k -> Map k a -> (Maybe a, Map k a)
Map.updateLookupWithKey (\EntityID
_ Dynamic
_ -> Maybe Dynamic
forall a. Maybe a
Nothing) EntityID
e Map EntityID Dynamic
cs
            dynS' :: DynamicStorage
dynS' = [Dynamic] -> DynamicStorage -> DynamicStorage
S.fromAscListDyn (Map EntityID Dynamic -> [Dynamic]
forall k a. Map k a -> [a]
Map.elems Map EntityID Dynamic
cs') DynamicStorage
dynS
            !dynAcc' :: Map ComponentID DynamicStorage
dynAcc' = case Maybe Dynamic
dynA of
              Just Dynamic
d -> ComponentID
-> DynamicStorage
-> Map ComponentID DynamicStorage
-> Map ComponentID DynamicStorage
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert ComponentID
cId (Dynamic -> DynamicStorage -> DynamicStorage
S.singletonDyn Dynamic
d DynamicStorage
dynS') Map ComponentID DynamicStorage
dynAcc
              Maybe Dynamic
Nothing -> Map ComponentID DynamicStorage
dynAcc
         in (Map ComponentID DynamicStorage
dynAcc', Archetype
archAcc {storages = Map.insert cId dynS' $ storages archAcc})
    )
    (Map ComponentID DynamicStorage
forall k a. Map k a
Map.empty, Archetype
arch)
    (Map ComponentID DynamicStorage -> [(ComponentID, DynamicStorage)]
forall k a. Map k a -> [(k, a)]
Map.toList (Map ComponentID DynamicStorage -> [(ComponentID, DynamicStorage)])
-> Map ComponentID DynamicStorage
-> [(ComponentID, DynamicStorage)]
forall a b. (a -> b) -> a -> b
$ Archetype -> Map ComponentID DynamicStorage
storages Archetype
arch)