MoonTools.ECS/src/World.cs

using System;
using System.Collections.Generic;
using System.Runtime.CompilerServices;
using MoonTools.ECS.Collections;

namespace MoonTools.ECS
{
	public class World : IDisposable
	{
		// Get TypeId from a Type
		private readonly Dictionary<Type, TypeId> TypeToId = new Dictionary<Type, TypeId>();

		#if DEBUG
		private Dictionary<TypeId, Type> IdToType = new Dictionary<TypeId, Type>();
		#endif

		// Get element size from a TypeId
		private readonly Dictionary<TypeId, int> ElementSizes = new Dictionary<TypeId, int>();

		// Filters
		internal readonly Dictionary<FilterSignature, Filter> FilterIndex = new Dictionary<FilterSignature, Filter>();
		private readonly Dictionary<TypeId, List<Filter>> TypeToFilter = new Dictionary<TypeId, List<Filter>>();

		// TODO: can we make the tag an native array of chars at some point?
		internal Dictionary<Entity, string> EntityTags = new Dictionary<Entity, string>();

		// Relation Storages
		internal Dictionary<TypeId, RelationStorage> RelationIndex = new Dictionary<TypeId, RelationStorage>();
		internal Dictionary<Entity, IndexableSet<TypeId>> EntityRelationIndex = new Dictionary<Entity, IndexableSet<TypeId>>();

		// Message Storages
		private Dictionary<TypeId, MessageStorage> MessageIndex = new Dictionary<TypeId, MessageStorage>();

		public FilterBuilder FilterBuilder => new FilterBuilder(this);

		internal readonly Dictionary<TypeId, ComponentStorage> ComponentIndex = new Dictionary<TypeId, ComponentStorage>();
		internal Dictionary<Entity, IndexableSet<TypeId>> EntityComponentIndex = new Dictionary<Entity, IndexableSet<TypeId>>();

		internal IdAssigner EntityIdAssigner = new IdAssigner();
		private IdAssigner TypeIdAssigner = new IdAssigner();

		private bool IsDisposed;

		internal TypeId GetTypeId<T>() where T : unmanaged
		{
			if (TypeToId.ContainsKey(typeof(T)))
			{
				return TypeToId[typeof(T)];
			}

			var typeId = new TypeId(TypeIdAssigner.Assign());
			TypeToId.Add(typeof(T), typeId);
			ElementSizes.Add(typeId, Unsafe.SizeOf<T>());

			#if DEBUG
			IdToType.Add(typeId, typeof(T));
			#endif

			return typeId;
		}

		internal TypeId GetComponentTypeId<T>() where T : unmanaged
		{
			var typeId = GetTypeId<T>();
			if (ComponentIndex.TryGetValue(typeId, out var componentStorage))
			{
				return typeId;
			}

			componentStorage = new ComponentStorage(typeId, ElementSizes[typeId]);
			ComponentIndex.Add(typeId, componentStorage);
			TypeToFilter.Add(typeId, new List<Filter>());
			return typeId;
		}

		[MethodImpl(MethodImplOptions.AggressiveInlining)]
		private ComponentStorage GetComponentStorage<T>() where T : unmanaged
		{
			var typeId = GetTypeId<T>();
			if (ComponentIndex.TryGetValue(typeId, out var componentStorage))
			{
				return componentStorage;
			}

			componentStorage = new ComponentStorage(typeId, ElementSizes[typeId]);
			ComponentIndex.Add(typeId, componentStorage);
			TypeToFilter.Add(typeId, new List<Filter>());
			return componentStorage;
		}

		// FILTERS

		internal Filter GetFilter(FilterSignature signature)
		{
			if (!FilterIndex.TryGetValue(signature, out var filter))
			{
				filter = new Filter(this, signature);

				foreach (var typeId in signature.Included)
				{
					TypeToFilter[typeId].Add(filter);
				}

				foreach (var typeId in signature.Excluded)
				{
					TypeToFilter[typeId].Add(filter);
				}

				FilterIndex.Add(signature, filter);
			}

			return filter;
		}

		// ENTITIES

		public Entity CreateEntity(string tag = "")
		{
			var entity = new Entity(EntityIdAssigner.Assign());

			if (!EntityComponentIndex.ContainsKey(entity))
			{
				EntityRelationIndex.Add(entity, new IndexableSet<TypeId>());
				EntityComponentIndex.Add(entity, new IndexableSet<TypeId>());
			}

			EntityTags[entity] = tag;

			return entity;
		}

		public void Tag(Entity entity, string tag)
		{
			EntityTags[entity] = tag;
		}

		public string GetTag(Entity entity)
		{
			return EntityTags[entity];
		}

		public void Destroy(in Entity entity)
		{
			// remove all components from storages
			foreach (var componentTypeIndex in EntityComponentIndex[entity])
			{
				var componentStorage = ComponentIndex[componentTypeIndex];
				componentStorage.Remove(entity);

				foreach (var filter in TypeToFilter[componentTypeIndex])
				{
					filter.RemoveEntity(entity);
				}
			}

			// remove all relations from storage
			foreach (var relationTypeIndex in EntityRelationIndex[entity])
			{
				var relationStorage = RelationIndex[relationTypeIndex];
				relationStorage.RemoveEntity(entity);
			}

			EntityComponentIndex[entity].Clear();
			EntityRelationIndex[entity].Clear();

			// recycle ID
			EntityIdAssigner.Unassign(entity.ID);
		}

