Net.Like.Xue.Tokyo/Assets/BITKit/Core/Quadtree/Quadtree.cs

using System;
using System.Buffers;
using System.Collections;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Net.Http.Headers;
using System.Threading;
using System.Threading.Tasks;
using Unity.Mathematics;

namespace Net.BITKit.Quadtree
{
    public class Quadtree
    {
        public QuadtreeNode Root { get; private set; }

        public IDictionary<int, float2> Positions
        {
            get
            {
                Expansion();
                return _positions;
            }
        }
        private readonly Dictionary<int, float2> _positions;
        private Memory<int> _memory;
        private readonly ConcurrentQueue<(int, float2)> _queue;
        public Quadtree(float2 center, float2 size)
        {
            _memory = new int[(int)math.max(size.x, size.y)];
            _positions = new Dictionary<int, float2>();
            _queue = new ConcurrentQueue<(int, float2)>();
            Root = new QuadtreeNode(this,center, size);
        }
        public void  Insert(in int objectId,in float2 position)
        {
            _queue.Enqueue((objectId, position));

            var root = Root;
        
            InsertRecursive(ref root,in objectId,in position);
        }

        private static void InsertRecursive(ref QuadtreeNode node,in int objectId,in float2 position)
        {
            if (!node.Contains(position))
                return;
            if (node.Objects.Count < 4 || node.Size.x <= 1f) // 假设最小节点大小为1
            {
                node.Objects.Add(objectId);
            }
            else
            {
                if (node.Children[0].Size.x == 0) // 如果子节点未初始化
                {
                    node.Split();
                }
                for (var i = 0; i < 4; i++)
                {
                    InsertRecursive(ref node.Children[i], objectId, position);
                }
            }
        }
        public Memory<int> Query(float2 position, float radius)
        {
            Expansion();
            
            var index = 0;

            var root = Root;
            
            
/*
            var array = MemoryPool<int>.Shared.Rent(_positions.Count);

            try
            {
                QueryRecursive(in root, in position, in radius, array.Memory.Span, ref index);

                return array.Memory;
            }
            finally
            {
                array.Dispose();
            }
            */
            
            QueryRecursive(in root,in position,in radius, _memory.Span, ref index);

            return _memory[..index];
        }

        private void Expansion()
        {
            while (_queue.TryDequeue(out var item))
            {
                _positions.TryAdd(item.Item1, item.Item2);
            }

            if (_positions.Count > _memory.Length)
            {
                _memory = new int[_positions.Count*2];
            }
        }

        private void QueryRecursive(in QuadtreeNode node,in float2 position,in float radius,Span<int> result,ref int index)
        {
            if (!Intersects(node.Center, node.Size, position, radius))
                return;
            
   
            
            foreach (var obj in node.Objects)
            {
                // 直接 TryGetValue，避免 `Positions[obj]` 可能的重复哈希查找
                if (!_positions.TryGetValue(obj, out var objPos)) continue;

                // 计算平方距离，避免额外变量
                if (math.dot(objPos - position, objPos - position) > radius * radius) continue;

                // 直接写入 result
                result[index] = obj;
                index++; 
            }
            

            for (var i = 0; i < 4; i++)
            {
                if (node.Children[i].Size.x > 0)
                {
                    QueryRecursive(in node.Children[i],in position,in radius,result,ref index);
                }
            }
        }

        private static bool Intersects(in float2 center,in float2 size,in float2 position,in float radius)
        {
            // 计算 AABB 的最小/最大点
            var min = center - size * 0.5f;
            var max = center + size * 0.5f;

            // 找到圆心到 AABB 的最近点
            var closest = math.clamp(position, min, max);

            // 计算圆心到最近点的距离
            var delta = position - closest;
            var distanceSq = math.dot(delta, delta);

            // 相交条件：如果距离平方 <= 半径平方
            return distanceSq <= (radius * radius);
        }
        
        
        // 删除对象
        public bool Remove(int objectId)
        {
            Expansion();
            
            if (_positions.TryGetValue(objectId, out var position) is false) return false;

            var root = Root;
            
            if (RemoveRecursive(ref root, objectId, position) is false) return false;
            _positions.Remove(objectId);
            return true;
        }

        private static bool RemoveRecursive(ref QuadtreeNode node, int objectId, float2 position)
        {
            if (!node.Contains(position))
                return false;

            // 尝试从当前节点删除
            if (node.Objects.Remove(objectId))
            {
                return true;
            }

            // 如果当前节点是叶子节点且未找到对象，返回false
            if (node.IsLeaf())
            {
                return false;
            }

            // 递归从子节点删除
            for (var i = 0; i < 4; i++)
            {
                if (RemoveRecursive(ref node.Children[i], objectId, position))
                {
                    // 检查子节点是否需要合并
                    TryMerge(ref node);
                    return true;
                }
            }

            return false;
        }

        // 尝试合并子节点
        private static void TryMerge(ref QuadtreeNode node)
        {
            if (node.IsLeaf())
                return;

            var totalObjects = node.Objects.Count;
            for (var i = 0; i < 4; i++)
            {
                totalObjects += node.Children[i].Objects.Count;
            }

            // 如果所有子节点的对象数量加上当前节点的对象数量小于等于4，则合并
            if (totalObjects <= 4)
            {
                for (var i = 0; i < 4; i++)
                {
                    node.Objects.UnionWith(node.Children[i].Objects);
                    node.Children[i] = new QuadtreeNode(); // 清空子节点
                }
            }
        }

    }
}