From 9eb4b947f0fee2b196131a3b1f1a9a5ed4dc3657 Mon Sep 17 00:00:00 2001
From: cosmonaut <evan@moonside.games>
Date: Fri, 18 Mar 2022 12:50:59 -0700
Subject: [PATCH] WithEntity message getters

---
 src/MessageDepot.cs   | 19 +++++++++++++++++--
 src/MessageStorage.cs | 37 +++++++++++++++++++++++++++++++++++++
 src/System.cs         | 19 +++++++++++++++++--
 3 files changed, 71 insertions(+), 4 deletions(-)

diff --git a/src/MessageDepot.cs b/src/MessageDepot.cs
index ef91f6d..28e8cc9 100644
--- a/src/MessageDepot.cs
+++ b/src/MessageDepot.cs
@@ -24,16 +24,31 @@ internal class MessageDepot
 		return Lookup<TMessage>().Some();
 	}
 
-	public ReadOnlySpan<TMessage> ReadAll<TMessage>() where TMessage : struct
+	public ReadOnlySpan<TMessage> All<TMessage>() where TMessage : struct
 	{
 		return Lookup<TMessage>().All();
 	}
 
-	public TMessage ReadFirst<TMessage>() where TMessage : struct
+	public TMessage First<TMessage>() where TMessage : struct
 	{
 		return Lookup<TMessage>().First();
 	}
 
+	public IEnumerable<TMessage> WithEntity<TMessage>(int entityID) where TMessage : struct, IHasEntity
+	{
+		return Lookup<TMessage>().WithEntity(entityID);
+	}
+
+	public ref readonly TMessage FirstWithEntity<TMessage>(int entityID) where TMessage : struct, IHasEntity
+	{
+		return ref Lookup<TMessage>().FirstWithEntity(entityID);
+	}
+
+	public bool SomeWithEntity<TMessage>(int entityID) where TMessage : struct, IHasEntity
+	{
+		return Lookup<TMessage>().SomeWithEntity(entityID);
+	}
+
 	public void Clear()
 	{
 		foreach (var storage in storages.Values)
diff --git a/src/MessageStorage.cs b/src/MessageStorage.cs
index 8bfe24f..9441500 100644
--- a/src/MessageStorage.cs
+++ b/src/MessageStorage.cs
@@ -10,6 +10,7 @@ internal class MessageStorage<TMessage> : MessageStorage where TMessage : struct
 	private int count = 0;
 	private int capacity = 128;
 	private TMessage[] messages;
+	private Dictionary<int, List<int>> entityToIndices = new Dictionary<int, List<int>>();
 
 	public MessageStorage()
 	{
@@ -25,6 +26,17 @@ internal class MessageStorage<TMessage> : MessageStorage where TMessage : struct
 		}
 
 		messages[count] = message;
+
+		if (message is IHasEntity entityMessage)
+		{
+			if (!entityToIndices.ContainsKey(entityMessage.Entity.ID))
+			{
+				entityToIndices.Add(entityMessage.Entity.ID, new List<int>());
+			}
+
+			entityToIndices[entityMessage.Entity.ID].Add(count);
+		}
+
 		count += 1;
 	}
 
@@ -43,8 +55,33 @@ internal class MessageStorage<TMessage> : MessageStorage where TMessage : struct
 		return messages[0];
 	}
 
+	public IEnumerable<TMessage> WithEntity(int entityID)
+	{
+		if (entityToIndices.ContainsKey(entityID))
+		{
+			foreach (var index in entityToIndices[entityID])
+			{
+				yield return messages[index];
+			}
+		}
+	}
+
+	public ref readonly TMessage FirstWithEntity(int entityID)
+	{
+		return ref messages[entityToIndices[entityID][0]];
+	}
+
+	public bool SomeWithEntity(int entityID)
+	{
+		return entityToIndices.ContainsKey(entityID) && entityToIndices[entityID].Count > 0;
+	}
+
 	public override void Clear()
 	{
 		count = 0;
+		foreach (var set in entityToIndices.Values)
+		{
+			set.Clear();
+		}
 	}
 }
diff --git a/src/System.cs b/src/System.cs
index 1c166a1..c58c4e1 100644
--- a/src/System.cs
+++ b/src/System.cs
@@ -38,12 +38,12 @@ public abstract class System : EntityComponentReader
 
 	protected ReadOnlySpan<TMessage> ReadMessages<TMessage>() where TMessage : struct
 	{
-		return MessageDepot.ReadAll<TMessage>();
+		return MessageDepot.All<TMessage>();
 	}
 
 	protected TMessage ReadMessage<TMessage>() where TMessage : struct
 	{
-		return MessageDepot.ReadFirst<TMessage>();
+		return MessageDepot.First<TMessage>();
 	}
 
 	protected bool SomeMessage<TMessage>() where TMessage : struct
@@ -51,6 +51,21 @@ public abstract class System : EntityComponentReader
 		return MessageDepot.Some<TMessage>();
 	}
 
+	protected IEnumerable<TMessage> ReadMessagesWithEntity<TMessage>(in Entity entity) where TMessage : struct, IHasEntity
+	{
+		return MessageDepot.WithEntity<TMessage>(entity.ID);
+	}
+
+	protected ref readonly TMessage ReadMessageWithEntity<TMessage>(in Entity entity) where TMessage : struct, IHasEntity
+	{
+		return ref MessageDepot.FirstWithEntity<TMessage>(entity.ID);
+	}
+
+	protected bool SomeMessageWithEntity<TMessage>(in Entity entity) where TMessage : struct, IHasEntity
+	{
+		return MessageDepot.SomeWithEntity<TMessage>(entity.ID);
+	}
+
 	protected void Send<TMessage>(in TMessage message) where TMessage : struct
 	{
 		MessageDepot.Add(message);