MoonTools.ECS/src/RelationStorage.cs

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

namespace MoonTools.ECS;

// TODO: implement this entire class with NativeMemory equivalents, can just memcpy for snapshots
internal class RelationStorage
{
	internal NativeArray relations;
	internal NativeArray relationDatas;
	internal Dictionary<(Entity, Entity), int> indices = new Dictionary<(Entity, Entity), int>(16);
	internal Dictionary<Entity, IndexableSet<Entity>> outRelations = new Dictionary<Entity, IndexableSet<Entity>>(16);
	internal Dictionary<Entity, IndexableSet<Entity>> inRelations = new Dictionary<Entity, IndexableSet<Entity>>(16);
	private Stack<IndexableSet<Entity>> listPool = new Stack<IndexableSet<Entity>>();

	private bool disposed;

	public RelationStorage(int relationDataSize)
	{
		relations = new NativeArray(Unsafe.SizeOf<(Entity, Entity)>());
		relationDatas = new NativeArray(relationDataSize);
	}

	public ReverseSpanEnumerator<(Entity, Entity)> All()
	{
		return new ReverseSpanEnumerator<(Entity, Entity)>(relations.ToSpan<(Entity, Entity)>());
	}

	public unsafe void Set<T>(in Entity entityA, in Entity entityB, in T relationData) where T : unmanaged
	{
		var relation = (entityA, entityB);

		if (indices.TryGetValue(relation, out var index))
		{
			relationDatas.Set(index, relationData);
			return;
		}

		if (!outRelations.ContainsKey(entityA))
		{
			outRelations[entityA] = AcquireHashSetFromPool();
		}
		outRelations[entityA].Add(entityB);

		if (!inRelations.ContainsKey(entityB))
		{
			inRelations[entityB] = AcquireHashSetFromPool();
		}
		inRelations[entityB].Add(entityA);

		relations.Append(relation);
		relationDatas.Append(relationData);
		indices.Add(relation, relations.Count - 1);
	}

	public ref T Get<T>(in Entity entityA, in Entity entityB) where T : unmanaged
	{
		var relationIndex = indices[(entityA, entityB)];
		return ref relationDatas.Get<T>(relationIndex);
	}

	public bool Has(Entity entityA, Entity entityB)
	{
		return indices.ContainsKey((entityA, entityB));
	}

	public ReverseSpanEnumerator<Entity> OutRelations(Entity Entity)
	{
		if (outRelations.TryGetValue(Entity, out var entityOutRelations))
		{
			return entityOutRelations.GetEnumerator();
		}
		else
		{
			return ReverseSpanEnumerator<Entity>.Empty;
		}
	}

	public Entity OutFirst(Entity Entity)
	{
		return OutNth(Entity, 0);
	}

	public Entity OutNth(Entity Entity, int n)
	{
#if DEBUG
		if (!outRelations.ContainsKey(Entity) || outRelations[Entity].Count == 0)
		{
			throw new KeyNotFoundException("No out relations to this entity!");
		}
#endif
		return outRelations[Entity][n];
	}

	public bool HasOutRelation(Entity Entity)
	{
		return outRelations.ContainsKey(Entity) && outRelations[Entity].Count > 0;
	}

	public int OutRelationCount(Entity Entity)
	{
		return outRelations.TryGetValue(Entity, out var entityOutRelations) ? entityOutRelations.Count : 0;
	}

	public ReverseSpanEnumerator<Entity> InRelations(Entity Entity)
	{
		if (inRelations.TryGetValue(Entity, out var entityInRelations))
		{
			return entityInRelations.GetEnumerator();
		}
		else
		{
			return ReverseSpanEnumerator<Entity>.Empty;
		}
	}

	public Entity InFirst(Entity Entity)
	{
		return InNth(Entity, 0);
	}

	public Entity InNth(Entity Entity, int n)
	{
#if DEBUG
		if (!inRelations.ContainsKey(Entity) || inRelations[Entity].Count == 0)
		{
			throw new KeyNotFoundException("No in relations to this entity!");
		}
#endif

		return inRelations[Entity][n];
	}

	public bool HasInRelation(Entity Entity)
	{
		return inRelations.ContainsKey(Entity) && inRelations[Entity].Count > 0;
	}

	public int InRelationCount(Entity Entity)
	{
		return inRelations.TryGetValue(Entity, out var entityInRelations) ? entityInRelations.Count : 0;
	}

	public (bool, bool) Remove(in Entity entityA, in Entity entityB)
	{
		var aEmpty = false;
		var bEmpty = false;
		var relation = (entityA, entityB);

		if (outRelations.TryGetValue(entityA, out var entityOutRelations))
		{
			entityOutRelations.Remove(entityB);
			if (outRelations[entityA].Count == 0)
			{
				aEmpty = true;
			}
		}

		if (inRelations.TryGetValue(entityB, out var entityInRelations))
		{
			entityInRelations.Remove(entityA);
			if (inRelations[entityB].Count == 0)
			{
				bEmpty = true;
			}
		}

		if (indices.TryGetValue(relation, out var index))
		{
			var lastElementIndex = relations.Count - 1;

			relationDatas.Delete(index);
			relations.Delete(index);

			// move an element into the hole
			if (index != lastElementIndex)
			{
				var lastRelation = relations.Get<(Entity, Entity)>(lastElementIndex);
				indices[lastRelation] = index;
			}

			indices.Remove(relation);
		}

		return (aEmpty, bEmpty);
	}

	public void RemoveEntity(in Entity entity)
	{
		if (outRelations.TryGetValue(entity, out var entityOutRelations))
		{
			foreach (var entityB in entityOutRelations)
			{
				Remove(entity, entityB);
			}

			ReturnHashSetToPool(entityOutRelations);
			outRelations.Remove(entity);
		}

		if (inRelations.TryGetValue(entity, out var entityInRelations))
		{
			foreach (var entityA in entityInRelations)
			{
				Remove(entityA, entity);
			}

			ReturnHashSetToPool(entityInRelations);
			inRelations.Remove(entity);
		}
	}

	internal IndexableSet<Entity> AcquireHashSetFromPool()
	{
		if (listPool.Count == 0)
		{
			listPool.Push(new IndexableSet<Entity>());
		}

		return listPool.Pop();
	}

	private void ReturnHashSetToPool(IndexableSet<Entity> hashSet)
	{
		hashSet.Clear();
		listPool.Push(hashSet);
	}

	public void Clear()
	{
		indices.Clear();

		foreach (var set in inRelations.Values)
		{
			ReturnHashSetToPool(set);
		}
		inRelations.Clear();

		foreach (var set in outRelations.Values)
		{
			ReturnHashSetToPool(set);
		}
		outRelations.Clear();

		relations.Clear();
		relationDatas.Clear();
	}

	protected virtual void Dispose(bool disposing)
	{
		if (!disposed)
		{
			Clear();

			if (disposing)
			{
				foreach (var set in listPool)
				{
					set.Dispose();
				}

				relations.Dispose();
				relationDatas.Dispose();
			}

			disposed = true;
		}
	}

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