Quadtree (wip)

quadtree.hpp

//  Copyright (c) 2014 Seiji Emery. All rights reserved.

#pragma once

#include "aabb.hpp"
#include "vec2.hpp"

// Quadtree impl.
// Note: This implementation relies on bulk insertion (list of objects is known ahead of time);
// incremental insertion is not yet supported, and tree rebalancing w/ changing object bounds is not,
// and will never be supported (just rebuild the tree instead).
// Also, T is expected to be a non-owning ptr type (T*, shared_ptr, or weak_ptr (recommended); NOT unique_ptr),
// and MUST have a bounds() method that returns an AABB with x1, x2, y1, y2 member values.
//
// Performance wise, full rebuilds should be fast since state gets reused (node / bucket memory doesn't
// get reallocated), and insertions are processed in bulk at each tree level in a cache friendly manner
// (as opposed to deep recursion).
//
// storeUnique: if enabled, each element is only stored once*in the tree, and iteration (via the forEach
// fcns) will never alias. This can cause performance problems, though, as any object overlapping the split
// point gets stored in that node (instead of a leaf), which means that you can have hundreds of objects in
// the root node that will never get ignored (could really slow down collision, etc). 
// If it's turned off, an alternate impl is used, where objects are only stored in leaf nodes (no nodes
// get stored in the root). Since objects can span multiple nodes however, multiple leaves can get references
// to the same object. This is allowed (since the tree doesn't own the objects anyways), and should speed
// up collision detection, *but* this can mess up iteration (unless planned for), so use at your own risk.
template <typename T, unsigned threshold = 16, bool storeUnique = true>
class Quadtree {
protected:
	struct Node {
		std::array<Node*, 4> children;
		std::vector<T> objects;
		AABB bounds, innerBounds;
		Vec2 center;
		bool isLeaf = true;

		// Note: when isLeaf is set, children contains garbage values (not necessarily nullptr),
		// and objects is just a bucket of objects. When it is not set (non-leaf), objects contains
		// only those objects that do not fit within a child node (ie. bounds overlaps split point).

		Node (const AABB & bounds)
			: bounds(bounds), center(bounds.center()) {}
		Node (Node && other)
			: children(std::move(other.children)), objects(std::move(other.objects)),
			  bounds(other.bounds), center(other.center), isLeaf(other.isLeaf) {}

		Node (const Node & other) = delete;				// disable copy constructor and assignment
		Node & operator= (const Node & other) = delete;

		// Resets the node state (should have same effect as calling the constructor again)
		// Other members can have garbage values atm and are expected to be set by Quadtree::rebuild()
		Node & rebuild (const AABB & bounds_) {
			bounds = bounds_;
			center = bounds_.center();
			objects.clear();
			return *this;
		}
	};

	// All nodes in the quadtree are stored in a vector, and are reused when the tree is
	// rebuilt (ie. none of the nodes get deallocated until the tree destructor is called)
	std::vector<Node> nodes; 

	// temp state used by rebuild
	std::vector<T> tmp;
	unsigned nextNode = 1;

	Node * createNode (const AABB & bounds) {
		if (nextNode < nodes.size()) {
			return &(nodes[nextNode++].rebuild(bounds));
		} else {
			nodes.emplace_back(bounds);
			++nextNode;
			return &(nodes.back());
		}
	}

