From 227aae9a1a6fb0096580b0dd0c4d1787351fbc6a Mon Sep 17 00:00:00 2001
From: cosmonaut <evan@moonside.games>
Date: Mon, 7 Mar 2022 12:05:06 -0800
Subject: [PATCH] register filter components with storage

---
 src/ComponentDepot.cs | 34 +++++++++++++++++-----------------
 src/FilterBuilder.cs  |  2 ++
 2 files changed, 19 insertions(+), 17 deletions(-)

diff --git a/src/ComponentDepot.cs b/src/ComponentDepot.cs
index 18057fe..3afb8b6 100644
--- a/src/ComponentDepot.cs
+++ b/src/ComponentDepot.cs
@@ -10,6 +10,14 @@ internal class ComponentDepot
 
 	private Dictionary<int, HashSet<Type>> entityComponentMap = new Dictionary<int, HashSet<Type>>();
 
+	internal void Register<TComponent>() where TComponent : struct
+	{
+		if (!storages.ContainsKey(typeof(TComponent)))
+		{
+			storages.Add(typeof(TComponent), new ComponentStorage<TComponent>());
+		}
+	}
+
 	private ComponentStorage Lookup(Type type)
 	{
 		return storages[type];
@@ -18,11 +26,7 @@ internal class ComponentDepot
 	private ComponentStorage<TComponent> Lookup<TComponent>() where TComponent : struct
 	{
 		// TODO: is it possible to optimize this?
-		if (!storages.ContainsKey(typeof(TComponent)))
-		{
-			storages.Add(typeof(TComponent), new ComponentStorage<TComponent>());
-		}
-
+		Register<TComponent>();
 		return storages[typeof(TComponent)] as ComponentStorage<TComponent>;
 	}
 
@@ -64,10 +68,7 @@ internal class ComponentDepot
 			{
 				foreach (var filterSignature in filterSignatures)
 				{
-					if (CheckFilter(filterSignature, entityID))
-					{
-						filterSignatureToEntityIDs[filterSignature].Add(entityID);
-					}
+					CheckFilter(filterSignature, entityID);
 				}
 			}
 		}
@@ -96,10 +97,7 @@ internal class ComponentDepot
 			{
 				foreach (var filterSignature in filterSignatures)
 				{
-					if (!CheckFilter(filterSignature, entityID))
-					{
-						filterSignatureToEntityIDs[filterSignature].Remove(entityID);
-					}
+					CheckFilter(filterSignature, entityID);
 				}
 			}
 		}
@@ -175,13 +173,14 @@ internal class ComponentDepot
 		}
 	}
 
-	private bool CheckFilter(FilterSignature filterSignature, int entityID)
+	private void CheckFilter(FilterSignature filterSignature, int entityID)
 	{
 		foreach (var type in filterSignature.Included)
 		{
 			if (!Has(type, entityID))
 			{
-				return false;
+				filterSignatureToEntityIDs[filterSignature].Remove(entityID);
+				return;
 			}
 		}
 
@@ -189,10 +188,11 @@ internal class ComponentDepot
 		{
 			if (Has(type, entityID))
 			{
-				return false;
+				filterSignatureToEntityIDs[filterSignature].Remove(entityID);
+				return;
 			}
 		}
 
-		return true;
+		filterSignatureToEntityIDs[filterSignature].Add(entityID);
 	}
 }
diff --git a/src/FilterBuilder.cs b/src/FilterBuilder.cs
index aec045d..753fb74 100644
--- a/src/FilterBuilder.cs
+++ b/src/FilterBuilder.cs
@@ -22,12 +22,14 @@ public struct FilterBuilder
 
 	public FilterBuilder Include<TComponent>() where TComponent : struct
 	{
+		ComponentDepot.Register<TComponent>();
 		Included.Add(typeof(TComponent));
 		return new FilterBuilder(ComponentDepot, Included, Excluded);
 	}
 
 	public FilterBuilder Exclude<TComponent>() where TComponent : struct
 	{
+		ComponentDepot.Register<TComponent>();
 		Excluded.Add(typeof(TComponent));
 		return new FilterBuilder(ComponentDepot, Included, Excluded);
 	}