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 = new Dictionary<Type, int>();
        BitArray queryArray;

        public void RegisterType<TComponent>() where TComponent : struct, IComponent
        {
            typeToIndex.Add(typeof(TComponent), typeToIndex.Count);
            foreach (var kvp in entities)
            {
                kvp.Value.Length = typeToIndex.Count;
            }
        }

        public void FinishRegistering()
        {
            queryArray = new BitArray(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
        {
            entities[entity].Set(typeToIndex[typeof(TComponent)], false);
        }

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

        public IEnumerable<Entity> EntitiesWithComponents(IEnumerable<Type> types)
        {
            foreach (var kvp in entities)
            {
                queryArray.SetAll(false);
                foreach (var type in types)
                {
                    queryArray.Set(typeToIndex[type], true);
                }
                queryArray.And(kvp.Value);
                var hasComponents = true;
                foreach (var type in types)
                {
                    if (!queryArray.Get(typeToIndex[type])) {
                        hasComponents = false;
                        break;
                    }
                }
                if (hasComponents) { yield return kvp.Key; }
            }
        }
    }
}