	void rebuild (Node & node) {
		if (node.objects.size() < threshold) {
			node.isLeaf = true;
		} else {
			node.isLeaf = false;

			node.children[0] = createNode({ bounds.x1, center.x, bounds.y1, center.y });
			node.children[1] = createNode({ center.x, bounds.x2, bounds.y1, center.y });
			node.children[2] = createNode({ bounds.x1, center.x, center.y, bounds.y2 });
			node.children[3] = createNode({ bounds.x, bounds.x2, center.y, bounds.y2 });

			if (storeUnique) {
				innerBounds.x1 = innerBounds.x2 = center.x;
				innerBounds.y1 = innerBounds.y2 = center.y;

				tmp.clear();

				for (auto obj : objects) {
					if (obj->bounds().x2 < node.center.x) {
						if (obj->bounds().y2 < node.center.y) {
							node.children[0]->objects.push_back(std::move(obj));
						} else if (obj->bounds().y1 > node.center.y) {
							node.children[2]->objects.push_back(std::move(obj));
						} else {
							innerBounds.grow(obj->bounds());
							tmp.push_back(std::move(obj));
						}
					} else if (obj->bounds().x1 > node.center.y) {
						if (obj->bounds().y2 < node.center.y) {
							node.children[1]->objects.push_back(std::move(obj));
						} else if (obj->bounds().y1 > node.center.y) {
							node.children[3]->objects.push_back(std::move(obj));
						} else {
							innerBounds.grow(obj->bounds());
							tmp.push_back(std::move(obj));
						}
					} else {
						innerBounds.grow(obj->bounds());
						tmp.push_back(std::move(obj));
					}
				}
				std::swap(tmp, node.objects);
			} else {
				// Alternate impl: objs are stored *only* in leaves.
				// The tree may contain duplicate ptrs to each obj, but that's fine since the tree
				// never has ownership over the objects it contains anyways.
				for (auto obj : objects) {
					if (obj->bounds().x1 < node.center.x) {
						if (obj->bounds().y1 < node.center.y)
							node.children[0]->objects.push_back(obj);
						if (obj->bounds().y2 > node.center.y)
							node.children[2]->objects.push_back(obj);
					}
					if (obj->bounds().x2 > node.center.x) {
						if (obj->bounds().y1 < node.center.y)
							node.children[1]->objects.push_back(obj);
						if (obj->bounds().y2 > node.center.y)
							node.children[3]->objects.push_back(obj);
					}
				}
				objects.clear();
			}

			for (auto i = 0; i < 4; ++i) {
				assert(node.children[i] != nullptr);
				rebuild(*(node.children[i]));
			}
		}

		void iterElementsOverlappingPoint (Node & node, Vec2 point, std::function<void(const T &)> callback) {
			// if (node.isLeaf || node.innerBounds.contains(point)) {
			if (node.isLeaf || (node.innerBounds.contains(point) && storeUnique)) {
				for (auto obj : node.objects) {
					if (obj->bounds().contains(point))
						callback(obj);
				}
			}

			if (!node.isLeaf) {
				if (point.x < center.x)
					if (point.y < center.y)
						iterElementsOverlappingPoint(node.children[0], point, callback);
					else
						iterElementsOverlappingPoint(node.children[2], point, callback);
				else
					if (point.y < center.y)
						iterElementsOverlappingPoint(node.children[1], point, callback);
					else
						iterElementsOverlappingPoint(node.children[3], point, callback);
			}
		}

		void iterElementsOverlappingBounds (Node & node, Vec2 point, std::function<void(const T &)> callback) {
			// if (node.isLeaf || node.innerBounds.contains(bounds)) {
			if (node.isLeaf || (node.innerBounds.contains(point) && storeUnique)) {
				for (auto obj : node.objects) {
					if (obj->bounds().contains(bounds))
						callback(obj);
				}
			}

			// if (!node.isLeaf) {
			else {
				for (auto i = 0; i < 4; ++i) {
					assert(node.children[i] != nullptr)
					if (bounds.contains(node.children[i]->bounds)) {
						iterElementsOverlappingBounds(*node.children[i], point, callback);
					}
				}
			}
		}
	}

	void iterPotentialCollisions (Node & node, const std::function<void(const T &, const T &)> & callback) {
		if (!storeUnique) {
			// Note how this MASSIVELY simplifies collision detection xD
			if (!node.isLeaf) {
				for (auto i = 0; i < 4; ++i) {
					assert(node.children[i] != nullptr);
					iterPotentialCollisions(*node.children[i], callback);
				}
			} else {
				for (auto i = 0; i < node.objects.size(); ++i) {
					for (auto j = i + 1; j < node.objects.size(); ++j) {
						if (node.objects[i].bounds().contains(node.objects[j].bounds())) {
							callback(node.objects[i], node.objects[j]);
						}
					}
				}
			}
		} else {
			// more complicated...
		}
	}

public:
	void rebuild (const std::vector<T> & objects) {
		// Determine bounds
		AABB bounds;
		for (auto obj : objects)
			bounds.grow(obj->bounds());

		// Allocate root node
		nextNode = 0;
		createNode(bounds);

		assert(nodes.size() > 0);
		std::copy(objects.begin(), objects.end(), nodes[0].objects.begin());
		rebuild(nodes[0], bounds);
	}

	void forEachElementInBounds (const AABB & bounds, const std::function<void (const T &)> & callback) {
		assert(nodes.size() > 0);
		iterElementsOverlappingBounds(nodes[0], bounds, callback);
		// Note: this now aliases. Fuck.

		// Possible workaround: add a visited flag to T that alternates w/ each pass. When visited, flip the flag,
		// and only execute the visitor logic when the flag is set to an expected value.
		// This is pretty intrusive though... maybe we should add a template switch to select between two impls
		// (either each obj only stored *once* (unique, but can have tons of objs stored in the root), or aliased
		// (iteration becomes problematic, but performance should be better if there's lots of objs that span any 
		// of the split points))
		// edit: sort of fixed (see storeUnique)
	}

	void forEachElementContainingPoint (const Vec2 & bounds, const std::function<void (const T &)> & callback) {
		assert(nodes.size() > 0);
		iterElementsOverlappingPoint(nodes[0], bounds, callback);
	}
};