From da35e992661d293bc85e68fb001fb5f2068136da Mon Sep 17 00:00:00 2001
From: cosmonaut <evan@moonside.games>
Date: Wed, 6 Apr 2022 12:53:50 -0700
Subject: [PATCH] entity relation system

---
 src/EntityComponentReader.cs |  21 ++++++
 src/Relation.cs              |  35 ++++++++++
 src/RelationDepot.cs         |  50 ++++++++++++++
 src/RelationStorage.cs       | 122 +++++++++++++++++++++++++++++++++++
 src/System.cs                |  11 ++++
 src/World.cs                 |  12 +++-
 6 files changed, 249 insertions(+), 2 deletions(-)
 create mode 100644 src/Relation.cs
 create mode 100644 src/RelationDepot.cs
 create mode 100644 src/RelationStorage.cs

diff --git a/src/EntityComponentReader.cs b/src/EntityComponentReader.cs
index 4146d94..2e3d981 100644
--- a/src/EntityComponentReader.cs
+++ b/src/EntityComponentReader.cs
@@ -4,6 +4,7 @@ public abstract class EntityComponentReader
 {
 	internal EntityStorage EntityStorage;
 	internal ComponentDepot ComponentDepot;
+	internal RelationDepot RelationDepot;
 	protected FilterBuilder FilterBuilder => new FilterBuilder(ComponentDepot);
 
 	internal void RegisterEntityStorage(EntityStorage entityStorage)
@@ -16,6 +17,11 @@ public abstract class EntityComponentReader
 		ComponentDepot = componentDepot;
 	}
 
+	internal void RegisterRelationDepot(RelationDepot relationDepot)
+	{
+		RelationDepot = relationDepot;
+	}
+
 	protected ReadOnlySpan<TComponent> ReadComponents<TComponent>() where TComponent : struct
 	{
 		return ComponentDepot.ReadComponents<TComponent>();
@@ -50,4 +56,19 @@ public abstract class EntityComponentReader
 	{
 		return EntityStorage.Exists(entity);
 	}
+
+	protected IEnumerable<Relation> Relations<TRelationKind>()
+	{
+		return RelationDepot.Relations<TRelationKind>();
+	}
+
+	protected IEnumerable<Entity> RelatedToA<TRelationKind>(in Entity entity)
+	{
+		return RelationDepot.RelatedToA<TRelationKind>(entity.ID);
+	}
+
+	protected IEnumerable<Entity> RelatedToB<TRelationKind>(in Entity entity)
+	{
+		return RelationDepot.RelatedToB<TRelationKind>(entity.ID);
+	}
 }
