From 16c8dd19980503f5f12ec46bbdde56afb76908fd Mon Sep 17 00:00:00 2001
From: josh <josh@stepbro.wtf>
Date: Fri, 26 Jan 2024 15:55:51 -0500
Subject: [PATCH] Implement QuadTree 2D

---
 include/J3ML/Geometry/QuadTree.h | 120 +++++++++++++++++++++++++++++++
 1 file changed, 120 insertions(+)

diff --git a/include/J3ML/Geometry/QuadTree.h b/include/J3ML/Geometry/QuadTree.h
index 4414793..66de42a 100644
--- a/include/J3ML/Geometry/QuadTree.h
+++ b/include/J3ML/Geometry/QuadTree.h
@@ -3,9 +3,11 @@
 #include <vector>
 #include <cstdint>
 #include <J3ML/LinearAlgebra/Vector2.h>
+#include "AABB2D.h"
 
 namespace Geometry
 {
+    using LinearAlgebra::Vector2;
     template <typename T>
     class QuadTree
     {
@@ -42,7 +44,125 @@ namespace Geometry
             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;
+    }
 }
\ No newline at end of file