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
	{
		// 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>();

		// Archetypes
		internal readonly Dictionary<ArchetypeSignature, Archetype> ArchetypeIndex = new Dictionary<ArchetypeSignature, Archetype>();
		internal readonly Dictionary<Entity, ArchetypeRecord> EntityIndex = new Dictionary<Entity, ArchetypeRecord>();
		internal readonly Archetype EmptyArchetype;

		// 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 IdAssigner EntityIdAssigner = new IdAssigner();
		private IdAssigner TypeIdAssigner = new IdAssigner();

		public World()
		{
			EmptyArchetype = CreateArchetype(ArchetypeSignature.Empty);
		}

		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;
		}

		[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(ElementSizes[typeId]);
			ComponentIndex.Add(typeId, componentStorage);
			return componentStorage;
		}

		private Archetype CreateArchetype(ArchetypeSignature signature)
		{
			var archetype = new Archetype(this, signature);
			ArchetypeIndex.Add(signature, archetype);
			return archetype;
		}

		// ENTITIES

		public Entity CreateEntity(string tag = "")
		{
			var entity = new Entity(EntityIdAssigner.Assign());
			EntityIndex.Add(entity, new ArchetypeRecord(EmptyArchetype, EmptyArchetype.Count));
			EmptyArchetype.Append(entity);
			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)
		{
			var record = EntityIndex[entity];
			var archetype = record.Archetype;
			var row = record.Row;

			// remove all components from storages
			for (int i = 0; i < archetype.Signature.Count; i += 1)
			{
				var componentStorage = ComponentIndex[archetype.Signature[i]];
				componentStorage.Remove(entity);
			}

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

			EntityRelationIndex[entity].Clear();

			// remove from archetype
			if (row != archetype.Count - 1)
			{
				var lastEntity = archetype.Entities[archetype.Count - 1];
				archetype.Entities[row] = lastEntity;
				EntityIndex[lastEntity] = new ArchetypeRecord(archetype, row);
			}
			archetype.Entities.RemoveLastElement();
			EntityIndex.Remove(entity);

			// 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);
		}

		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))
			{
				UpdateArchetype<T>(entity, true);
			}
		}

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

			if (componentStorage.Remove(entity))
			{
				UpdateArchetype<T>(entity, false);
			}
		}

		private void UpdateArchetype<T>(in Entity entity, bool insert)
		{
			Archetype? nextArchetype;

			var componentTypeId = TypeToId[typeof(T)];
			var record = EntityIndex[entity];
			var archetype = record.Archetype;

			if (archetype.Edges.TryGetValue(TypeToId[typeof(T)], out var edge))
			{
				nextArchetype = insert ? edge.Add : edge.Remove;
			}
			else
			{
				var nextSignature = new ArchetypeSignature(archetype.Signature.Count + 1);
				archetype.Signature.CopyTo(nextSignature);

				if (insert)
				{
					nextSignature.Insert(componentTypeId);
				}
				else
				{
					nextSignature.Remove(componentTypeId);
				}

				if (!ArchetypeIndex.TryGetValue(nextSignature, out nextArchetype))
				{
					nextArchetype = CreateArchetype(nextSignature);
				}

				var newEdge = new ArchetypeEdge(nextArchetype, archetype);
				archetype.Edges.Add(componentTypeId, newEdge);
				nextArchetype.Edges.Add(componentTypeId, newEdge);
			}

			var newRow = archetype.Transfer(record.Row, nextArchetype);
			EntityIndex[entity] = new ArchetypeRecord(nextArchetype, newRow);
		}

		// 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>();
		}

		// 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, EntityIndex[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 ArchetypeRecord Record;
			private int ComponentIndex;

			public ComponentTypeEnumerator GetEnumerator() => this;

			internal ComponentTypeEnumerator(
				World world,
				ArchetypeRecord record
			)
			{
				World = world;
				Record = record;
				ComponentIndex = -1;
			}

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

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