encompass-cs/encompass-cs/ComponentBitSet.cs

using System;
using System.Collections;
using System.Collections.Generic;

namespace Encompass
{
    internal class ComponentBitSet
    {
        BitArrayPool bitArrayPool = new BitArrayPool(32768); // todo: set entity cap
        Dictionary<Entity, BitArray> entities = new Dictionary<Entity, BitArray>();
        Dictionary<Type, int> TypeToIndex { get; set; }
        BitArray withQueryArray;
        BitArray withoutQueryArray;
        BitArray emptyArray;

        public ComponentBitSet(Dictionary<Type, int> typeToIndex)
        {
            TypeToIndex = typeToIndex;
        }

        public void FinishRegistering()
        {
            withQueryArray = new BitArray(TypeToIndex.Count);
            withoutQueryArray = new BitArray(TypeToIndex.Count);
            emptyArray = new BitArray(TypeToIndex.Count);
            foreach (var kvp in entities)
            {
                kvp.Value.Length = TypeToIndex.Count;
            }
        }

        public void Clear()
        {
            foreach (var kvp in entities)
            {
                bitArrayPool.Free(kvp.Value);
            }
            entities.Clear();
        }

        public void AddEntity(Entity entity)
        {
            var bitArray = bitArrayPool.Obtain(TypeToIndex.Count);
            entities.Add(entity, bitArray);
        }

        public void Set<TComponent>(Entity entity) where TComponent : struct, IComponent
        {
            if (!entities.ContainsKey(entity)) { AddEntity(entity); }
            entities[entity].Set(TypeToIndex[typeof(TComponent)], true);
        }

        public void RemoveComponent<TComponent>(Entity entity) where TComponent : struct, IComponent
        {
            if (entities.ContainsKey(entity))
            {
                entities[entity].Set(TypeToIndex[typeof(TComponent)], false);
            }
        }

        public void RemoveEntity(Entity entity)
        {
            if (entities.ContainsKey(entity))
            {
                bitArrayPool.Free(entities[entity]);
                entities.Remove(entity);
            }
        }

        public BitArray EntityBitArray(Entity entity)
        {
            return entities.ContainsKey(entity) ? entities[entity] : emptyArray;
        }
    }
}