122 lines
2.6 KiB
Rust
122 lines
2.6 KiB
Rust
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,
|
|
),
|
|
]));
|
|
}
|
|
}
|