GhostEngine/Ghost.Entities/Archetype.cs

using Ghost.Core;
using Misaki.HighPerformance.LowLevel.Buffer;
using Misaki.HighPerformance.LowLevel.Collections;
using Misaki.HighPerformance.LowLevel.Utilities;
using System.Diagnostics;
using System.Runtime.CompilerServices;

namespace Ghost.Entities;

internal unsafe sealed class ChunkDebugView
{
    [DebuggerDisplay("{Name,nq}: {Data}")]
    internal class ComponentArrayView
    {
        public string Name { get; }
        [DebuggerBrowsable(DebuggerBrowsableState.RootHidden)]
        public object Data { get; }

        public ComponentArrayView(string name, object data)
        {
            Name = name;
            Data = data;
        }
    }

    private Chunk _chunk;

    public ChunkDebugView(Chunk chunk)
    {
        _chunk = chunk;
    }

    [DebuggerBrowsable(DebuggerBrowsableState.RootHidden)]
    public object[] Items => GetItems(in _chunk);

    private static T[] ReadComponentArray<T>(long pData, int offsetInChunk, int count)
        where T : unmanaged
    {
        var result = new T[count];
        unsafe
        {
            var basePtr = (byte*)pData + offsetInChunk;
            var span = new Span<T>(basePtr, count);
            span.CopyTo(result);
        }

        return result;
    }

    private static object[] GetItems(ref readonly Chunk chunk)
    {
#if !(DEBUG || GHOST_EDITOR)
        return [];
#else
        var pData = chunk.GetUnsafePtr();
        var count = chunk._count;
        var capacity = chunk._capacity;
        var worldID = chunk._worldID;
        var archetypeID = chunk._archetypeID;

        if (count == 0)
        {
            return [];
        }

        var views = new List<object>();
        var world = World.GetWorld(worldID);
        if (world is null)
        {
            return [];
        }

        ref var archetype = ref world.ComponentManager.GetArchetypeReference(archetypeID);
        var it = archetype._signature.GetIterator();
        while (it.Next(out var index))
        {
            var type = ComponentRegistry.s_runtimeIDToType[index];
            if (type == null)
            {
                continue;
            }
            var layout = archetype.GetLayout(index).Value;
            var readMethod = typeof(ChunkDebugView)
                .GetMethod(nameof(ReadComponentArray), System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Static)!
                .MakeGenericMethod(type);

            // 3. Invoke it to get a Position[] or Velocity[]
            var array = readMethod.Invoke(null, [(long)pData, layout.offset, count]);
            if (array == null)
            {
                continue;
            }

            // 4. Wrap it in a nice label so the debugger shows "Position[]"
            views.Add(new ComponentArrayView(type.Name, array));
        }

        return [.. views];
#endif
    }
}

[DebuggerTypeProxy(typeof(ChunkDebugView))]
internal unsafe struct Chunk : IDisposable
{
    public const int CHUNK_BUFFER_SIZE = 16384; // 16 KB
    public const int BIT_ALIGNMENT = 8;
    public const int BIT_SHIFT = 3; // log2(BIT_ALIGNMENT)
    public const int BIT_ALIGNMENT_MINUS_ONE = BIT_ALIGNMENT - 1;

    private UnsafeArray<byte> _data;
    private UnsafeArray<int> _versions;

    // TODO: Add structual change versioning, similar to DidOrderChange in unity ecs.
    internal int _structuralVersion;

    internal int _count;
    internal readonly int _capacity;

#if DEBUG || GHOST_EDITOR
    // For debugging purpose
    internal int _worldID;
    internal int _archetypeID;
#endif

