Quadtree/Quadtree.cs

using BoundingBox;
using System.Collections.Generic;
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 overlap.</para>
    /// </summary>
    /// <typeparam name="T">Type of the reference to a world object.</typeparam>
    // 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
    {
        /// <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 m_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 m_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> m_vertices;

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

        private BoundingBox2 m_rootBoundingBox;

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

        private Stack<QuadtreeQuery> _queryStack;

        public Quadtree(int maxLeafSize, int maxTreeDepth)
        {
            m_maxLeafSize = maxLeafSize;
            m_maxTreeDepth = maxTreeDepth;
            m_vertices = new List<QuadtreeVertex>() { new QuadtreeVertex(-1, 0) };
            m_items = new List<QuadtreeItem<IWorldObject>>();
            m_rootBoundingBox = new BoundingBox2(Vector2.Zero, Vector2.Zero);
            m_maxObjectRadius = 0f;
            _queryStack = new Stack<QuadtreeQuery>();
        }

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

            if (m_vertices.Count == 1 && m_vertices[0].ChildCount < m_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 (!m_rootBoundingBox.Contains(obj.Position))
                {
                    ExpandRoot(obj.Position);
                }

                (int leafIndex, int depth, BoundingBox2 bounds) = FindLeaf(0, m_rootBoundingBox, 0, obj.Position);
                while (m_vertices[leafIndex].ChildCount >= m_maxLeafSize && depth < m_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, m_rootBoundingBox, 0, obj.Position);

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

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

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

                int childIndex = m_vertices[query.VertexIndex].FirstChildIndex;
                BoundingBox2 halfBounds = new BoundingBox2(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()
        {
            m_vertices.Clear();
            m_vertices.Add(new QuadtreeVertex(-1, 0));
            m_items.Clear();
            m_rootBoundingBox = new BoundingBox2(Vector2.Zero, Vector2.Zero);
            m_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 QuadtreeVertex(m_vertices.Count, -1);
            int oldRootIndex = m_vertices.Count;
            if (position.X > m_rootBoundingBox.Max.X)
            {
                oldRootIndex++;
                m_rootBoundingBox.Max.X += m_rootBoundingBox.Max.X - m_rootBoundingBox.Min.X;
            }
            else
            {
                m_rootBoundingBox.Min.X += m_rootBoundingBox.Min.X - m_rootBoundingBox.Max.X;
            }
            if (position.Y > m_rootBoundingBox.Max.Y)
            {
                oldRootIndex += 2;
                m_rootBoundingBox.Max.Y += m_rootBoundingBox.Max.Y - m_rootBoundingBox.Min.Y;
            }
            else
            {
                m_rootBoundingBox.Min.Y += m_rootBoundingBox.Min.Y - m_rootBoundingBox.Max.Y;
            }

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

        private void AddToRootLeaf(int itemIndex)
        {
            Vector2 position = m_items[itemIndex].WorldObject.Position;
            if (m_vertices[0].ChildCount == 0)
            {
                m_vertices[0] = new QuadtreeVertex(itemIndex, 1);
                m_rootBoundingBox.Min = (m_rootBoundingBox.Max = position);
            }
            else
            {
                var item = m_items[itemIndex];
                item.Next = m_vertices[0].FirstChildIndex;
                m_items[itemIndex] = item;
                m_vertices[0] = new QuadtreeVertex(itemIndex, m_vertices[0].ChildCount + 1);
                IncludeInBoundingBox(ref m_rootBoundingBox, position);
            }
        }

        private void AddToLeaf(int leafIndex, int itemIndex)
        {
            int next;
            if (m_vertices[leafIndex].ChildCount == 0)
            {
                next = -1;
                m_vertices[leafIndex] = new QuadtreeVertex(itemIndex, 1);
            }
            else
            {
                next = m_vertices[leafIndex].FirstChildIndex;
                m_vertices[leafIndex] = new QuadtreeVertex(itemIndex, m_vertices[leafIndex].ChildCount + 1);
            }
            m_items[itemIndex] = new QuadtreeItem<IWorldObject>(m_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 = m_vertices[leafIndex];
            m_vertices[leafIndex] = new QuadtreeVertex(m_vertices.Count, -1);
            m_vertices.Add(new QuadtreeVertex(-1, 0));
            m_vertices.Add(new QuadtreeVertex(-1, 0));
            m_vertices.Add(new QuadtreeVertex(-1, 0));
            m_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 = m_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 = m_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 (m_vertices[vertexIndex].ChildCount == -1)
            {
                depth++;
                var center = bounds.Center;
                vertexIndex = m_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 (m_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 (m_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)
        {
            // TODO: Stack as reused member variable?
            Stack<int> stack = new Stack<int>();
            stack.Push(vertexIndex);

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

        private void QueryAppendContainedLeafItems(int vertexIndex, BoundingBox2 box, List<IWorldObject> resultList)
        {
            int index = m_vertices[vertexIndex].FirstChildIndex;
            while (index != -1)
            {
                IWorldObject obj = m_items[index].WorldObject;
                if (box.Intersects(obj.BoundingBox) != IntersectionType.Disjoint)
                {
                    resultList.Add(obj);
                }
                index = m_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 "m_items.Count - 1" to "index".
        /// </summary>
        /// <param name="index">Index of the item to be removed.</param>
        private void RemoveFromItems(int index)
        {
            m_items.RemoveUnorderedAt(index);
            if (index != m_items.Count)
            {
                (int leafIndex, _, _) = FindLeaf(0, m_rootBoundingBox, 0, m_items[index].WorldObject.Position);
                int current = m_vertices[leafIndex].FirstChildIndex;
                if (current == m_items.Count)
                {
                    m_vertices[leafIndex] = new QuadtreeVertex(index, m_vertices[leafIndex].ChildCount);
                }
                else
                {
                    while (m_items[current].Next != m_items.Count)
                    {
                        current = m_items[current].Next;
                    }
                    m_items[current] = new QuadtreeItem<IWorldObject>(m_items[current].WorldObject, index);
                }
            }
        }

        /// <summary>
        /// Slightly more efficient implementation of BoundingBox.Include(), assuming that Min.X <= Max.X,
        /// Min.Z <= Max.Z, and all Y = 0.
        /// </summary>
        private void IncludeInBoundingBox(ref BoundingBox2 boundingBox, Vector2 position)
        {
            if (boundingBox.Min.X > position.X)
            {
                boundingBox.Min.X = position.X;
            }
            else if (boundingBox.Max.X < position.X)
            {
                boundingBox.Max.X = position.X;
            }
            if (boundingBox.Min.Y > position.Y)
            {
                boundingBox.Min.Y = position.Y;
            }
            else if (boundingBox.Max.Y < position.Y)
            {
                boundingBox.Max.Y = position.Y;
            }
        }
    }
}