Quadtree/Quadtree.cs

using System.Numerics;

namespace Quadtree
{
    /// <summary>
    ///   <para>Quadtree for fast add, remove, and query of world objects. The world size does not have to be known in
    ///   advance, instead the boundaries will be determined dynamically as objects are added.</para>
    ///   <para>Objects are expected to have similar size, and to be somewhat small in relation to the distances in
    ///   between them. This allows simple queries. If very large objects would have to be stored, it could become
    ///   necessary to store objects not in a single leaf, but in each leaf they intersect with.</para>
    /// </summary>
    // TODO: Support for moving objects without having to remove and re-add them.
    // TODO: Add method to prune empty leaves. An empty leaf is not removed automatically when its last object is removed.
    // TODO: Add unit tests.
    public class Quadtree
    {
        // TODO: It should be possible to calculate whether average query time will be better with or without split, without resorting to max leaf size and max tree depth.
        /// <summary>
        /// Maximum number of items in a leaf vertex before it is split, unless the depth of the leaf is greater or
        /// equal to the maximum tree depth.
        /// </summary>
        private readonly int _maxLeafSize;

        /// <summary>
        /// Maximum depth of the quadtree, that is the maximum allowed length of the path from the root to a new leaf
        /// when determining whether a leaf vertex may be split.
        /// </summary>
        private readonly int _maxTreeDepth;

        /// <summary>
        /// The list of quadtree vertices. The root vertex is always the first item. The four child vertices of a branch
        /// are always created together and stored contiguously.
        /// </summary>
        private readonly List<QuadtreeVertex> _vertices;

        /// <summary>
        /// The list of all items stored in the quadtree.
        /// </summary>
        private readonly List<QuadtreeItem<IWorldObject>> _items;

        private BoundingBox2 _rootBoundingBox;

        /// <summary>
        /// The maximum bounding radius of all stored objects, used to inflate the bounding box for queries.
        /// </summary>
        private float _maxObjectRadius;

        /// <summary>
        /// Stack of vertex indices with their bounding boxes. It is reused for each query.
        /// </summary>
        private readonly Stack<QuadtreeQuery> _queryStack;

        /// <summary>
        /// Stack of vertex indices. It is reused for each query.
        /// </summary>
        private readonly Stack<int> _vertexStack;

        public Quadtree(int maxLeafSize, int maxTreeDepth)
        {
            _maxLeafSize = maxLeafSize;
            _maxTreeDepth = maxTreeDepth;
            _vertices = new List<QuadtreeVertex>() { new(-1, 0) };
            _items = new List<QuadtreeItem<IWorldObject>>();
            _rootBoundingBox = new BoundingBox2(Vector2.Zero, Vector2.Zero);
            _maxObjectRadius = 0f;
            _queryStack = new Stack<QuadtreeQuery>();
            _vertexStack = new Stack<int>();
        }

        public void Add(IWorldObject obj)
        {
            var itemIndex = _items.Count;
            _items.Add(new QuadtreeItem<IWorldObject>(obj));
            if (_maxObjectRadius < obj.BoundingRadius)
            {
                _maxObjectRadius = obj.BoundingRadius;
            }

            if (_vertices.Count == 1 && _vertices[0].ChildCount < _maxLeafSize)
            {
                // Before doing any splitting, the initial quadtree root vertex is filled up in order to get a
                // reasonable guess for what the world space is. This space can be expanded later if that assumption
                // turns out to be false.
                AddToRootLeaf(itemIndex);
            }
            else
            {
                // Expands the root until the new world object fits.
                while (!_rootBoundingBox.Contains(obj.Position))
                {
                    ExpandRoot(obj.Position);
                }

                (int leafIndex, int depth, BoundingBox2 bounds) = FindLeaf(0, _rootBoundingBox, 0, obj.Position);
                while (_vertices[leafIndex].ChildCount >= _maxLeafSize && depth < _maxTreeDepth)
                {
                    // Splits the vertex and decends into one of the new leaves.
                    SplitLeaf(leafIndex, bounds);
                    (leafIndex, depth, bounds) = FindLeaf(leafIndex, bounds, depth, obj.Position);
                }
                AddToLeaf(leafIndex, itemIndex);
            }
        }

        public bool Remove(IWorldObject obj)
        {
            // Finds the leaf that should contain the world object.
            (int leafIndex, _, _) = FindLeaf(0, _rootBoundingBox, 0, obj.Position);

            // Tries to find the item of the world object.
            int previous = -1;
            int current = _vertices[leafIndex].FirstChildIndex;
            while (current != -1)
            {
                QuadtreeItem<IWorldObject> item = _items[current];
                if (item.WorldObject == obj)
                {
                    // Removes the found item from the leaf vertex.
                    QuadtreeVertex leaf = _vertices[leafIndex];
                    leaf.ChildCount--;
                    if (previous != -1)
                    {
                        _items[previous] = new QuadtreeItem<IWorldObject>(_items[previous].WorldObject, item.Next);
                    }
                    else
                    {
                        leaf.FirstChildIndex = item.Next;
                    }
                    _vertices[leafIndex] = leaf;

                    // Removes the found item from the item list.
                    RemoveFromItems(current);

                    return true;
                }
                previous = current;
                current = item.Next;
            }

            // World object was not found.
            return false;
        }