diff --git a/src/Relation.cs b/src/Relation.cs
new file mode 100644
index 0000000..2023801
--- /dev/null
+++ b/src/Relation.cs
@@ -0,0 +1,35 @@
+namespace MoonTools.ECS
+{
+	public struct Relation : IEquatable<Relation>
+	{
+		public Entity A { get; }
+		public Entity B { get; }
+
+		internal Relation(Entity entityA, Entity entityB)
+		{
+			A = entityA;
+			B = entityB;
+		}
+
+		internal Relation(int idA, int idB)
+		{
+			A = new Entity(idA);
+			B = new Entity(idB);
+		}
+
+		public override bool Equals(object? obj)
+		{
+			return obj is Relation relation && Equals(relation);
+		}
+
+		public bool Equals(Relation other)
+		{
+			return A.Equals(other.A) && B.Equals(other.B);
+		}
+
+		public override int GetHashCode()
+		{
+			return HashCode.Combine(A, B);
+		}
+	}
+}
diff --git a/src/RelationDepot.cs b/src/RelationDepot.cs
new file mode 100644
index 0000000..a293864
--- /dev/null
+++ b/src/RelationDepot.cs
@@ -0,0 +1,50 @@
+namespace MoonTools.ECS
+{
+	internal class RelationDepot
+	{
+		private Dictionary<Type, RelationStorage> storages = new Dictionary<Type, RelationStorage>();
+
+		private RelationStorage Lookup<TRelationKind>()
+		{
+			return storages[typeof(TRelationKind)];
+		}
+
+		public void Register<TRelationKind>()
+		{
+			storages[typeof(TRelationKind)] = new RelationStorage();
+		}
+
+		public void Add<TRelationKind>(Relation relation)
+		{
+			Lookup<TRelationKind>().Add(relation);
+		}
+
+		public void Remove<TRelationKind>(Relation relation)
+		{
+			Lookup<TRelationKind>().Remove(relation);
+		}
+
+		public void OnEntityDestroy(int entityID)
+		{
+			foreach (var storage in storages.Values)
+			{
+				storage.OnEntityDestroy(entityID);
+			}
+		}
+
+		public IEnumerable<Relation> Relations<TRelationKind>()
+		{
+			return Lookup<TRelationKind>().All();
+		}
+
+		public IEnumerable<Entity> RelatedToA<TRelationKind>(int entityID)
+		{
+			return Lookup<TRelationKind>().RelatedToA(entityID);
+		}
+
+		public IEnumerable<Entity> RelatedToB<TRelationKind>(int entityID)
+		{
+			return Lookup<TRelationKind>().RelatedToB(entityID);
+		}
+	}
+}
diff --git a/src/RelationStorage.cs b/src/RelationStorage.cs
new file mode 100644
index 0000000..9df0a35
--- /dev/null
+++ b/src/RelationStorage.cs
@@ -0,0 +1,122 @@
+namespace MoonTools.ECS
+{
+	internal class RelationStorage
+	{
+		private HashSet<Relation> relations = new HashSet<Relation>(16);
+		private Dictionary<int, HashSet<int>> entitiesRelatedToA = new Dictionary<int, HashSet<int>>(16);
+		private Dictionary<int, HashSet<int>> entitiesRelatedToB = new Dictionary<int, HashSet<int>>(16);
+		private Stack<HashSet<int>> listPool = new Stack<HashSet<int>>();
+
+		public IEnumerable<Relation> All()
+		{
+			foreach (var relation in relations)
+			{
+				yield return relation;
+			}
+		}
+
+		public void Add(Relation relation)
+		{
+			if (relations.Contains(relation)) { return; }
+			var idA = relation.A.ID;
+			var idB = relation.B.ID;
+
+			if (!entitiesRelatedToA.ContainsKey(idA))
+			{
+				entitiesRelatedToA[idA] = AcquireHashSetFromPool();
+			}
+			entitiesRelatedToA[idA].Add(idB);
+
+			if (!entitiesRelatedToB.ContainsKey(idB))
+			{
+				entitiesRelatedToB[idB] = AcquireHashSetFromPool();
+			}
+			entitiesRelatedToB[idB].Add(idA);
+
+			relations.Add(relation);
+		}
+
+		public IEnumerable<Entity> RelatedToA(int entityID)
+		{
+			if (entitiesRelatedToA.ContainsKey(entityID))
+			{
+				foreach (var id in entitiesRelatedToA[entityID])
+				{
+					yield return new Entity(id);
+				}
+			}
+		}
+
+		public IEnumerable<Entity> RelatedToB(int entityID)
+		{
+			if (entitiesRelatedToB.ContainsKey(entityID))
+			{
+				foreach (var id in entitiesRelatedToB[entityID])
+				{
+					yield return new Entity(id);
+				}
+			}
+		}
+
+		public bool Remove(Relation relation)
+		{
+			if (entitiesRelatedToA.ContainsKey(relation.A.ID))
+			{
+				entitiesRelatedToA[relation.A.ID].Remove(relation.B.ID);
+			}
+
+			if (entitiesRelatedToB.ContainsKey(relation.B.ID))
+			{
+				entitiesRelatedToB[relation.B.ID].Remove(relation.A.ID);
+			}
+
+			return relations.Remove(relation);
+		}
+
+		// this exists so we don't recurse in OnEntityDestroy
+		private bool DestroyRemove(Relation relation)
+		{
+			return relations.Remove(relation);
+		}
+
+		public void OnEntityDestroy(int entityID)
+		{
+			if (entitiesRelatedToA.ContainsKey(entityID))
+			{
+				foreach (var entityB in entitiesRelatedToA[entityID])
+				{
+					DestroyRemove(new Relation(entityID, entityB));
+				}
+
+				ReturnHashSetToPool(entitiesRelatedToA[entityID]);
+				entitiesRelatedToA.Remove(entityID);
+			}
+
+			if (entitiesRelatedToB.ContainsKey(entityID))
+			{
+				foreach (var entityA in entitiesRelatedToB[entityID])
+				{
+					DestroyRemove(new Relation(entityA, entityID));
+				}
+
+				ReturnHashSetToPool(entitiesRelatedToB[entityID]);
+				entitiesRelatedToB.Remove(entityID);
+			}
+		}
+
+		private HashSet<int> AcquireHashSetFromPool()
+		{
+			if (listPool.Count == 0)
+			{
+				listPool.Push(new HashSet<int>());
+			}
+
+			return listPool.Pop();
+		}
+
+		private void ReturnHashSetToPool(HashSet<int> hashSet)
+		{
+			listPool.Push(hashSet);
+		}
+	}
+}
diff --git a/src/System.cs b/src/System.cs
index c58c4e1..3cb11cd 100644
--- a/src/System.cs
+++ b/src/System.cs
@@ -71,9 +71,20 @@ public abstract class System : EntityComponentReader
 		MessageDepot.Add(message);
 	}
 
+	protected void Relate<TRelationKind>(in Entity entityA, in Entity entityB)
+	{
+		RelationDepot.Add<TRelationKind>(new Relation(entityA, entityB));
+	}
+
+	protected void Unrelate<TRelationKind>(in Entity entityA, in Entity entityB)
+	{
+		RelationDepot.Remove<TRelationKind>(new Relation(entityA, entityB));
+	}
+
 	protected void Destroy(in Entity entity)
 	{
 		ComponentDepot.OnEntityDestroy(entity.ID);
+		RelationDepot.OnEntityDestroy(entity.ID);
 		EntityStorage.Destroy(entity);
 	}
 }
diff --git a/src/World.cs b/src/World.cs
index a05abc2..dc497c2 100644
--- a/src/World.cs
+++ b/src/World.cs
@@ -1,22 +1,30 @@
-namespace MoonTools.ECS;
+namespace MoonTools.ECS;
 
 public class World
 {
 	private readonly EntityStorage EntityStorage = new EntityStorage();
 	private readonly ComponentDepot ComponentDepot = new ComponentDepot();
-	private MessageDepot MessageDepot = new MessageDepot();
+	private readonly MessageDepot MessageDepot = new MessageDepot();
+	private readonly RelationDepot RelationDepot = new RelationDepot();
 
 	internal void AddSystem(System system)
 	{
 		system.RegisterEntityStorage(EntityStorage);
 		system.RegisterComponentDepot(ComponentDepot);
 		system.RegisterMessageDepot(MessageDepot);
+		system.RegisterRelationDepot(RelationDepot);
 	}
 
 	internal void AddRenderer(Renderer renderer)
 	{
 		renderer.RegisterEntityStorage(EntityStorage);
 		renderer.RegisterComponentDepot(ComponentDepot);
+		renderer.RegisterRelationDepot(RelationDepot);
+	}
+
+	public void AddRelationKind<TRelationKind>()
+	{
+		RelationDepot.Register<TRelationKind>();
 	}
 
 	public Entity CreateEntity()