MoonWorks/src/Collision/Fixed/SpatialHash2D.cs

using System.Collections.Generic;
using MoonWorks.Math.Fixed;

namespace MoonWorks.Collision.Fixed
{
	/// <summary>
	/// Used to quickly check if two shapes are potentially overlapping.
	/// </summary>
	/// <typeparam name="T">The type that will be used to uniquely identify shape-transform pairs.</typeparam>
	public class SpatialHash2D<T, U> where T : struct, System.IEquatable<T>
	{
		private readonly Fix64 cellSize;

		private readonly Dictionary<long, HashSet<T>> hashDictionary = new Dictionary<long, HashSet<T>>();
		private readonly Dictionary<T, AABB2D> IDBoxLookup = new Dictionary<T, AABB2D>();
		private readonly Dictionary<T, U> IDDataLookup = new Dictionary<T, U>();

		private readonly HashSet<T> DynamicIDs = new HashSet<T>();

		private int MinX;
		private int MaxX;
		private int MinY;
		private int MaxY;

		private Queue<HashSet<T>> hashSetPool = new Queue<HashSet<T>>();

		public SpatialHash2D(int cellSize)
		{
			this.cellSize = new Fix64(cellSize);
		}

		private (int, int) Hash(Vector2 position)
		{
			return ((int) (position.X / cellSize), (int) (position.Y / cellSize));
		}

		/// <summary>
		/// Inserts an element into the SpatialHash.
		/// </summary>
		/// <param name="id">A unique ID for the shape-transform pair.</param>
		public void Insert(T id, AABB2D aabb, Transform2D transform2D, U data, bool dynamic = true)
		{
			Remove(id);

			var box = AABB2D.Transformed(aabb, transform2D);
			var minHash = Hash(box.Min);
			var maxHash = Hash(box.Max);

			foreach (var key in Keys(minHash.Item1, minHash.Item2, maxHash.Item1, maxHash.Item2))
			{
				if (!hashDictionary.ContainsKey(key))
				{
					hashDictionary.Add(key, new HashSet<T>());
				}

				hashDictionary[key].Add(id);
				IDDataLookup[id] = data;
			}

			MinX = System.Math.Min(MinX, minHash.Item1);
			MinY = System.Math.Min(MinY, minHash.Item2);
			MaxX = System.Math.Max(MaxX, maxHash.Item1);
			MaxY = System.Math.Max(MaxY, maxHash.Item2);

			if (dynamic)
			{
				DynamicIDs.Add(id);
			}

			IDBoxLookup[id] = box;
		}

		/// <summary>
		/// Retrieves all the potential collisions of a shape-transform pair. Excludes any shape-transforms with the given ID.
		/// </summary>
		public RetrieveEnumerator Retrieve<V>(T id, V hasAABB, Transform2D transform2D) where V : IHasAABB2D
		{
			var box = AABB2D.Transformed(hasAABB.AABB, transform2D);
			var (minX, minY) = Hash(box.Min);
			var (maxX, maxY) = Hash(box.Max);

			if (minX < MinX) { minX = MinX; }
			if (maxX > MaxX) { maxX = MaxX; }
			if (minY < MinY) { minY = MinY; }
			if (maxY > MaxY) { maxY = MaxY; }

			return new RetrieveEnumerator(
				this,
				Keys(minX, minY, maxX, maxY),
				id
			);
		}

		/// <summary>
		/// Retrieves all the potential collisions of a shape-transform pair.
		/// </summary>
		public RetrieveEnumerator Retrieve<V>(V hasAABB, Transform2D transform2D) where V : IHasAABB2D
		{
			var box = AABB2D.Transformed(hasAABB.AABB, transform2D);
			return Retrieve(box);
		}

		/// <summary>
		/// Retrieves objects based on a pre-transformed AABB.
		/// </summary>
		/// <param name="aabb">A transformed AABB.</param>
		/// <returns></returns>
		public RetrieveEnumerator Retrieve(T id, AABB2D aabb)
		{
			var (minX, minY) = Hash(aabb.Min);
			var (maxX, maxY) = Hash(aabb.Max);

			if (minX < MinX) { minX = MinX; }
			if (maxX > MaxX) { maxX = MaxX; }
			if (minY < MinY) { minY = MinY; }
			if (maxY > MaxY) { maxY = MaxY; }

			return new RetrieveEnumerator(
				this,
				Keys(minX, minY, maxX, maxY),
				id
			);
		}

		/// <summary>
		/// Retrieves objects based on a pre-transformed AABB.
		/// </summary>
		/// <param name="aabb">A transformed AABB.</param>
		/// <returns></returns>
		public RetrieveEnumerator Retrieve(AABB2D aabb)
		{
			var (minX, minY) = Hash(aabb.Min);
			var (maxX, maxY) = Hash(aabb.Max);

			if (minX < MinX) { minX = MinX; }
			if (maxX > MaxX) { maxX = MaxX; }
			if (minY < MinY) { minY = MinY; }
			if (maxY > MaxY) { maxY = MaxY; }

			return new RetrieveEnumerator(
				this,
				Keys(minX, minY, maxX, maxY)
			);
		}

