riversong-code-showcase/Source/Riversong/Game/Collections/SpatialLookup.cs

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

namespace DanieleMarotta.RiversongCodeShowcase
{
    public class SpatialLookup<T>
    {
        private static readonly Predicate<T> NoopFilter = _ => true;

        private int _cellSize;

        private ListMultiDictionary<int2, Item> _lookup = new();

        public SpatialLookup(int cellSize)
        {
            _cellSize = cellSize;
        }

        public void Add(T obj, TileRect rect, int key)
        {
            foreach (var cell in CellRange(rect)) _lookup.Add(cell, new Item(obj, rect, key));
        }

        public void Remove(TileRect rect, int key)
        {
            foreach (var cell in CellRange(rect)) _lookup.Remove(cell, new Item(key));
        }

        private TileRange CellRange(TileRect rect)
        {
            return TileRange.From(rect.Min / _cellSize, rect.Max / _cellSize);
        }

        public void RemoveAll(TileRect rect, Predicate<T> filter = null, List<T> result = null)
        {
            filter ??= NoopFilter;

            using var closedScope = HashSetPool<int>.Get(out var closed);
            using var toRemoveScope = ListPool<(int2, Item)>.Get(out var toRemove);

            var cellMin = rect.Min / _cellSize;
            var cellMax = rect.Max / _cellSize;

            int2 cell;
            for (cell.x = cellMin.x; cell.x <= cellMax.x; cell.x++)
            for (cell.y = cellMin.y; cell.y <= cellMax.y; cell.y++)
            {
                if (!_lookup.TryGetValues(cell, out var list)) continue;

                if (cell.x == cellMin.x || cell.x == cellMax.x || cell.y == cellMin.y || cell.y == cellMax.y)
                {
                    foreach (var item in list)
                    {
                        if (!closed.Add(item.EqualityKey) || !rect.Intersects(item.Rect) || !filter.Invoke(item.Obj)) continue;
                        toRemove.Add((cell, item));
                        result?.Add(item.Obj);
                    }
                }
                else
                {
                    if (result != null)
                        foreach (var item in list)
                            if (closed.Add(item.EqualityKey) || !filter.Invoke(item.Obj))
                                result.Add(item.Obj);
                    _lookup.Clear(cell);
                }
            }

            foreach (var (key, value) in toRemove) _lookup.Remove(key, value);
        }

        public void Find(TileRect rect, List<T> result)
        {
            using var closedScope = HashSetPool<int>.Get(out var closed);

            var cellMin = rect.Min / _cellSize;
            var cellMax = rect.Max / _cellSize;

            int2 cell;
            for (cell.x = cellMin.x; cell.x <= cellMax.x; cell.x++)
            for (cell.y = cellMin.y; cell.y <= cellMax.y; cell.y++)
            {
                if (!_lookup.TryGetValues(cell, out var list)) continue;

                if (cell.x == cellMin.x || cell.x == cellMax.x || cell.y == cellMin.y || cell.y == cellMax.y)
                    foreach (var item in list)
                    {
                        if (!rect.Intersects(item.Rect) || !closed.Add(item.EqualityKey)) continue;
                        result.Add(item.Obj);
                    }
                else
                    foreach (var item in list)
                        if (closed.Add(item.EqualityKey))
                            result.Add(item.Obj);
            }
        }

        public bool FindMax(TileRect rect, Func<T, int> scoreFunc, out T max)
        {
            using var resultScope = ListPool<T>.Get(out var result);

            Find(rect, result);

            max = default;
            var maxScore = int.MinValue;

            foreach (var item in result)
            {
                var score = scoreFunc.Invoke(item);
                if (score > maxScore)
                {
                    maxScore = score;
                    max = item;
                }
            }

            return maxScore > int.MinValue;
        }

        private struct Item : IEquatable<Item>
        {
            public T Obj;

            public TileRect Rect;

            public int EqualityKey;

            public Item(T obj, TileRect rect, int equalityKey)
            {
                Obj = obj;
                Rect = rect;
                EqualityKey = equalityKey;
            }

            public Item(int equalityKey) : this(default, default, equalityKey)
            {
            }

            public bool Equals(Item other)
            {
                return EqualityKey == other.EqualityKey;
            }
        }
    }
}