MoonTools.ECS/src/Filter.cs

using System;
using System.Collections.Generic;

namespace MoonTools.ECS;

// TODO: do we want to get fancy with queries beyond Include and Exclude?
public class Filter
{
	private Archetype EmptyArchetype;
	private HashSet<TypeId> Included;
	private HashSet<TypeId> Excluded;

	public EntityEnumerator Entities => new EntityEnumerator(this);
	internal ArchetypeEnumerator Archetypes => new ArchetypeEnumerator(this);
	public RandomEntityEnumerator EntitiesInRandomOrder => new RandomEntityEnumerator(this);

	public bool Empty
	{
		get
		{
			var empty = true;

			foreach (var archetype in Archetypes)
			{
				if (archetype.Count > 0)
				{
					return false;
				}
			}

			return empty;
		}
	}

	public int Count
	{
		get
		{
			var count = 0;

			foreach (var archetype in Archetypes)
			{
				count += archetype.Count;
			}

			return count;
		}
	}

	public Entity RandomEntity
	{
		get
		{
			var randomIndex = RandomManager.Next(Count);
			return NthEntity(randomIndex);
		}
	}

	// WARNING: this WILL crash if the index is out of range!
	public Entity NthEntity(int index)
	{
		foreach (var archetype in Archetypes)
		{
			if (index < archetype.Count)
			{
				return archetype.Entities[index];
			}

			index -= archetype.Count;
		}

		throw new InvalidOperationException("Filter index out of range!");
	}

	public void DestroyAllEntities()
	{
		foreach (var archetype in Archetypes)
		{
			archetype.ClearAll();
		}
	}

	internal Filter(Archetype emptyArchetype, HashSet<TypeId> included, HashSet<TypeId> excluded)
	{
		EmptyArchetype = emptyArchetype;
		Included = included;
		Excluded = excluded;
	}

	internal ref struct ArchetypeEnumerator
	{
		private Archetype CurrentArchetype;

		// TODO: pool these
		private Queue<Archetype> ArchetypeQueue = new Queue<Archetype>();
		private Queue<Archetype> ArchetypeSearchQueue = new Queue<Archetype>();
		private HashSet<Archetype> Explored = new HashSet<Archetype>();

		public ArchetypeEnumerator GetEnumerator() => this;

		public ArchetypeEnumerator(Filter filter)
		{
			var empty = filter.EmptyArchetype;
			ArchetypeSearchQueue.Enqueue(empty);

			// TODO: can we cache this search effectively?
			while (ArchetypeSearchQueue.TryDequeue(out var current))
			{
				// exclude the empty archetype
				var satisfiesFilter = filter.Included.Count != 0;

				foreach (var componentId in filter.Included)
				{
					if (!current.Signature.Contains(componentId))
					{
						satisfiesFilter = false;
					}
				}

				foreach (var componentId in filter.Excluded)
				{
					if (current.Signature.Contains(componentId))
					{
						satisfiesFilter = false;
					}
				}

				if (satisfiesFilter)
				{
					ArchetypeQueue.Enqueue(current);
				}

				// breadth-first search
				// ignore excluded component edges
				foreach (var (componentId, archetype) in current.AddEdges)
				{
					if (!Explored.Contains(archetype) && !filter.Excluded.Contains(componentId))
					{
						Explored.Add(archetype);
						ArchetypeSearchQueue.Enqueue(archetype);
					}
				}

				foreach (var (componentId, archetype) in current.RemoveEdges)
				{
					if (!Explored.Contains(archetype))
					{
						Explored.Add(archetype);
						ArchetypeSearchQueue.Enqueue(archetype);
					}
				}
			}
		}

		public bool MoveNext()
		{
			return ArchetypeQueue.TryDequeue(out CurrentArchetype!);
		}

		public Archetype Current => CurrentArchetype;
	}

	public ref struct EntityEnumerator
	{
		private Entity CurrentEntity;

		public EntityEnumerator GetEnumerator() => this;

		// TODO: pool this
		Queue<Entity> EntityQueue = new Queue<Entity>();

		internal EntityEnumerator(Filter filter)
		{
			var archetypeEnumerator = new ArchetypeEnumerator(filter);

			foreach (var archetype in archetypeEnumerator)
			{
				foreach (var entity in archetype.Entities)
				{
					EntityQueue.Enqueue(entity);
				}
			}
		}

		public bool MoveNext()
		{
			return EntityQueue.TryDequeue(out CurrentEntity);
		}

		public Entity Current => CurrentEntity;
	}

	public ref struct RandomEntityEnumerator
	{
		private Filter Filter;
		private LinearCongruentialEnumerator LinearCongruentialEnumerator;

		public RandomEntityEnumerator GetEnumerator() => this;

		internal RandomEntityEnumerator(Filter filter)
		{
			Filter = filter;
			LinearCongruentialEnumerator =
				RandomManager.LinearCongruentialSequence(filter.Count);
		}

		public bool MoveNext() => LinearCongruentialEnumerator.MoveNext();
		public Entity Current => Filter.NthEntity(LinearCongruentialEnumerator.Current);
	}
}