encompass-cs/encompass-cs/EntityManager.cs

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

namespace Encompass
{
    internal class EntityManager
    {
        private readonly ConcurrentDictionary<Guid, Entity> IDToEntity = new ConcurrentDictionary<Guid, Entity>();

        private readonly HashSet<Guid> entitiesMarkedForDestroy = new HashSet<Guid>();

        private readonly Dictionary<Type, HashSet<IEntityTracker>> componentTypeToEntityTrackers = new Dictionary<Type, HashSet<IEntityTracker>>();
        private readonly Dictionary<Guid, HashSet<IEntityTracker>> entityToEntityTrackers = new Dictionary<Guid, HashSet<IEntityTracker>>();

        private readonly ConcurrentDictionary<Guid, byte> entitiesWithAddedComponents;
        private readonly ConcurrentDictionary<Guid, byte> entitiesWithRemovedComponents;

        private readonly ComponentManager componentManager;

        public EntityManager(
            ComponentManager componentManager,
            ConcurrentDictionary<Guid, byte> entitiesWithAddedComponents,
            ConcurrentDictionary<Guid, byte> entitiesWithRemovedComponents
        )
        {
            this.componentManager = componentManager;
            this.entitiesWithAddedComponents = entitiesWithAddedComponents;
            this.entitiesWithRemovedComponents = entitiesWithRemovedComponents;
        }

        public Entity CreateEntity()
        {
            var id = NextID();
            var entity = new Entity(id);
            IDToEntity[id] = entity;
            componentManager.RegisterEntity(id);
            return entity;
        }

        public bool EntityExists(Guid id)
        {
            return IDToEntity.ContainsKey(id);
        }

        public Entity GetEntity(Guid id)
        {
            return IDToEntity[id];
        }

        public void MarkForDestroy(Guid entityID)
        {
            entitiesMarkedForDestroy.Add(entityID);
        }

        public void DestroyMarkedEntities()
        {
            foreach (var entityID in entitiesMarkedForDestroy)
            {
                componentManager.MarkAllComponentsOnEntityForRemoval(entityID);
                IDToEntity.TryRemove(entityID, out _);
                entityToEntityTrackers.Remove(entityID);
                componentManager.RegisterDestroyedEntity(entityID);
            }

            entitiesMarkedForDestroy.Clear();
        }

        private Guid NextID()
        {
            return Guid.NewGuid();
        }

        public void RegisterEntityTracker(IEntityTracker entityTracker)
        {
            foreach (var componentType in entityTracker.ComponentTypes)
            {
                if (!componentTypeToEntityTrackers.ContainsKey(componentType))
                {
                    componentTypeToEntityTrackers.Add(componentType, new HashSet<IEntityTracker>());
                }

                componentTypeToEntityTrackers[componentType].Add(entityTracker);
            }

            if (entityTracker is EntityRenderer)
            {
                var entityRenderer = entityTracker as EntityRenderer;
                if (!componentTypeToEntityTrackers.ContainsKey(entityRenderer.DrawComponentType))
                {
                    componentTypeToEntityTrackers.Add(entityRenderer.DrawComponentType, new HashSet<IEntityTracker>());
                }

                componentTypeToEntityTrackers[entityRenderer.DrawComponentType].Add(entityRenderer);
            }
        }

        public void RegisterDirtyEntityWithAddedComponents(Guid entityID)
        {
            entitiesWithAddedComponents.TryAdd(entityID, 0);
        }

        public void RegisterDirtyEntityWithRemovedComponents(Guid entityID)
        {
            entitiesWithRemovedComponents.TryAdd(entityID, 0);
        }

        public void CheckEntitiesWithAddedComponents()
        {
            foreach (var entityID in entitiesWithAddedComponents.Keys)
            {
                CheckAndRegisterEntity(entityID);
            }

            entitiesWithAddedComponents.Clear();
        }

        public void CheckEntitiesWithRemovedComponents()
        {
            foreach (var entityID in entitiesWithRemovedComponents.Keys)
            {
                if (entityToEntityTrackers.ContainsKey(entityID))
                {
                    foreach (var engine in entityToEntityTrackers[entityID])
                    {
                        engine.CheckAndUntrackEntity(entityID);
                    }
                }
            }

            entitiesWithRemovedComponents.Clear();
        }

        private void CheckAndRegisterEntity(Guid entityID)
        {
            var componentTypes = componentManager.GetAllComponentTypesOfEntity(entityID);
            foreach (var componentType in componentTypes)
            {
                if (componentTypeToEntityTrackers.ContainsKey(componentType))
                {
                    foreach (var entityTracker in componentTypeToEntityTrackers[componentType])
                    {
                        if (entityTracker.CheckAndTrackEntity(entityID))
                        {
                            if (!entityToEntityTrackers.ContainsKey(entityID))
                            {
                                entityToEntityTrackers.Add(entityID, new HashSet<IEntityTracker>());
                            }
                            entityToEntityTrackers[entityID].Add(entityTracker);
                        }
                    }
                }
            }
        }
    }
}