		// COMPONENTS

		public bool Has<T>(in Entity entity) where T : unmanaged
		{
			var storage = GetComponentStorage<T>();
			return storage.Has(entity);
		}

		internal bool Has(in Entity entity, in TypeId typeId)
		{
			return EntityComponentIndex[entity].Contains(typeId);
		}

		public bool Some<T>() where T : unmanaged
		{
			var storage = GetComponentStorage<T>();
			return storage.Any();
		}

		public ref T Get<T>(in Entity entity) where T : unmanaged
		{
			var storage = GetComponentStorage<T>();
			return ref storage.Get<T>(entity);
		}

		public ref T GetSingleton<T>() where T : unmanaged
		{
			var storage = GetComponentStorage<T>();
			return ref storage.GetFirst<T>();
		}

		public Entity GetSingletonEntity<T>() where T : unmanaged
		{
			var storage = GetComponentStorage<T>();
			return storage.FirstEntity();
		}

		public void Set<T>(in Entity entity, in T component) where T : unmanaged
		{
			var componentStorage = GetComponentStorage<T>();

			if (!componentStorage.Set(entity, component))
			{
				EntityComponentIndex[entity].Add(componentStorage.TypeId);

				foreach (var filter in TypeToFilter[componentStorage.TypeId])
				{
					filter.Check(entity);
				}
			}
		}

		public void Remove<T>(in Entity entity) where T : unmanaged
		{
			var componentStorage = GetComponentStorage<T>();

			if (componentStorage.Remove(entity))
			{
				EntityComponentIndex[entity].Remove(componentStorage.TypeId);

				foreach (var filter in TypeToFilter[componentStorage.TypeId])
				{
					filter.Check(entity);
				}
			}
		}

		// RELATIONS

		private RelationStorage RegisterRelationType(TypeId typeId)
		{
			var relationStorage = new RelationStorage(ElementSizes[typeId]);
			RelationIndex.Add(typeId, relationStorage);
			return relationStorage;
		}

		[MethodImpl(MethodImplOptions.AggressiveInlining)]
		private RelationStorage GetRelationStorage<T>() where T : unmanaged
		{
			var typeId = GetTypeId<T>();
			if (RelationIndex.TryGetValue(typeId, out var relationStorage))
			{
				return relationStorage;
			}

			return RegisterRelationType(typeId);
		}

		public void Relate<T>(in Entity entityA, in Entity entityB, in T relation) where T : unmanaged
		{
			var relationStorage = GetRelationStorage<T>();
			relationStorage.Set(entityA, entityB, relation);
			EntityRelationIndex[entityA].Add(TypeToId[typeof(T)]);
			EntityRelationIndex[entityB].Add(TypeToId[typeof(T)]);
		}

		public void Unrelate<T>(in Entity entityA, in Entity entityB) where T : unmanaged
		{
			var relationStorage = GetRelationStorage<T>();
			relationStorage.Remove(entityA, entityB);
		}

		public void UnrelateAll<T>(in Entity entity) where T : unmanaged
		{
			var relationStorage = GetRelationStorage<T>();
			relationStorage.RemoveEntity(entity);
		}

		public bool Related<T>(in Entity entityA, in Entity entityB) where T : unmanaged
		{
			var relationStorage = GetRelationStorage<T>();
			return relationStorage.Has(entityA, entityB);
		}

		public T GetRelationData<T>(in Entity entityA, in Entity entityB) where T : unmanaged
		{
			var relationStorage = GetRelationStorage<T>();
			return relationStorage.Get<T>(entityA, entityB);
		}

		public ReverseSpanEnumerator<(Entity, Entity)> Relations<T>() where T : unmanaged
		{
			var relationStorage = GetRelationStorage<T>();
			return relationStorage.All();
		}

		public ReverseSpanEnumerator<Entity> OutRelations<T>(Entity entity) where T : unmanaged
		{
			var relationStorage = GetRelationStorage<T>();
			return relationStorage.OutRelations(entity);
		}

		public Entity OutRelationSingleton<T>(in Entity entity) where T : unmanaged
		{
			var relationStorage = GetRelationStorage<T>();
			return relationStorage.OutFirst(entity);
		}

		public bool HasOutRelation<T>(in Entity entity) where T : unmanaged
		{
			var relationStorage = GetRelationStorage<T>();
			return relationStorage.HasOutRelation(entity);
		}

		public int OutRelationCount<T>(in Entity entity) where T : unmanaged
		{
			var relationStorage = GetRelationStorage<T>();
			return relationStorage.OutRelationCount(entity);
		}

		public Entity NthOutRelation<T>(in Entity entity, int n) where T : unmanaged
		{
			var relationStorage = GetRelationStorage<T>();
			return relationStorage.OutNth(entity, n);
		}