    public Chunk(int bufferSize, int capacity, int componentCount, int globalVersion)
    {
        _data = new UnsafeArray<byte>(bufferSize, Allocator.Persistent, AllocationOption.Clear);
        _versions = new UnsafeArray<int>(componentCount, Allocator.Persistent);
        _capacity = capacity;
        _count = 0;

        _versions.AsSpan().Fill(globalVersion);
        _structuralVersion = globalVersion;
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public readonly byte* GetUnsafePtr()
    {
        return (byte*)_data.GetUnsafePtr();
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public readonly int* GetVersionUnsafePtr()
    {
        return (int*)_versions.GetUnsafePtr();
    }

    public void Dispose()
    {
        _data.Dispose();
        _versions.Dispose();
    }
}

internal unsafe struct Archetype : IDisposable
{
    internal struct ComponentMemoryLayout
    {
        public int componentID;
        public int size;
        public int offset;
        public int enableBitsOffset;
        public int versionIndex;
    }

    private struct Edge
    {
        public int componentID;
        public int targetArchetype; // can't use Identifier<Archetype> because cycle causer
    }

    internal UnsafeBitSet _signature;
    internal UnsafeList<Chunk> _chunks;
    internal UnsafeArray<ComponentMemoryLayout> _layouts;

    private UnsafeArray<int> _componentIDToLayoutIndex;
    // TODO: Is hash map better?
    private UnsafeList<Edge> _edgesAdd;
    private UnsafeList<Edge> _edgesRemove;

    private readonly Identifier<Archetype> _id;
    private readonly Identifier<World> _worldID;

    private readonly int _hash;
    private int _entityCapacity;
    private int _maxComponentID;
    private int _entityIdsOffset;

    public readonly Identifier<Archetype> ID => _id;
    public readonly Identifier<World> WorldID => _worldID;

    public readonly int EntityCapacity => _entityCapacity;
    public readonly int ChunkCount => _chunks.Count;
    public readonly int EntityIDsOffset => _entityIdsOffset;

    public Archetype(Identifier<Archetype> id, Identifier<World> worldID, ReadOnlySpan<Identifier<IComponent>> componentIds)
    {
        _id = id;
        _worldID = worldID;

        _chunks = new UnsafeList<Chunk>(4, Allocator.Persistent);
        _edgesAdd = new UnsafeList<Edge>(4, Allocator.Persistent);
        _edgesRemove = new UnsafeList<Edge>(4, Allocator.Persistent);

        if (componentIds.IsEmpty)
        {
            _signature = new UnsafeBitSet(1, Allocator.Persistent, AllocationOption.Clear);
            _hash = 0;

            _signature.ClearAll();
            _entityCapacity = Chunk.CHUNK_BUFFER_SIZE / sizeof(Entity);

            return;
        }

        var highestComponentID = 0;
        for (var i = 0; i < componentIds.Length; i++)
        {
            if (componentIds[i] > highestComponentID)
            {
                highestComponentID = componentIds[i];
            }
        }

        _signature = new UnsafeBitSet(highestComponentID + 1, Allocator.Persistent, AllocationOption.Clear);
        _hash = _signature.GetHashCode();

        CalculateLayout(componentIds);
    }

    private void CalculateLayout(ReadOnlySpan<Identifier<IComponent>> componentIds)
    {
        var entitySize = sizeof(Entity);
        var entityAlign = (int)MemoryUtility.AlignOf<Entity>();

        var components = (Span<ComponentInfo>)stackalloc ComponentInfo[componentIds.Length];
        for (var i = 0; i < componentIds.Length; i++)
        {
            _signature.SetBit(componentIds[i]);
            components[i] = ComponentRegistry.GetComponentInfo(componentIds[i]);
        }

        // Calculate total size per entity to get an initial capacity estimate
        var bytesPerEntity = entitySize;
        var maxComponentID = 0;
        for (var i = 0; i < components.Length; i++)
        {
            var comp = components[i];
            bytesPerEntity += comp.size;
            if (comp.id > maxComponentID)
            {
                maxComponentID = comp.id;
            }
        }

        _maxComponentID = maxComponentID;
        _entityCapacity = Chunk.CHUNK_BUFFER_SIZE / bytesPerEntity;
        _layouts = new UnsafeArray<ComponentMemoryLayout>(components.Length, Allocator.Persistent);
        _componentIDToLayoutIndex = new UnsafeArray<int>(_maxComponentID + 1, Allocator.Persistent);

        _componentIDToLayoutIndex.AsSpan().Fill(-1);

        components.Sort((a, b) => b.alignment.CompareTo(a.alignment));
        var tempOffsets = stackalloc int[components.Length];
        var tempBitmaskOffsets = stackalloc int[components.Length];

        while (_entityCapacity > 0)
        {
            var currentOffset = 0;
            var fits = true;

            currentOffset = (currentOffset + entityAlign - 1) & ~(entityAlign - 1);

            _entityIdsOffset = currentOffset;
            currentOffset += _entityCapacity * entitySize;

            for (var i = 0; i < components.Length; i++)
            {
                var size = components[i].size;
                var align = components[i].alignment;

                currentOffset = (currentOffset + align - 1) & ~(align - 1);
                tempOffsets[i] = currentOffset;
                currentOffset += _entityCapacity * size;

                var bitmaskOffset = -1;
                if (components[i].isEnableable)
                {
                    var bitmaskSize = (_entityCapacity + Chunk.BIT_ALIGNMENT_MINUS_ONE) / Chunk.BIT_ALIGNMENT;
                    // Reserve space for the bitmask (1 bit per entity)

                    currentOffset = (currentOffset + Chunk.BIT_ALIGNMENT_MINUS_ONE) & ~Chunk.BIT_ALIGNMENT_MINUS_ONE; // Align
                    bitmaskOffset = currentOffset;
                    currentOffset += bitmaskSize;
                }

                tempBitmaskOffsets[i] = bitmaskOffset;

                if (currentOffset > Chunk.CHUNK_BUFFER_SIZE)
                {
                    fits = false;
                    break;
                }
            }

            if (fits)
            {
                for (var i = 0; i < components.Length; i++)
                {
                    _layouts[i] = new ComponentMemoryLayout
                    {
                        componentID = components[i].id,
                        offset = tempOffsets[i],
                        size = components[i].size,
                        enableBitsOffset = tempBitmaskOffsets[i],
                        versionIndex = i
                    };

                    _componentIDToLayoutIndex[components[i].id] = i;
                }

                return;
            }

            _entityCapacity--;
        }
    }

    public void AllocateEntity(out int chunkIndex, out int rowIndex)
    {
        var world = World.GetWorldUncheck(_worldID);

        for (var i = 0; i < _chunks.Count; i++)
        {
            ref var chunk = ref _chunks[i];
            if (chunk._count < _entityCapacity)
            {
                rowIndex = chunk._count;
                chunk._count++;
                chunk._structuralVersion = world.Version;
                chunkIndex = i;

                return;
            }
        }

        // Need to allocate a new chunk
        var newChunk = new Chunk(Chunk.CHUNK_BUFFER_SIZE, _entityCapacity, _layouts.Count, world.Version);
#if DEBUG || GHOST_EDITOR
        newChunk._worldID = _worldID;
        newChunk._archetypeID = _id;
#endif

        // Set all enable to true by default for enableable components
        for (var i = 0; i < _layouts.Count; i++)
        {
            var layout = _layouts[i];
            if (layout.enableBitsOffset != -1)
            {
                var pChunk = newChunk.GetUnsafePtr();
                var pBits = pChunk + layout.enableBitsOffset;
                MemoryUtility.MemSet(pBits, 0xFF, (nuint)((_entityCapacity + Chunk.BIT_ALIGNMENT_MINUS_ONE) / Chunk.BIT_ALIGNMENT));
            }
        }

        rowIndex = 0;
        newChunk._count++;
        chunkIndex = _chunks.Count;

        _chunks.Add(newChunk);
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public readonly Entity GetEntity(int chunkIndex, int rowIndex)
    {
        var chunk = _chunks[chunkIndex];
        var chunkBase = chunk.GetUnsafePtr();
        var src = chunkBase + _entityIdsOffset + (sizeof(Entity) * rowIndex);

        return *(Entity*)src;
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public readonly void SetEntity(int chunkIndex, int rowIndex, Entity entity)
    {
        var chunk = _chunks[chunkIndex];
        var chunkBase = chunk.GetUnsafePtr();
        var dst = chunkBase + _entityIdsOffset + (sizeof(Entity) * rowIndex);

        MemoryUtility.MemCpy(dst, &entity, (nuint)sizeof(Entity));
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public readonly Error SetComponentData(int chunkIndex, int rowIndex, Identifier<IComponent> componentID, void* pComponent)
    {
        var r = GetLayout(componentID);
        if (r.Error != Error.None)
        {
            return r.Error;
        }

        var offset = r.Value.offset;
        ref var chunk = ref _chunks[chunkIndex];

        var chunkBase = chunk.GetUnsafePtr();
        var size = ComponentRegistry.GetComponentInfo(componentID).size;
        var dst = chunkBase + offset + (size * rowIndex);

        MemoryUtility.MemCpy(dst, pComponent, (nuint)size);

        var world = World.GetWorldUncheck(_worldID);
        MarkChanged(chunkIndex, componentID, world.Version);

        return Error.None;
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public readonly void* GetComponentData(int chunkIndex, int rowIndex, Identifier<IComponent> componentID)
    {
        var r = GetLayout(componentID);
        if (r.Error != Error.None)
        {
            return null;
        }

        var offset = r.Value.offset;
        var chunk = _chunks[chunkIndex];

        var chunkBase = chunk.GetUnsafePtr();
        var size = ComponentRegistry.GetComponentInfo(componentID).size;
        return chunkBase + offset + (size * rowIndex);
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public ref Chunk GetChunkReference(int index)
    {
        return ref _chunks[index];
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public readonly Result<ComponentMemoryLayout, Error> GetLayout(int componentID)
    {
        if (componentID >= _componentIDToLayoutIndex.Count)
        {
            return Error.InvalidArgument;
        }

        var layoutIndex = _componentIDToLayoutIndex[componentID];
        if (layoutIndex == -1)
        {
            return Error.NotFound;
        }

        return _layouts[layoutIndex];
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public readonly Error MarkChanged(int chunkIndex, int componentTypeId, int globalVersion)
    {
        var layoutResult = GetLayout(componentTypeId);
        if (layoutResult.IsFailure)
        {
            return layoutResult.Error;
        }

        ref var chunk = ref _chunks[chunkIndex];
        chunk.GetVersionUnsafePtr()[layoutResult.Value.versionIndex] = globalVersion;

        return Error.None;
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public readonly Result<int, Error> GetVersion(int chunkIndex, int componentTypeId)
    {
        var layoutResult = GetLayout(componentTypeId);
        if (layoutResult.Error != Error.None)
        {
            return layoutResult.Error;
        }

        ref var chunk = ref _chunks[chunkIndex];
        return chunk.GetVersionUnsafePtr()[layoutResult.Value.versionIndex];
    }

    public Error RemoveEntity(int chunkIndex, int rowIndex)
    {
        if (chunkIndex < 0 || chunkIndex >= _chunks.Count)
        {
            return Error.InvalidArgument;
        }

        var world = World.GetWorldUncheck(_worldID);
        ref var chunk = ref _chunks[chunkIndex];
        var lastIndex = chunk._count - 1;

        // If we are NOT removing the very last entity, we must swap.
        if (rowIndex != lastIndex)
        {
            var chunkBase = chunk.GetUnsafePtr();
            var pLastEntity = chunkBase + _entityIdsOffset + (sizeof(Entity) * lastIndex);
            var pRowEntity = chunkBase + _entityIdsOffset + (sizeof(Entity) * rowIndex);

            var result = world.EntityManager.UpdateEntityLocation(*(Entity*)pLastEntity, _id, chunkIndex, rowIndex);
            if (result != Error.None)
            {
                return result;
            }

            // Only operate the swap back after the update is succeed.
            MemoryUtility.MemCpy(pRowEntity, pLastEntity, (nuint)sizeof(Entity));

            for (var i = 0; i < _layouts.Count; i++)
            {
                var layout = _layouts[i];

                var pRow = chunk.GetUnsafePtr() + layout.offset + (layout.size * rowIndex);
                var pLast = chunk.GetUnsafePtr() + layout.offset + (layout.size * lastIndex);

                MemoryUtility.MemCpy(pRow, pLast, (nuint)layout.size);
            }
        }

        chunk._count--;
        chunk._structuralVersion = world.Version;

        return Error.None;
    }

    public Error RemoveEntities(int chunkIndex, ReadOnlySpan<int> sortedIndicesToRemove)
    {
        if (chunkIndex < 0 || chunkIndex >= _chunks.Count)
        {
            return Error.InvalidArgument;
        }

        if (sortedIndicesToRemove.Length == 0)
        {
            return Error.None;
        }

        ref var chunk = ref _chunks[chunkIndex];

        int oldCount = chunk._count;
        int removeCount = sortedIndicesToRemove.Length;
        int newCount = oldCount - removeCount; // The boundary between "Keep" and "Drop"

        var chunkBase = chunk.GetUnsafePtr();
        var world = World.GetWorldUncheck(_worldID); // Typo fixed from 'wrold'

        // Pointers for the swap logic
        // 1. 'holePtr' tracks which index in the sorted list we are processing
        int holePtr = 0;

        // 2. 'candidateIndex' starts at the end of the OLD array and moves backward
        int candidateIndex = oldCount - 1;

        // 3. 'removalTailPtr' tracks removals at the end of the array to skip them
        int removalTailPtr = sortedIndicesToRemove.Length - 1;

        // Iterate through the holes that are strictly INSIDE the new valid range
        while (holePtr < removeCount)
        {
            int holeIndex = sortedIndicesToRemove[holePtr];

            // If the current hole is beyond the new count, it's in the "Drop Zone".
            // Since the list is sorted, all subsequent holes are also in the drop zone.
            // We are done filling holes.
            if (holeIndex >= newCount)
                break;

            // --- Find a Valid Filler ---
            // We look for an entity at the end of the array that IS NOT scheduled for removal.
            while (candidateIndex >= newCount)
            {
                // Check if the current candidate is actually marked for removal
                bool isCandidateRemoved = false;

                // Because sortedIndices is sorted, we check the end of the list 
                // to see if the candidateIndex matches a removal request.
                if (removalTailPtr >= 0 && sortedIndicesToRemove[removalTailPtr] == candidateIndex)
                {
                    isCandidateRemoved = true;
                    removalTailPtr--; // Consume this removal
                }

                if (!isCandidateRemoved)
                {
                    // Found a valid filler!
                    break;
                }

                // This candidate was also removed, so skip it and keep looking left
                candidateIndex--;
            }

            // --- Perform The Swap ---
            // Move 'candidateIndex' (Filler) into 'holeIndex' (Hole)

            var pFillerEntity = chunkBase + _entityIdsOffset + (sizeof(Entity) * candidateIndex);
            var pHoleEntity = chunkBase + _entityIdsOffset + (sizeof(Entity) * holeIndex);

            // 1. Update the Map (Critical Step)
            // We tell the world: "The entity that WAS at 'candidateIndex' is now at 'holeIndex'"
            var result = world.EntityManager.UpdateEntityLocation(*(Entity*)pFillerEntity, _id, chunkIndex, holeIndex);
            if (result != Error.None)
            {
                return result;
            }

            // 2. Overwrite Entity ID
            MemoryUtility.MemCpy(pHoleEntity, pFillerEntity, (nuint)sizeof(Entity));

            // 3. Overwrite Components
            for (var i = 0; i < _layouts.Count; i++)
            {
                var layout = _layouts[i];
                var pRow = chunkBase + layout.offset + (layout.size * holeIndex);
                var pLast = chunkBase + layout.offset + (layout.size * candidateIndex);
                MemoryUtility.MemCpy(pRow, pLast, (nuint)layout.size);
            }

            // Prepare for next hole
            holePtr++;
            candidateIndex--;
        }

        chunk._count = newCount;
        chunk._structuralVersion = world.Version;

        return Error.None;
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public readonly bool HasComponent(Identifier<IComponent> componentID)
    {
        return _signature.IsSet(componentID);
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public void AddEdgeAdd(Identifier<IComponent> componentID, Identifier<Archetype> targetArchetype)
    {
        _edgesAdd.Add(new Edge
        {
            componentID = componentID,
            targetArchetype = targetArchetype
        });
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public readonly Identifier<Archetype> GetEdgeAdd(Identifier<IComponent> componentID)
    {
        for (var i = 0; i < _edgesAdd.Count; i++)
        {
            var edge = _edgesAdd[i];
            if (edge.componentID == componentID)
            {
                return edge.targetArchetype;
            }
        }

        return Identifier<Archetype>.Invalid;
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public void AddEdgeRemove(Identifier<IComponent> componentID, Identifier<Archetype> targetArchetype)
    {
        _edgesRemove.Add(new Edge
        {
            componentID = componentID,
            targetArchetype = targetArchetype
        });
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public readonly Identifier<Archetype> GetEdgeRemove(Identifier<IComponent> componentID)
    {
        for (var i = 0; i < _edgesRemove.Count; i++)
        {
            var edge = _edgesRemove[i];
            if (edge.componentID == componentID)
            {
                return edge.targetArchetype;
            }
        }

        return Identifier<Archetype>.Invalid;
    }

    public override readonly int GetHashCode()
    {
        return _hash;
    }

    public void Dispose()
    {
        if (_chunks.IsCreated)
        {
            foreach (ref var chunk in _chunks)
            {
                chunk.Dispose();
            }
        }

        _signature.Dispose();
        _chunks.Dispose();
        _componentIDToLayoutIndex.Dispose();
        _layouts.Dispose();
        _edgesAdd.Dispose();
        _edgesRemove.Dispose();
    }
}