		/// <summary>
		/// Removes a specific ID from the SpatialHash.
		/// </summary>
		public void Remove(T id)
		{
			if (IDBoxLookup.TryGetValue(id, out var aabb))
			{
				var minHash = Hash(aabb.Min);
				var maxHash = Hash(aabb.Max);

				foreach (var key in Keys(minHash.Item1, minHash.Item2, maxHash.Item1, maxHash.Item2))
				{
					if (hashDictionary.TryGetValue(key, out HashSet<T> value))
					{
						value.Remove(id);
					}
				}

				IDDataLookup.Remove(id);
				IDBoxLookup.Remove(id);
			}

			DynamicIDs.Remove(id);
		}

		/// <summary>
		/// Removes everything that has been inserted into the SpatialHash.
		/// </summary>
		public void Clear()
		{
			foreach (var hash in hashDictionary.Values)
			{
				hash.Clear();
			}

			IDDataLookup.Clear();
			IDBoxLookup.Clear();
		}

		/// <summary>
		/// Removes
		/// </summary>
		public void ClearDynamic()
		{
			foreach (var id in DynamicIDs)
			{
				Remove(id);
			}
		}

		private static long MakeLong(int left, int right)
		{
			return ((long) left << 32) | ((uint) right);
		}

		internal static KeysEnumerator Keys(int minX, int minY, int maxX, int maxY)
		{
			return new KeysEnumerator(minX, minY, maxX, maxY);
		}

		internal HashSet<T> AcquireHashSet()
		{
			if (hashSetPool.Count == 0)
			{
				hashSetPool.Enqueue(new HashSet<T>());
			}

			var hashSet = hashSetPool.Dequeue();
			hashSet.Clear();
			return hashSet;
		}

		internal void FreeHashSet(HashSet<T> hashSet)
		{
			hashSetPool.Enqueue(hashSet);
		}

		internal ref struct KeysEnumerator
		{
			private int MinX;
			private int MinY;
			private int MaxX;
			private int MaxY;
			private int i, j;

			public KeysEnumerator GetEnumerator() => this;

			public KeysEnumerator(int minX, int minY, int maxX, int maxY)
			{
				MinX = minX;
				MinY = minY;
				MaxX = maxX;
				MaxY = maxY;
				i = minX;
				j = minY - 1;
			}

			public bool MoveNext()
			{
				if (j < MaxY)
				{
					j += 1;
					return true;
				}
				else if (i < MaxX)
				{
					i += 1;
					j = MinY;
					return true;
				}

				return false;
			}

			public long Current
			{
				get
				{
					return MakeLong(i, j);
				}
			}
		}

		public ref struct RetrieveEnumerator
		{
			public SpatialHash2D<T, U> SpatialHash;
			private KeysEnumerator KeysEnumerator;
			private HashSet<T>.Enumerator HashSetEnumerator;
			private bool HashSetEnumeratorActive;
			private HashSet<T> Duplicates;
			private T? ID;

			public RetrieveEnumerator GetEnumerator() => this;

			internal RetrieveEnumerator(
				SpatialHash2D<T, U> spatialHash,
				KeysEnumerator keysEnumerator,
				T id
			) {
				SpatialHash = spatialHash;
				KeysEnumerator = keysEnumerator;
				HashSetEnumerator = default;
				HashSetEnumeratorActive = false;
				Duplicates = SpatialHash.AcquireHashSet();
				ID = id;
			}

			internal RetrieveEnumerator(
				SpatialHash2D<T, U> spatialHash,
				KeysEnumerator keysEnumerator
			) {
				SpatialHash = spatialHash;
				KeysEnumerator = keysEnumerator;
				HashSetEnumerator = default;
				HashSetEnumeratorActive = false;
				Duplicates = SpatialHash.AcquireHashSet();
				ID = null;
			}

			public bool MoveNext()
			{
				if (!HashSetEnumeratorActive || !HashSetEnumerator.MoveNext())
				{
					if (!KeysEnumerator.MoveNext())
					{
						SpatialHash.FreeHashSet(Duplicates);
						return false;
					}

					if (SpatialHash.hashDictionary.TryGetValue(KeysEnumerator.Current, out var hashset))
					{
						HashSetEnumerator = hashset.GetEnumerator();
						HashSetEnumeratorActive = true;
					}

					return MoveNext();
				}

				// conditions
				var t = HashSetEnumerator.Current;

				if (Duplicates.Contains(t))
				{
					return MoveNext();
				}

				if (ID.HasValue)
				{
					if (ID.Value.Equals(t))
					{
						return MoveNext();
					}
				}

				Duplicates.Add(t);
				return true;
			}

			public (T, U) Current
			{
				get
				{
					var t = HashSetEnumerator.Current;
					var u = SpatialHash.IDDataLookup[t];
					return (t, u);
				}
			}
		}
	}
}