j3ml-fork/include/J3ML/Geometry/QuadTree.h

#pragma once

#include <vector>
#include <cstdint>
#include <J3ML/LinearAlgebra/Vector2.h>
#include "AABB2D.h"

namespace Geometry
{
    using LinearAlgebra::Vector2;
    template <typename T>
    class QuadTree
    {
    public:
        struct Node
        {
            Node *parent;
            uint32_t childIndex;
            std::vector<T> objects;

            Vector2 center;
            Vector2 radius;

            bool IsLeaf() const { return childIndex == 0xFFFFFFFF;}

            uint32_t TopLeftChildIndex() const { return childIndex;}
            uint32_t TopRightChildIndex() const { return childIndex+1;}
            uint32_t BottomLeftChildIndex() const { return childIndex+2;}
            uint32_t BottomRightChildIndex() const { return childIndex+3;}

            /// This assumes that the QuadTree contains unique objects and never duplicates
            void Remove(const T &object)
            {
                for (size_t i = 0; i < objects.size(); ++i){
                    if (objects[i] == object) {
                        AssociateQuadTreeNode(object, 0); // Mark in the object that it has been removed from the QuadTree.
                        std::swap(objects[i], objects.back());
                        objects.pop_back();
                        return;
                    }
                }
            }

            AABB2D ComputeAABB()
            {}

            float DistanceSq(const Vector2& point) const
            {
                Vector2 centered = point - center;
                Vector2 closestPoint = centered.Clamp(-radius, radius);
                return closestPoint.DistanceSq(centered);
            }

        };

        // Helper struct used when traversing through the tree
        struct TraversalStackItem
        {
            Node *node;
        };

        QuadTree():
        rootNodeIndex(-1),
        boundingAABB(Vector2(0,0), Vector2(0,0))
        {
            // TODO: currently storing persistent raw pointers to this array outside the array
            nodes.reserve(200000);
        }

        /// Removes all nodes and objects in this tree and reintializes the tree to a single root node.
        void Clear(const Vector2& minXY = Vector2(-1, -1), const Vector2& maxXY = Vector2(1, 1));

        /// Places the given object onto the proper (leaf) node of the tree. After placing, if the leaf split rule is
        /// satisfied, subdivides the leaf node into 4 subquadrants and reassings the objects to the new leaves.
        void Add(const T& object);

        /// Removes the given object from this tree.
        /// To call this function, you must define a function QuadTree<T>::Node *GetQuadTreeNode(const T& object)
        /// which returns the node of this quadtree where the object resides in.
        void Remove(const T& object);

        /// @return The bounding rectangle for the whole tree.
        /// @note This bounding rectangle does not tightly bound the objects themselves, only the root node of the tree.
        AABB2D BoundingAABB() const { return boundingAABB; }

        /// @return The topmost node in the tree.
        Node *Root();
        const Node *Root() const;

        /// Returns the total number of nodes (all nodes, i.e. inner nodes + leaves) in the tree.
        /// Runs in constant time.
        int NumNodes() const;

        /// Returns the total number of leaf nodes in the tree.
        /// @warning Runs in time linear 'O(n)' to the number of nodes in the tree.
        int NumLeaves() const;

        /// Returns the total number of inner nodes in the tree.
        /// @warning Runs in time linear 'O(n)' to the number of nodes in the tree.
        int NumInnerNodes() const;

        /// Returns the total number of objects stored in the tree.
        /// @warning Runs in time linear 'O(n)' to the number of nodes in the tree.
        int NumObjects() const;

        /// Returns the maximum height of the whole tree (the path from the root to the farthest leaf node).
        int TreeHeight() const;

        /// Returns the height of the subtree rooted at 'node'.
        int TreeHeight(const Node *node) const;

        /// Performs an AABB intersection query in this Quadtreee, and calls the given callback function for each non-empty
        /// node of the tree which intersects the given AABB.
        /** @param aabb The axis-aligned bounding box to intersect this QuadTree with.
            @param callback A function or a function object of prototype
                bool callbackFunction(QuadTree<T> &tree, const AABB2D &queryAABB, QuadTree<T>::Node &node);
            If the callback function returns true, the execution of the query is stopped and this function immediately
            returns afterwards. If the callback function returns false, the execution of the query continues. */
        template<typename Func>
        inline void AABBQuery(const AABB2D &aabb, Func &callback);

        /// Finds all object pairs inside the given AABB which have colliding AABBs. For each such pair, calls the
        /// specified callback function.
        template<typename Func>
        inline void CollidingPairsQuery(const AABB2D &aabb, Func &callback);
        /// Performs various consistency checks on the given node. Use only for debugging purposes.
        void DebugSanityCheckNode(Node *n);

    private:
        void Add(const T &object, Node *n);

        /// Allocates a sequential 4-tuple of QuadtreeNodes, contiguous in memory.
        int AllocateNodeGroup(Node *parent);

        void SplitLeaf(Node *leaf);

        std::vector<Node> nodes;
        int rootNodeIndex;
        AABB2D boundingAABB;

        void GrowRootTopLeft();
        void GrowRootTopRight();
        void GrowRootBottomLeft();
        void GrowRootBottomRight();

        void GrowImpl(int quadrantForRoot);
    };

    // NOTE: Keep members defined here. Template-parameterized classes
    // can't be split across header and implementation files
    // but the presence of the implementation file is a requirement


    template <typename T>
    void QuadTree<T>::Clear(const Vector2& minXY, const Vector2& maxXY)
    {
        nodes.clear();

        boundingAABB.minPoint = minXY;
        boundingAABB.maxPoint = maxXY;

        rootNodeIndex = AllocateNodeGroup(0);

        Node *root = Root();
        root->center = (minXY + maxXY) * 0.5f;
        root->radius = maxXY - root->center;
    }
}