        public void Query(BoundingBox2 box, List<IWorldObject> resultList)
        {
            BoundingBox2 inflatedBox = box;
            inflatedBox.Expand(_maxObjectRadius);

            QueryProcessChildVertex(inflatedBox, box, 0, _rootBoundingBox, resultList, _queryStack);

            while (_queryStack.Count > 0)
            {
                QuadtreeQuery query = _queryStack.Pop();

                int childIndex = _vertices[query.VertexIndex].FirstChildIndex;
                BoundingBox2 halfBounds = new(query.VertexBounds.Center, query.VertexBounds.Max);
                Vector2 halfSize = halfBounds.Size;
                QueryProcessChildVertex(inflatedBox, box, childIndex++, halfBounds, resultList, _queryStack);

                halfBounds.Translate(new Vector2(-halfSize.X, 0f));
                QueryProcessChildVertex(inflatedBox, box, childIndex++, halfBounds, resultList, _queryStack);

                halfBounds.Translate(new Vector2(halfSize.X, -halfSize.Y));
                QueryProcessChildVertex(inflatedBox, box, childIndex++, halfBounds, resultList, _queryStack);

                halfBounds.Translate(new Vector2(-halfSize.X, 0f));
                QueryProcessChildVertex(inflatedBox, box, childIndex, halfBounds, resultList, _queryStack);
            }
        }

        public void Clear()
        {
            _vertices.Clear();
            _vertices.Add(new QuadtreeVertex(-1, 0));
            _items.Clear();
            _rootBoundingBox = new BoundingBox2(Vector2.Zero, Vector2.Zero);
            _maxObjectRadius = 0f;
        }

        /// <summary>
        /// Expands the root bounding box by turning the root into a child of a new root vertex.
        /// </summary>
        /// <param name="position">Position outside the root bounding box to be included in the quadtree.</param>
        private void ExpandRoot(Vector2 position)
        {
            // Finds the new index for the old root vertex and expands the bounds.
            QuadtreeVertex newRoot = new(_vertices.Count, -1);
            int oldRootIndex = _vertices.Count;
            if (position.X > _rootBoundingBox.Max.X)
            {
                oldRootIndex++;
                _rootBoundingBox.Max.X += _rootBoundingBox.Max.X - _rootBoundingBox.Min.X;
            }
            else
            {
                _rootBoundingBox.Min.X += _rootBoundingBox.Min.X - _rootBoundingBox.Max.X;
            }
            if (position.Y > _rootBoundingBox.Max.Y)
            {
                oldRootIndex += 2;
                _rootBoundingBox.Max.Y += _rootBoundingBox.Max.Y - _rootBoundingBox.Min.Y;
            }
            else
            {
                _rootBoundingBox.Min.Y += _rootBoundingBox.Min.Y - _rootBoundingBox.Max.Y;
            }

            // Adds the new leaves and updates the old and new root.
            for (int i = 0; i < 4; i++)
            {
                if (oldRootIndex == _vertices.Count)
                {
                    _vertices.Add(_vertices[0]);
                }
                else
                {
                    _vertices.Add(new QuadtreeVertex(-1, 0));
                }
            }
            _vertices[0] = newRoot;
        }

        private void AddToRootLeaf(int itemIndex)
        {
            Vector2 position = _items[itemIndex].WorldObject.Position;
            if (_vertices[0].ChildCount == 0)
            {
                _vertices[0] = new QuadtreeVertex(itemIndex, 1);
                _rootBoundingBox.Min = (_rootBoundingBox.Max = position);
            }
            else
            {
                var item = _items[itemIndex];
                item.Next = _vertices[0].FirstChildIndex;
                _items[itemIndex] = item;
                _vertices[0] = new QuadtreeVertex(itemIndex, _vertices[0].ChildCount + 1);
                _rootBoundingBox.Include(position);
            }
        }

        private void AddToLeaf(int leafIndex, int itemIndex)
        {
            int next;
            if (_vertices[leafIndex].ChildCount == 0)
            {
                next = -1;
                _vertices[leafIndex] = new QuadtreeVertex(itemIndex, 1);
            }
            else
            {
                next = _vertices[leafIndex].FirstChildIndex;
                _vertices[leafIndex] = new QuadtreeVertex(itemIndex, _vertices[leafIndex].ChildCount + 1);
            }
            _items[itemIndex] = new QuadtreeItem<IWorldObject>(_items[itemIndex].WorldObject, next);
        }

