UnityECS/LocalPackages/com.nimin.tinyecs/Runtime/EntityCollection.cs

using System;
using Unity.Collections;
using System.Diagnostics;
using Unity.Collections.LowLevel.Unsafe;

namespace GameCore.TinyECS
{
    public struct EntityCollection<TEnum> : IDisposable
        where TEnum : unmanaged, IConvertible
    {
        private struct Buffer
        {
            public int index;
            public int generation;
        }

        public readonly TEnum entityType;
        
        private NativeList<Buffer> entityBuffer;
        private NativeList<int> entityIndexToDataIndex;

        public EntityCollection(TEnum type, int capacity, AllocatorManager.AllocatorHandle allocator)
        {
            entityType = type;
            entityBuffer = new (capacity, allocator);
            entityIndexToDataIndex = new (capacity, allocator);
        }

        public void Dispose()
        {
            entityBuffer.Dispose();
            entityIndexToDataIndex.Dispose();
        }

        public int Count => entityBuffer.Length;

        public Entity<TEnum> this[int index]
        {
            get
            {
                var buf = entityBuffer[index];
                return new(entityType, buf.index, buf.generation);
            }
        }

        public int Add(Entity<TEnum> entity)
        {
            CheckEntityType(entity);

            var dataIndex = entityBuffer.Length;
            var oldEntityCount = entityIndexToDataIndex.Length;
            if (entity.index >= oldEntityCount)
            {
                entityIndexToDataIndex.Resize(entity.index + 1, NativeArrayOptions.UninitializedMemory);
                for (int i = oldEntityCount, e = entity.index - 1; i < e; ++i)
                {
                    entityIndexToDataIndex[i] = -1;
                }
            }

            entityIndexToDataIndex[entity.index] = dataIndex;
            entityBuffer.Add(new Buffer{index = entity.index, generation = entity.generation });
            return dataIndex;
        }

        public int GetIndex(Entity<TEnum> entity)
        {
            CheckEntityType(entity);
            
            if (entity.index >= entityIndexToDataIndex.Length)
            {
                throw new Exception($"Invalid entity: {entity}");
            }

            var index = entityIndexToDataIndex[entity.index];
            var buf = entityBuffer[index];
            if (buf.index != entity.index || buf.generation != entity.generation)
            {
                throw new Exception($"Invalid entity: {entity}");
            }

            return index;
        }

        public bool TryGetIndex(Entity<TEnum> entity, out int index)
        {
            CheckEntityType(entity);

            index = 0;
            if (entity.index >= entityIndexToDataIndex.Length)
            {
                return false;
            }

            index = entityIndexToDataIndex[entity.index];
            if (index < 0 || index >= entityBuffer.Length)
            {
                return false;
            }

            var buf = entityBuffer[index];
            return buf.index == entity.index && buf.generation == entity.generation;
        }

        public bool Remove(Entity<TEnum> entity, out int index, out bool removeFromBack)
        {
            CheckEntityType(entity);

            if (!TryGetIndex(entity, out index))
            {
                removeFromBack = false;
                return false;
            }

            removeFromBack = RemoveInternal(entity, index);
            return true;
        }

        private bool RemoveInternal(Entity<TEnum> entity, int index)
        {
            entityIndexToDataIndex[entity.index] = -1;

            var backIndex = entityBuffer.Length - 1;
            if (index == backIndex)
            {
                entityBuffer.Resize(backIndex, NativeArrayOptions.UninitializedMemory);
                return true;
            }

            var backBuf = entityBuffer[backIndex];
            entityBuffer.RemoveAtSwapBack(index);
            entityIndexToDataIndex[backBuf.index] = index;

            return false;
        }

        [Conditional("DEBUG")]
        void CheckEntityType(Entity<TEnum> entity)
        {
            if (!UnsafeUtility.EnumEquals(entityType, entity.type))
                throw new Exception($"Wrong entity type: {entity.type}");
        }

        public void Clear()
        {
            entityBuffer.Clear();
            entityIndexToDataIndex.Clear();
        }

    }
}