encompass-cs/encompass-cs/ComponentManager.cs

using System;
using System.Collections.Generic;
using System.Linq;
using Encompass.Exceptions;
using Collections.Pooled;

namespace Encompass
{
    internal class ComponentManager
    {
        private readonly DrawLayerManager drawLayerManager;

        private readonly Dictionary<Guid, Type> componentIDToType = new Dictionary<Guid, Type>();
        private readonly ComponentStore IDToComponent = new ComponentStore();
        private readonly Dictionary<Guid, PooledSet<Guid>> entityIDToComponentIDs = new Dictionary<Guid, PooledSet<Guid>>();
        private readonly Dictionary<Guid, Guid> componentIDToEntityID = new Dictionary<Guid, Guid>();

        private readonly Dictionary<Guid, PooledDictionary<Type, Guid>> entityIDToComponentTypeToComponentID = new Dictionary<Guid, PooledDictionary<Type, Guid>>();

        private readonly Dictionary<Type, HashSet<Guid>> typeToComponentIDs = new Dictionary<Type, HashSet<Guid>>();

        private readonly Dictionary<(Entity, Type), (Guid, IComponent)> componentWriteData = new Dictionary<(Entity, Type), (Guid, IComponent)>();
        private readonly Dictionary<(Entity, Type), int> componentWritePriorities = new Dictionary<(Entity, Type), int>();
        private readonly HashSet<Guid> componentIDsMarkedForWrite = new HashSet<Guid>();
        private readonly HashSet<Guid> componentsMarkedForRemoval = new HashSet<Guid>();

        public ComponentManager(DrawLayerManager drawLayerManager)
        {
            this.drawLayerManager = drawLayerManager;
        }

        internal void RegisterEntity(Guid entityID)
        {
            entityIDToComponentIDs.Add(entityID, new PooledSet<Guid>());
            entityIDToComponentTypeToComponentID.Add(entityID, new PooledDictionary<Type, Guid>());
        }

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

        internal Guid MarkComponentForWrite<TComponent>(Entity entity, TComponent component, int priority) where TComponent : struct, IComponent
        {
            Guid id;
            if (EntityHasComponentOfType<TComponent>(entity))
            {
                id = GetComponentByEntityAndType<TComponent>(entity).Item1;
            }
            else
            {
                id = NextID();
            }

            if (componentWriteData.ContainsKey((entity, typeof(TComponent))))
            {
                var currentPriority = componentWritePriorities[(entity, typeof(TComponent))];
                if (priority < currentPriority)
                {
                    componentWriteData[(entity, typeof(TComponent))] = (id, component);
                    componentWritePriorities[(entity, typeof(TComponent))] = priority;
                    componentIDsMarkedForWrite.Add(id);
                }
            }
            else
            {
                componentWriteData.Add((entity, typeof(TComponent)), (id, component));
                componentWritePriorities[(entity, typeof(TComponent))] = priority;
                componentIDsMarkedForWrite.Add(id);
            }

            return id;
        }

        internal void RegisterDrawableComponent<TComponent>(Guid componentID, TComponent component) where TComponent : IDrawableComponent
        {
            drawLayerManager.RegisterComponentWithLayer(componentID, component.Layer);
        }

        internal void AddComponent(Entity entity, Type type, Guid componentID, IComponent component)
        {
            IDToComponent.Set(type, componentID, component);
            componentIDToEntityID[componentID] = entity.ID;
            componentIDToType[componentID] = type;
            entityIDToComponentTypeToComponentID[entity.ID][type] = componentID;
            if (!typeToComponentIDs.ContainsKey(type))
            {
                typeToComponentIDs.Add(type, new HashSet<Guid>());
            }
            typeToComponentIDs[type].Add(componentID);
            entityIDToComponentIDs[entity.ID].Add(componentID);
        }

        internal void UpdateComponent(Guid componentID, IComponent component)
        {
            IDToComponent[componentID] = component;
        }

        internal void WriteComponents()
        {
            foreach (var keyValuePair in componentWriteData)
            {
                var (entity, type) = keyValuePair.Key;
                var (componentID, component) = keyValuePair.Value;

                if (!componentIDsMarkedForWrite.Contains(componentID) || !entityIDToComponentTypeToComponentID.ContainsKey(entity.ID)) { continue; }

                if (entityIDToComponentTypeToComponentID[entity.ID].ContainsKey(type))
                {
                    UpdateComponent(componentID, component);
                }
                else
                {
                    AddComponent(entity, type, componentID, component);
                }
            }

            componentWriteData.Clear();
            componentIDsMarkedForWrite.Clear();
            componentWritePriorities.Clear();
        }

