Spatial Hash Grid for the Unity Game Engine.

SpatialHashGrid.cs

using System.Collections.Generic;
using Unity.Mathematics;
using UnityEngine;

public class SpatialHashGrid<T> where T : MonoBehaviour
{
    public SpatialHashGrid(int xSize, int ySize, int zSize) => CellSize = new int3(xSize, ySize, zSize);

    public SpatialHashGrid(int3 cellSize) => CellSize = cellSize;

    public int3 CellSize { get; } 

    public Dictionary<T, int3> Units = new();

    public Dictionary<int3, Cell> Cells = new();

    public class Cell
    {
        public Cell(SpatialHashGrid<T> parent, int3 key)
        {
            Parent = parent;
            Key = key;
        }

        public SpatialHashGrid<T> Parent { get; }
        public int3 Key { get; }
        public HashSet<T> Contents { get; } = new();

        public void AddUnit (T unit)
        {
            if(unit == null || Contents.Contains(unit)) return;
            Contents.Add(unit);
        }

        public void RemoveUnit (T unit)
        {
            if(unit == null || !Contents.Contains(unit)) return;
            Contents.Remove(unit);
            if(Contents.Count == 0) Parent.Cells.Remove(Key);
        }
    }

    public int3 GetKey (Vector3 position)
    {
        int x = Mathf.RoundToInt(position.x / (float)CellSize.x) * CellSize.x;
        int y = Mathf.RoundToInt(position.y / (float)CellSize.y) * CellSize.y;
        int z = Mathf.RoundToInt(position.z / (float)CellSize.z) * CellSize.z;

        return new(x, y, z);
    }

    Cell GetCell (int3 key)
    {
        if(!Cells.ContainsKey(key)) Cells.Add(key, new(this, key));
        return Cells[key];
    }

    public void AddUnit (T unit, int3? key = null)
    {
        key ??= GetKey(unit.transform.position);

        Units.Add(unit, key.Value);

        GetCell(key.Value).AddUnit(unit);
    }

    public void RemoveUnit (T unit)
    {
        int3 key = GetKey(unit.transform.position);

        Units.Remove(unit);

        GetCell(key).RemoveUnit(unit);
    }

    public void UpdateUnit (T unit)
    {
        if(unit == null) return;

        int3 key = GetKey(unit.transform.position);

        if(Units.ContainsKey(unit))
        {
            if(key.Equals(Units[unit])) return;

            RemoveUnit(unit);
        }

        AddUnit(unit, key);
    }

    HashSet<T> queryResults = new();
    int3 previousQueryKey;

    public HashSet<T> Query (Vector3 position, int3 bounds)
    {
        int3 offset = GetKey(position);

        if(offset.Equals(previousQueryKey) && queryResults.Count != 0) return queryResults;
        previousQueryKey = offset;

        queryResults.Clear();

        for (int x = -bounds.x; x <= bounds.x; x++)
        {
            for(int y = -bounds.y; y <= bounds.y; y++)
            {
                for(int z = -bounds.z; z <= bounds.z; z++)
                {
                    int3 cellPos = new int3(x, y, z) * CellSize;

                    int3 key = cellPos + offset;

                    if(math.distance(key, offset) > bounds.x * CellSize.x) continue;

                    if(!Cells.ContainsKey(key)) continue;

                    queryResults.UnionWith(Cells[key].Contents);
                }
            }
        }

        return queryResults;
    }
}