        private void SplitLeaf(int leafIndex, BoundingBox2 leafBounds)
        {
            // Splits the leaf vertex by turning it into a branch and adding four new leaves.
            var oldLeaf = _vertices[leafIndex];
            _vertices[leafIndex] = new QuadtreeVertex(_vertices.Count, -1);
            _vertices.Add(new QuadtreeVertex(-1, 0));
            _vertices.Add(new QuadtreeVertex(-1, 0));
            _vertices.Add(new QuadtreeVertex(-1, 0));
            _vertices.Add(new QuadtreeVertex(-1, 0));

            // Distributes the existing items over the new leaves.
            var center = leafBounds.Center;
            int next = oldLeaf.FirstChildIndex;
            while (next != -1)
            {
                var item = _items[next];
                int newLeafIndex = FindNextVertex(leafIndex, center, item.WorldObject.Position);
                AddToLeaf(newLeafIndex, next);
                next = item.Next;
            }
        }

        private int FindNextVertex(int vertexIndex, Vector2 boundsCenter, Vector2 position)
        {
            vertexIndex = _vertices[vertexIndex].FirstChildIndex;
            if (position.X < boundsCenter.X)
            {
                vertexIndex++;
            }
            if (position.Y < boundsCenter.Y)
            {
                vertexIndex += 2;
            }
            return vertexIndex;
        }

        private (int leafIndex, int depth, BoundingBox2 bounds) FindLeaf(int vertexIndex, BoundingBox2 bounds,
            int depth, Vector2 position)
        {
            while (_vertices[vertexIndex].ChildCount == -1)
            {
                depth++;
                var center = bounds.Center;
                vertexIndex = _vertices[vertexIndex].FirstChildIndex;
                if (position.X < center.X)
                {
                    vertexIndex++;
                    bounds.Max.X = center.X;
                }
                else
                {
                    bounds.Min.X = center.X;
                }
                if (position.Y < center.Y)
                {
                    vertexIndex += 2;
                    bounds.Max.Y = center.Y;
                }
                else
                {
                    bounds.Min.Y = center.Y;
                }
            }
            return (vertexIndex, depth, bounds);
        }

        private void QueryProcessChildVertex(BoundingBox2 inflatedBox, BoundingBox2 box, int vertexIndex,
            BoundingBox2 vertexBounds, List<IWorldObject> resultList, Stack<QuadtreeQuery> stack)
        {
            switch (inflatedBox.Intersects(vertexBounds))
            {
                case IntersectionType.Contains:
                    if (_vertices[vertexIndex].ChildCount == -1)
                    {
                        // Found contained vertex, appends items of all child vertices.
                        QueryAppendContainedItems(vertexIndex, box, resultList);
                    }
                    else
                    {
                        // Found leaf, appends items.
                        QueryAppendContainedLeafItems(vertexIndex, box, resultList);
                    }
                    break;
                case IntersectionType.Intersects:
                    if (_vertices[vertexIndex].ChildCount == -1)
                    {
                        // Branches the query.
                        stack.Push(new QuadtreeQuery(vertexIndex, vertexBounds));
                    }
                    else
                    {
                        // Found leaf, appends items.
                        QueryAppendContainedLeafItems(vertexIndex, box, resultList);
                    }
                    break;
            }
        }

        private void QueryAppendContainedItems(int vertexIndex, BoundingBox2 box, List<IWorldObject> resultList)
        {
            _vertexStack.Push(vertexIndex);

            while (_vertexStack.Count > 0)
            {
                int current = _vertexStack.Pop();
                if (_vertices[current].ChildCount == -1)
                {
                    // Branches into all vertices until leaves are found.
                    int childIndex = _vertices[current].FirstChildIndex;
                    _vertexStack.Push(childIndex++);
                    _vertexStack.Push(childIndex++);
                    _vertexStack.Push(childIndex++);
                    _vertexStack.Push(childIndex);
                }
                else
                {
                    // Found leaf, appends items.
                    QueryAppendContainedLeafItems(current, box, resultList);
                }
            }
        }

        private void QueryAppendContainedLeafItems(int vertexIndex, BoundingBox2 box, List<IWorldObject> resultList)
        {
            int index = _vertices[vertexIndex].FirstChildIndex;
            while (index != -1)
            {
                IWorldObject obj = _items[index].WorldObject;
                if (box.Intersects(obj.BoundingBox) != IntersectionType.Disjoint)
                {
                    resultList.Add(obj);
                }
                index = _items[index].Next;
            }
        }

        /// <summary>
        /// Removes the found item from the item list. Since the list is being reordered for fast removal, the reference
        /// to the relocated item has to be updated because its list index changes from "_items.Count - 1" to "index".
        /// </summary>
        /// <param name="index">Index of the item to be removed.</param>
        private void RemoveFromItems(int index)
        {
            _items.RemoveUnorderedAt(index);
            if (index != _items.Count)
            {
                (int leafIndex, _, _) = FindLeaf(0, _rootBoundingBox, 0, _items[index].WorldObject.Position);
                int current = _vertices[leafIndex].FirstChildIndex;
                if (current == _items.Count)
                {
                    _vertices[leafIndex] = new QuadtreeVertex(index, _vertices[leafIndex].ChildCount);
                }
                else
                {
                    while (_items[current].Next != _items.Count)
                    {
                        current = _items[current].Next;
                    }
                    _items[current] = new QuadtreeItem<IWorldObject>(_items[current].WorldObject, index);
                }
            }
        }
    }
}