        internal IEnumerable<Guid> GetComponentIDsByEntityID(Guid entityID)
        {
            if (entityIDToComponentIDs.TryGetValue(entityID, out PooledSet<Guid> idSet))
            {
                return idSet;
            }
            return Enumerable.Empty<Guid>();
        }

        internal IEnumerable<(Guid, TComponent)> GetComponentsByType<TComponent>() where TComponent : struct, IComponent
        {
            if (typeToComponentIDs.TryGetValue(typeof(TComponent), out HashSet<Guid> idSet))
            {
                return idSet.Select(id => (id, IDToComponent.Get<TComponent>(id)));
            }
            return Enumerable.Empty<(Guid, TComponent)>();
        }

        internal (Guid, TComponent) GetComponentByEntityAndType<TComponent>(Entity entity) where TComponent : struct, IComponent
        {
            if (entityIDToComponentTypeToComponentID.ContainsKey(entity.ID) && entityIDToComponentTypeToComponentID[entity.ID].TryGetValue(typeof(TComponent), out Guid id))
            {
                return (id, IDToComponent.Get<TComponent>(id));
            }

            throw new NoComponentOfTypeOnEntityException("No Component of type {0} exists on Entity {1}", typeof(TComponent).Name, entity.ID);
        }

        internal bool EntityHasComponentOfType<TComponent>(Entity entity) where TComponent : struct, IComponent
        {
            return (entityIDToComponentTypeToComponentID.ContainsKey(entity.ID) && entityIDToComponentTypeToComponentID[entity.ID].ContainsKey(typeof(TComponent)));
        }

        internal bool ComponentOfTypeExists<TComponent>() where TComponent : struct, IComponent
        {
            if (typeToComponentIDs.TryGetValue(typeof(TComponent), out HashSet<Guid> idSet))
            {
                return idSet.Count > 0;
            }

            return false;
        }

        internal IComponent GetComponentByID(Guid componentID)
        {
            if (IDToComponent.Has<TComponent>(componentID))
            {
                return IDToComponent[componentID];
            }
            else
            {
                throw new ComponentNotFoundException("Component with ID {0} does not exist.", componentID);
            }
        }

        internal Type GetComponentTypeByID(Guid componentID)
        {
            if (componentIDToType.ContainsKey(componentID))
            {
                return componentIDToType[componentID];
            }
            else
            {
                throw new ComponentNotFoundException("Component with ID {0} does not exist.", componentID);
            }
        }

        internal Guid GetEntityIDByComponentID(Guid componentID)
        {
            if (componentIDToEntityID.ContainsKey(componentID))
            {
                return componentIDToEntityID[componentID];
            }
            else
            {
                throw new ComponentNotFoundException("Component with ID {0} does not exist.", componentID);
            }
        }

        internal void MarkAllComponentsOnEntityForRemoval(Guid entityID)
        {
            foreach (var componentID in GetComponentIDsByEntityID(entityID))
            {
                MarkForRemoval(componentID);
            }
        }

        internal void MarkForRemoval(Guid componentID)
        {
            componentsMarkedForRemoval.Add(componentID);
        }

        internal void RemoveMarkedComponents()
        {
            foreach (var componentID in componentsMarkedForRemoval)
            {
                if (componentIDsMarkedForWrite.Contains(componentID))
                {
                    componentIDsMarkedForWrite.Remove(componentID);
                }

                if (IDToComponent.ContainsKey(componentID))
                {
                    Remove(componentID);
                }
            }

            componentsMarkedForRemoval.Clear();
        }

        private void Remove<TComponent>(Guid componentID) where TComponent : struct, IComponent
        {
            var type = componentIDToType[componentID];

            var entityID = componentIDToEntityID[componentID];
            if (entityIDToComponentIDs.ContainsKey(entityID))
            {
                entityIDToComponentIDs[entityID].Remove(componentID);
            }

            if (entityIDToComponentTypeToComponentID.ContainsKey(entityID))
            {
                entityIDToComponentTypeToComponentID[entityID].Remove(type);
            }

            IDToComponent.Remove<TComponent>(componentID);
            componentIDToType.Remove(componentID);
            componentIDToEntityID.Remove(componentID);
            typeToComponentIDs[type].Remove(componentID);

            drawLayerManager.UnRegisterComponentWithLayer(componentID);
        }

        public void RegisterDestroyedEntity(Guid entityID)
        {
            entityIDToComponentIDs[entityID].Dispose();
            entityIDToComponentIDs.Remove(entityID);

            entityIDToComponentTypeToComponentID[entityID].Dispose();
            entityIDToComponentTypeToComponentID.Remove(entityID);
        }
    }
}