From efed7dcdbd6df5bb570e88c05364410bb87aa2c4 Mon Sep 17 00:00:00 2001
From: Mikhail <99481254+DCFApixels@users.noreply.github.com>
Date: Fri, 23 Feb 2024 20:44:11 +0800
Subject: [PATCH] update PoolMediator

---
 src/EcsWorld.pools.cs | 62 ++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 61 insertions(+), 1 deletion(-)

diff --git a/src/EcsWorld.pools.cs b/src/EcsWorld.pools.cs
index 376eeef..4ea01f0 100644
--- a/src/EcsWorld.pools.cs
+++ b/src/EcsWorld.pools.cs
@@ -202,6 +202,50 @@ namespace DCFApixels.DragonECS
             if (count < 0) Throw.World_InvalidIncrementComponentsBalance();
 #endif
         }
+
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private bool TryRegisterEntityComponent(int entityID, int componentTypeID, EcsMaskChunck maskBit)
+        {
+            ref int chunk = ref _entityComponentMasks[entityID * _entityComponentMaskLength + maskBit.chankIndex];
+            int newChunk = chunk | maskBit.mask;
+            if (chunk != newChunk)
+            {
+                chunk = newChunk;
+                _poolComponentCounts[componentTypeID]++;
+                _componentCounts[entityID]++;
+                return true;
+            }
+            return false;
+        }
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private bool TryUnregisterEntityComponent(int entityID, int componentTypeID, EcsMaskChunck maskBit)
+        {
+            ref int chunk = ref _entityComponentMasks[entityID * _entityComponentMaskLength + maskBit.chankIndex];
+            int newChunk = chunk & ~maskBit.mask;
+            if (chunk != newChunk)
+            {
+                _poolComponentCounts[componentTypeID]--;
+                var count = --_componentCounts[entityID];
+                chunk = newChunk;
+
+                if (count == 0 && IsUsed(entityID))
+                {
+                    DelEntity(entityID);
+                }
+#if (DEBUG && !DISABLE_DEBUG) || ENABLE_DRAGONECS_ASSERT_CHEKS
+                if (count < 0) Throw.World_InvalidIncrementComponentsBalance();
+#endif
+                return true;
+            }
+            return false;
+        }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private int GetPoolComponentCount(int componentTypeID)
+        {
+            return _poolComponentCounts[componentTypeID];
+        }
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         private bool HasEntityComponent(int entityID, EcsMaskChunck maskBit)
         {
@@ -217,7 +261,7 @@ namespace DCFApixels.DragonECS
             {
                 if (world == null || world._poolsMediator._world != null)
                 {
-                    throw new MethodAccessException();
+                    throw new InvalidOperationException();
                 }
                 _world = world;
             }
@@ -232,6 +276,22 @@ namespace DCFApixels.DragonECS
                 _world.UnregisterEntityComponent(entityID, componentTypeID, maskBit);
             }
             [MethodImpl(MethodImplOptions.AggressiveInlining)]
+            public bool TryRegisterComponent(int entityID, int componentTypeID, EcsMaskChunck maskBit)
+            {
+                return _world.TryRegisterEntityComponent(entityID, componentTypeID, maskBit);
+            }
+            [MethodImpl(MethodImplOptions.AggressiveInlining)]
+            public bool TryUnregisterComponent(int entityID, int componentTypeID, EcsMaskChunck maskBit)
+            {
+                return _world.TryUnregisterEntityComponent(entityID, componentTypeID, maskBit);
+            }
+  
+            [MethodImpl(MethodImplOptions.AggressiveInlining)]
+            public int GetComponentCount(int componentTypeID)
+            {
+                return _world.GetPoolComponentCount(componentTypeID);
+            }
+            [MethodImpl(MethodImplOptions.AggressiveInlining)]
             public bool HasComponent(int entityID, EcsMaskChunck maskBit)
             {
                 return _world.HasEntityComponent(entityID, maskBit);