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

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

namespace GameCore.TinyECS
{
    public partial struct EntityDataBlockMeta
    {
        public FixedList64Bytes<byte> componentIdToIndex;
        public FixedList128Bytes<short> componentIdToOffset;
        public int elementSize;
        public int elementAlign;
    }

    public unsafe struct EntityDataBlockChain : IDisposable
    {
        public readonly EntityDataBlockMeta blockMeta;
        private FixedList512Bytes<IntPtr> blockPtrList;
        private int maxElementCountPerBlock;
        private int elementCount;

        public int BlockCount => blockPtrList.Length;
        public int BlockCapacity => blockPtrList.Capacity;
        public int ElementCount => elementCount;
        public int MaxElementCount => MaxElementCountPerBlock * BlockCapacity;
        public int MaxElementCountPerBlock => maxElementCountPerBlock;

        public bool IsFull => elementCount == MaxElementCount;
        public bool IsEmpty => elementCount == 0;

        public EntityDataBlockChain(EntityDataBlockMeta blockMeta, int maxElementCountPerBlock)
        {
            this.blockMeta = blockMeta;
            blockPtrList = default;
            this.maxElementCountPerBlock = maxElementCountPerBlock;
            elementCount = 0;
        }

        public void Dispose()
        {
            for (var i = BlockCount - 1; i >= 0; --i)
            {
                ReleaseBlock(i);
            }
        }

        void AllocBlock()
        {
            var bufferSize = blockMeta.elementSize * maxElementCountPerBlock;
            var ptr = UnsafeUtility.Malloc(bufferSize, blockMeta.elementAlign, Allocator.Persistent);

            blockPtrList.Add((IntPtr)ptr);
        }

        void ReleaseBlock(int blockIndex)
        {
            var pBlock = (void*)blockPtrList[blockIndex];
            if (pBlock != null)
            {
                UnsafeUtility.Free(pBlock, Allocator.Persistent);
                blockPtrList.RemoveAt(blockIndex);
            }
        }

        void* ElementPtrAt(int elementIndex)
        {
            var blockIndex = elementIndex / maxElementCountPerBlock;
            var pBlock = blockPtrList[blockIndex];
            return (byte*)pBlock + elementIndex * blockMeta.elementSize;
        }

        void* ComponentPtrAt<T>(int elementIndex) where T : unmanaged
        {
            var componentId = ComponentId.Get<T>();
            var componentIndex = blockMeta.componentIdToIndex[componentId];
            if (componentIndex > 0)
            {
                var pElement = ElementPtrAt(elementIndex);
                return (byte*)pElement + blockMeta.componentIdToOffset[componentId];
            }

            return null;
        }

        [Conditional("DEBUG")]
        void CheckElementIndex(int index)
        {
            if (index < 0 || index >= elementCount)
                throw new IndexOutOfRangeException($"Invalid index: {index}");
        }

        [Conditional("DEBUG")]
        void CheckBlockIndex(int index)
        {
            if (index < 0 || index > (elementCount / maxElementCountPerBlock))
                throw new IndexOutOfRangeException($"Invalid index: {index}");
        }

        public void Clear()
        {
            elementCount = 0;
        }

        public void Add()
        {
            if (elementCount % maxElementCountPerBlock == 0)
            {
                if (IsFull)
                    throw new OutOfMemoryException("Block chain is full!");

                var blockIndex = elementCount / maxElementCountPerBlock;
                if (blockIndex >= BlockCount)
                    AllocBlock();
            }

            void* ptr = ElementPtrAt(elementCount);
            UnsafeUtility.MemClear(ptr, blockMeta.elementSize);

            elementCount++;
        }

        public void Remove(int index)
        {
            CheckElementIndex(index);
            if (index != elementCount - 1)
            {
                var src = ElementPtrAt(elementCount - 1);
                var dst = ElementPtrAt(index);
                UnsafeUtility.MemCpy(dst, src, blockMeta.elementSize);
            }

            elementCount--;
        }

        public ref T GetData<T>(int index) where T : unmanaged
        {
            CheckElementIndex(index);
            var pComponent = ComponentPtrAt<T>(index);
            if (pComponent == null)
                throw new Exception($"Invalid component type: {typeof(T)}");
            return ref UnsafeUtility.AsRef<T>(pComponent);
        }

        public EntityDataBlockView AsBlockView(int blockIndex)
        {
            CheckBlockIndex(blockIndex);
            var ptr = (void*)blockPtrList[blockIndex];
            return new(ptr, blockMeta, maxElementCountPerBlock * blockIndex, elementCount);
        }
    }

    public unsafe readonly struct EntityDataBlockView
    {
        private readonly void* ptr;
        public readonly EntityDataBlockMeta meta;
        public readonly int elementStartIndex;
        public readonly int elementCount;

        public EntityDataBlockView(void* ptr, EntityDataBlockMeta meta, int elementStartIndex, int elementCount)
        {
            this.ptr = ptr;
            this.meta = meta;
            this.elementStartIndex = elementStartIndex;
            this.elementCount = elementCount;
        }

        void* ElementPtrAt(int elementIndex)
        {
            return (byte*)ptr + elementIndex * meta.elementSize;
        }

        void* ComponentPtrAt<T>(int elementIndex) where T : unmanaged
        {
            var componentId = ComponentId.Get<T>();
            var componentIndex = meta.componentIdToIndex[componentId];
            if (componentIndex > 0)
            {
                var pElement = ElementPtrAt(elementIndex);
                return (byte*)pElement + meta.componentIdToOffset[componentId];
            }

            return null;
        }

        [Conditional("DEBUG")]
        void CheckIndex(int index)
        {
            if (index < elementStartIndex || index >= elementStartIndex + elementCount)
                throw new IndexOutOfRangeException($"Invalid index: {index}");
        }

        public ref T GetData<T>(int index) where T : unmanaged
        {
            var pComponent = ComponentPtrAt<T>(index);
            if (pComponent == null)
                throw new Exception($"Invalid component type: {typeof(T)}");
            return ref UnsafeUtility.AsRef<T>(pComponent);
        }
    }
}