perf!: replace Archetype Map with IntMap

aztecs-hs · Feb 27, 2025 · 349da82 · 349da82
1 parent 8ce8731
commit 349da82
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 28 deletions.
diff --git a/src/Aztecs/ECS/World.hs b/src/Aztecs/ECS/World.hs
@@ -29,6 +29,7 @@ import Aztecs.ECS.World.Entities (Entities)
 import qualified Aztecs.ECS.World.Entities as E
 import Control.DeepSeq
 import Data.Dynamic
+import Data.IntMap (IntMap)
 import Data.Map (Map)
 import GHC.Generics
 import Prelude hiding (lookup)
@@ -89,5 +90,5 @@ removeWithId :: forall a. (Component a) => EntityID -> ComponentID -> World -> (
 removeWithId e cId w = let (a, es) = E.removeWithId e cId (entities w) in (a, w {entities = es})
 
 -- | Despawn an entity, returning its components.
-despawn :: EntityID -> World -> (Map ComponentID Dynamic, World)
+despawn :: EntityID -> World -> (IntMap Dynamic, World)
 despawn e w = let (a, es) = E.despawn e (entities w) in (a, w {entities = es})
diff --git a/src/Aztecs/ECS/World/Archetype.hs b/src/Aztecs/ECS/World/Archetype.hs
@@ -36,26 +36,28 @@ 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.IntMap (IntMap)
+import qualified Data.IntMap.Strict 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
 
 data Archetype = Archetype
-  { storages :: !(Map ComponentID DynamicStorage),
+  { storages :: !(IntMap DynamicStorage),
     entities :: !(Set EntityID)
   }
   deriving (Show, Generic, NFData)
 
 empty :: Archetype
-empty = Archetype {storages = Map.empty, entities = Set.empty}
+empty = Archetype {storages = IntMap.empty, entities = Set.empty}
 
 {-# INLINE lookupStorage #-}
 lookupStorage :: (Component a) => ComponentID -> Archetype -> Maybe (StorageT a)
 lookupStorage cId w = do
-  dynS <- Map.lookup cId $ storages w
+  !dynS <- IntMap.lookup (unComponentId cId) $ storages w
   fromDynamic $ storageDyn dynS
 
 {-# INLINE lookupComponent #-}
@@ -80,52 +82,52 @@ insertComponent :: forall a. (Component a) => EntityID -> ComponentID -> a -> Ar
 insertComponent e cId c arch =
   let !storage =
         S.fromAscList @a @(StorageT a) . Map.elems . Map.insert e c $ lookupComponents cId arch
-   in arch {storages = Map.insert cId (dynStorage @a storage) (storages arch)}
+   in arch {storages = IntMap.insert (unComponentId cId) (dynStorage @a storage) (storages arch)}
 
 member :: ComponentID -> Archetype -> Bool
-member cId arch = Map.member cId (storages arch)
+member cId = IntMap.member (unComponentId cId) . storages
 
 -- | Insert a list of components into the archetype, sorted in ascending order by their `EntityID`.
 {-# INLINE insertAscList #-}
 insertAscList :: forall a. (Component a) => ComponentID -> [a] -> Archetype -> Archetype
 insertAscList cId as arch =
   let !storage = dynStorage @a $ S.fromAscList @a @(StorageT a) as
-   in arch {storages = Map.insert cId storage (storages arch)}
+   in arch {storages = IntMap.insert (unComponentId cId) storage $ storages arch}
 
-remove :: EntityID -> Archetype -> (Map ComponentID Dynamic, Archetype)
+remove :: EntityID -> Archetype -> (IntMap Dynamic, Archetype)
 remove e arch =
   foldl'
     ( \(dynAcc, archAcc) (cId, dynS) ->
         let cs = Map.fromAscList . zip (Set.toList $ entities arch) $ toAscListDyn dynS
             !(dynA, cs') = Map.updateLookupWithKey (\_ _ -> Nothing) e cs
             dynS' = S.fromAscListDyn (Map.elems cs') dynS
             !dynAcc' = case dynA of
-              Just d -> Map.insert cId d dynAcc
+              Just d -> IntMap.insert cId d dynAcc
               Nothing -> dynAcc
-         in (dynAcc', archAcc {storages = Map.insert cId dynS' $ storages archAcc})
+         in (dynAcc', archAcc {storages = IntMap.insert cId dynS' $ storages archAcc})
     )
-    (Map.empty, arch)
-    (Map.toList $ storages arch)
+    (IntMap.empty, arch)
+    (IntMap.toList $ storages arch)
 
-removeStorages :: EntityID -> Archetype -> (Map ComponentID DynamicStorage, Archetype)
+removeStorages :: EntityID -> Archetype -> (IntMap DynamicStorage, Archetype)
 removeStorages e arch =
   foldl'
     ( \(dynAcc, archAcc) (cId, dynS) ->
         let cs = Map.fromAscList . zip (Set.toList $ entities arch) $ toAscListDyn dynS
             !(dynA, cs') = Map.updateLookupWithKey (\_ _ -> Nothing) e cs
             dynS' = S.fromAscListDyn (Map.elems cs') dynS
             !dynAcc' = case dynA of
-              Just d -> Map.insert cId (S.singletonDyn d dynS') dynAcc
+              Just d -> IntMap.insert cId (S.singletonDyn d dynS') dynAcc
               Nothing -> dynAcc
-         in (dynAcc', archAcc {storages = Map.insert cId dynS' $ storages archAcc})
+         in (dynAcc', archAcc {storages = IntMap.insert cId dynS' $ storages archAcc})
     )
-    (Map.empty, arch)
-    (Map.toList $ storages arch)
+    (IntMap.empty, arch)
+    (IntMap.toList $ storages arch)
 
-insertComponents :: EntityID -> Map ComponentID Dynamic -> Archetype -> Archetype
+insertComponents :: EntityID -> IntMap Dynamic -> Archetype -> Archetype
 insertComponents e cs arch =
   let f archAcc (itemCId, dyn) =
-        let storages' = Map.adjust go itemCId (storages archAcc)
+        let storages' = IntMap.adjust go itemCId (storages archAcc)
             go s = fromAscListDyn (Map.elems . Map.insert e dyn . Map.fromAscList . zip (Set.toList $ entities archAcc) $ toAscListDyn s) s
          in archAcc {storages = storages'}
-   in foldl' f arch (Map.toList cs)
+   in foldl' f arch (IntMap.toList cs)
diff --git a/src/Aztecs/ECS/World/Archetypes.hs b/src/Aztecs/ECS/World/Archetypes.hs
@@ -36,6 +36,7 @@ import Aztecs.ECS.World.Storage.Dynamic (fromAscListDyn, toAscListDyn)
 import Control.DeepSeq (NFData (..))
 import Data.Dynamic
 import Data.Foldable (foldl')
+import qualified Data.IntMap.Strict as IntMap
 import Data.Map.Strict (Map)
 import qualified Data.Map.Strict as Map
 import Data.Maybe
@@ -179,18 +180,18 @@ remove e aId cId arches = case lookup aId arches of
     Just nextAId ->
       let !(cs, arch') = A.remove e (nodeArchetype node)
           !arches' = arches {nodes = Map.insert aId node {nodeArchetype = arch'} (nodes arches)}
-          (a, cs') = Map.updateLookupWithKey (\_ _ -> Nothing) cId cs
+          (a, cs') = IntMap.updateLookupWithKey (\_ _ -> Nothing) (unComponentId cId) cs
           go' archAcc (itemCId, dyn) =
             let adjustStorage s = fromAscListDyn (Map.elems . Map.insert e dyn . Map.fromAscList . zip (Set.toList $ entities archAcc) $ toAscListDyn s) s
-             in archAcc {storages = Map.adjust adjustStorage itemCId (storages archAcc)}
+             in archAcc {storages = IntMap.adjust adjustStorage itemCId (storages archAcc)}
           go nextNode =
-            nextNode {nodeArchetype = foldl' go' (nodeArchetype nextNode) (Map.toList cs')}
+            nextNode {nodeArchetype = foldl' go' (nodeArchetype nextNode) (IntMap.toList cs')}
        in ( (,nextAId) <$> (a >>= fromDynamic),
             arches' {nodes = Map.adjust go nextAId (nodes arches')}
           )
     Nothing ->
       let !(cs, arch') = A.removeStorages e (nodeArchetype node)
-          (a, cs') = Map.updateLookupWithKey (\_ _ -> Nothing) cId cs
+          (a, cs') = IntMap.updateLookupWithKey (\_ _ -> Nothing) (unComponentId cId) cs
           !n =
             Node
               { nodeComponentIds = Set.insert cId (nodeComponentIds node),

diff --git a/src/Aztecs/ECS/World/Entities.hs b/src/Aztecs/ECS/World/Entities.hs
@@ -32,6 +32,8 @@ 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
@@ -199,7 +201,7 @@ removeWithId e cId w = case Map.lookup e (entities w) of
   Nothing -> (Nothing, w)
 
 -- | Despawn an entity, returning its components.
-despawn :: EntityID -> Entities -> (Map ComponentID Dynamic, Entities)
+despawn :: EntityID -> Entities -> (IntMap Dynamic, Entities)
 despawn e w =
   let res = do
         !aId <- Map.lookup e $ entities w
@@ -214,4 +216,4 @@ despawn e w =
                     entities = Map.delete e (entities w)
                   }
               )
-        Nothing -> (Map.empty, w)
+        Nothing -> (IntMap.empty, w)