libraries/path-bool/src/util/quad_tree.rs

use crate::aabb::Aabb;
use std::collections::HashSet;

pub struct QuadTree<T> {
	bounding_box: Aabb,
	depth: usize,
	inner_node_capacity: usize,
	subtrees: Option<Box<[QuadTree<T>; 4]>>,
	pairs: Vec<(Aabb, T)>,
}

impl<T: Clone> QuadTree<T> {
	pub fn new(bounding_box: Aabb, depth: usize, inner_node_capacity: usize) -> Self {
		QuadTree {
			bounding_box,
			depth,
			inner_node_capacity,
			subtrees: None,
			pairs: Vec::new(),
		}
	}

	pub fn insert(&mut self, bounding_box: Aabb, value: T) -> bool {
		if !crate::aabb::bounding_boxes_overlap(&bounding_box, &self.bounding_box) {
			return false;
		}

		if self.depth > 0 && self.pairs.len() >= self.inner_node_capacity {
			self.ensure_subtrees();
			for tree in self.subtrees.as_mut().unwrap().iter_mut() {
				tree.insert(bounding_box, value.clone());
			}
		} else {
			self.pairs.push((bounding_box, value));
		}

		true
	}

	pub fn find(&self, bounding_box: &Aabb) -> HashSet<T>
	where
		T: Eq + std::hash::Hash + Clone,
	{
		let mut set = HashSet::new();
		self.find_internal(bounding_box, &mut set);
		set
	}

	fn find_internal(&self, bounding_box: &Aabb, set: &mut HashSet<T>)
	where
		T: Eq + std::hash::Hash + Clone,
	{
		if !crate::aabb::bounding_boxes_overlap(bounding_box, &self.bounding_box) {
			return;
		}

		for (key, value) in &self.pairs {
			if crate::aabb::bounding_boxes_overlap(bounding_box, key) {
				set.insert(value.clone());
			}
		}

		if let Some(subtrees) = &self.subtrees {
			for tree in subtrees.iter() {
				tree.find_internal(bounding_box, set);
			}
		}
	}

	fn ensure_subtrees(&mut self) {
		if self.subtrees.is_some() {
			return;
		}

		let mid_x = (self.bounding_box.left + self.bounding_box.right) / 2.;
		let mid_y = (self.bounding_box.top + self.bounding_box.bottom) / 2.;

		self.subtrees = Some(Box::new([
			QuadTree::new(
				Aabb {
					top: self.bounding_box.top,
					right: mid_x,
					bottom: mid_y,
					left: self.bounding_box.left,
				},
				self.depth - 1,
				self.inner_node_capacity,
			),
			QuadTree::new(
				Aabb {
					top: self.bounding_box.top,
					right: self.bounding_box.right,
					bottom: mid_y,
					left: mid_x,
				},
				self.depth - 1,
				self.inner_node_capacity,
			),
			QuadTree::new(
				Aabb {
					top: mid_y,
					right: mid_x,
					bottom: self.bounding_box.bottom,
					left: self.bounding_box.left,
				},
				self.depth - 1,
				self.inner_node_capacity,
			),
			QuadTree::new(
				Aabb {
					top: mid_y,
					right: self.bounding_box.right,
					bottom: self.bounding_box.bottom,
					left: mid_x,
				},
				self.depth - 1,
				self.inner_node_capacity,
			),
		]));
	}
}