		public ReverseSpanEnumerator<Entity> InRelations<T>(Entity entity) where T : unmanaged
		{
			var relationStorage = GetRelationStorage<T>();
			return relationStorage.InRelations(entity);
		}

		public Entity InRelationSingleton<T>(in Entity entity) where T : unmanaged
		{
			var relationStorage = GetRelationStorage<T>();
			return relationStorage.InFirst(entity);
		}

		public bool HasInRelation<T>(in Entity entity) where T : unmanaged
		{
			var relationStorage = GetRelationStorage<T>();
			return relationStorage.HasInRelation(entity);
		}

		public int InRelationCount<T>(in Entity entity) where T : unmanaged
		{
			var relationStorage = GetRelationStorage<T>();
			return relationStorage.InRelationCount(entity);
		}

		public Entity NthInRelation<T>(in Entity entity, int n) where T : unmanaged
		{
			var relationStorage = GetRelationStorage<T>();
			return relationStorage.InNth(entity, n);
		}

		// MESSAGES

		private TypeId GetMessageTypeId<T>() where T : unmanaged
		{
			var typeId = GetTypeId<T>();

			if (!MessageIndex.ContainsKey(typeId))
			{
				MessageIndex.Add(typeId, new MessageStorage(Unsafe.SizeOf<T>()));
			}

			return typeId;
		}

		public void Send<T>(in T message) where T : unmanaged
		{
			var typeId = GetMessageTypeId<T>();
			MessageIndex[typeId].Add(message);
		}

		public bool SomeMessage<T>() where T : unmanaged
		{
			var typeId = GetMessageTypeId<T>();
			return MessageIndex[typeId].Some();
		}

		public ReadOnlySpan<T> ReadMessages<T>() where T : unmanaged
		{
			var typeId = GetMessageTypeId<T>();
			return MessageIndex[typeId].All<T>();
		}

		public T ReadMessage<T>() where T : unmanaged
		{
			var typeId = GetMessageTypeId<T>();
			return MessageIndex[typeId].First<T>();
		}

		public void ClearMessages<T>() where T : unmanaged
		{
			var typeId = GetMessageTypeId<T>();
			MessageIndex[typeId].Clear();
		}

		// TODO: temporary component storage?
		public void FinishUpdate()
		{
			foreach (var (_, messageStorage) in MessageIndex)
			{
				messageStorage.Clear();
			}
		}

		// DEBUG
		// NOTE: these methods are very inefficient
		// they should only be used in debugging contexts!!
		#if DEBUG
		public ComponentTypeEnumerator Debug_GetAllComponentTypes(Entity entity)
		{
			return new ComponentTypeEnumerator(this, EntityComponentIndex[entity]);
		}

		public IEnumerable<Entity> Debug_GetEntities(Type componentType)
		{
			var storage = ComponentIndex[TypeToId[componentType]];
			return storage.Debug_GetEntities();
		}

		public IEnumerable<Type> Debug_SearchComponentType(string typeString)
		{
			foreach (var type in TypeToId.Keys)
			{
				if (type.ToString().ToLower().Contains(typeString.ToLower()))
				{
					yield return type;
				}
			}
		}

		public ref struct ComponentTypeEnumerator
		{
			private World World;
			private IndexableSet<TypeId> Types;
			private int ComponentIndex;

			public ComponentTypeEnumerator GetEnumerator() => this;

			internal ComponentTypeEnumerator(
				World world,
				IndexableSet<TypeId> types
			)
			{
				World = world;
				Types = types;
				ComponentIndex = -1;
			}

			public bool MoveNext()
			{
				ComponentIndex += 1;
				return ComponentIndex < Types.Count;
			}

			public unsafe Type Current => World.IdToType[Types[ComponentIndex]];
		}
		#endif

		protected virtual void Dispose(bool disposing)
		{
			if (!IsDisposed)
			{
				if (disposing)
				{
					foreach (var componentStorage in ComponentIndex.Values)
					{
						componentStorage.Dispose();
					}

					foreach (var relationStorage in RelationIndex.Values)
					{
						relationStorage.Dispose();
					}

					foreach (var messageStorage in MessageIndex.Values)
					{
						messageStorage.Dispose();
					}

					foreach (var typeSet in EntityComponentIndex.Values)
					{
						typeSet.Dispose();
					}

					foreach (var typeSet in EntityRelationIndex.Values)
					{
						typeSet.Dispose();
					}

					foreach (var filter in FilterIndex.Values)
					{
						filter.Dispose();
					}

					EntityIdAssigner.Dispose();
					TypeIdAssigner.Dispose();
				}

				IsDisposed = true;
			}
		}

		// // TODO: override finalizer only if 'Dispose(bool disposing)' has code to free unmanaged resources
		// ~World()
		// {
		//     // Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method
		//     Dispose(disposing: false);
		// }

		public void Dispose()
		{
			// Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method
			Dispose(disposing: true);
			GC.SuppressFinalize(this);
		}
	}
}