150 lines
3.9 KiB
Rust
150 lines
3.9 KiB
Rust
use bevy::{
|
|
math::bounding::{Aabb2d, BoundingVolume, IntersectsVolume},
|
|
prelude::Vec2,
|
|
};
|
|
|
|
use crate::geom::Point;
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub enum Entry {
|
|
Leaf { mbr: Aabb2d, object: Point },
|
|
Node { mbr: Aabb2d, child: TreeNode },
|
|
}
|
|
|
|
impl Entry {
|
|
/// Returns a reference to the minimum bounding volume for this entry.
|
|
pub fn mbr(&self) -> &Aabb2d {
|
|
match self {
|
|
Entry::Leaf { mbr, .. } | Entry::Node { mbr, .. } => mbr,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct TreeNode {
|
|
pub entries: Vec<Entry>,
|
|
pub is_leaf: bool,
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct RStarTree {
|
|
root: TreeNode,
|
|
max_entries: usize,
|
|
}
|
|
|
|
impl Entry {
|
|
fn as_leaf_obj(&self) -> Option<&Point> {
|
|
match self {
|
|
Entry::Leaf { object, .. } => Some(object),
|
|
_ => None,
|
|
}
|
|
}
|
|
fn child(&self) -> Option<&TreeNode> {
|
|
match self {
|
|
Entry::Node { child, .. } => Some(child),
|
|
_ => None,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl TreeNode {
|
|
fn is_leaf(&self) -> bool {
|
|
self.is_leaf
|
|
}
|
|
|
|
fn entries(&self) -> &[Entry] {
|
|
&self.entries
|
|
}
|
|
|
|
fn range_search_bbox<'s>(&'s self, bbox: &Aabb2d, out: &mut Vec<&'s Point>) {
|
|
if self.is_leaf() {
|
|
for entry in self.entries() {
|
|
if let Some(obj) = entry.as_leaf_obj()
|
|
&& entry.mbr().intersects(bbox)
|
|
{
|
|
out.push(obj);
|
|
}
|
|
}
|
|
} else {
|
|
for entry in self.entries() {
|
|
if let Some(child) = entry.child()
|
|
&& entry.mbr().intersects(bbox)
|
|
{
|
|
child.range_search_bbox(bbox, out);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
fn mbr(&self) -> Option<Aabb2d> {
|
|
entries_mbr(self.entries())
|
|
}
|
|
}
|
|
|
|
impl RStarTree {
|
|
pub fn new(max_entries: usize) -> Self {
|
|
let max_entries = max_entries.max(4);
|
|
RStarTree {
|
|
root: TreeNode {
|
|
entries: Vec::new(),
|
|
is_leaf: true,
|
|
},
|
|
max_entries,
|
|
}
|
|
}
|
|
|
|
pub fn range_search(&self, query_point: &Point, radius: f32) -> Vec<&Point> {
|
|
let query_bbox = Aabb2d::new(query_point.point, Vec2::splat(radius));
|
|
let r2 = radius.powi(2);
|
|
let candidates = self.range_search_bbox(&query_bbox);
|
|
candidates
|
|
.into_iter()
|
|
.filter(|&other| query_point.point.distance_squared(other.point) <= r2)
|
|
.collect()
|
|
}
|
|
|
|
pub fn range_search_bbox(&self, query_bbox: &Aabb2d) -> Vec<&Point> {
|
|
let mut res = Vec::with_capacity(self.root.entries.len() / 10);
|
|
self.root.range_search_bbox(query_bbox, &mut res);
|
|
res
|
|
}
|
|
|
|
pub fn insert_bulk(&mut self, objects: Vec<Point>) {
|
|
if objects.is_empty() {
|
|
return;
|
|
}
|
|
|
|
let mut entries: Vec<Entry> = objects
|
|
.into_iter()
|
|
.map(|point| Entry::Leaf {
|
|
mbr: point.mbr(),
|
|
object: point,
|
|
})
|
|
.collect();
|
|
|
|
while entries.len() > self.max_entries {
|
|
let mut new_level_entries = Vec::new();
|
|
let chunks = entries.chunks(self.max_entries);
|
|
|
|
for chunk in chunks {
|
|
let child = TreeNode {
|
|
entries: chunk.to_vec(),
|
|
is_leaf: self.root.is_leaf,
|
|
};
|
|
if let Some(mbr) = child.mbr() {
|
|
new_level_entries.push(Entry::Node { mbr, child });
|
|
}
|
|
}
|
|
entries = new_level_entries;
|
|
self.root.is_leaf = false;
|
|
}
|
|
|
|
self.root.entries.extend(entries);
|
|
}
|
|
}
|
|
|
|
fn entries_mbr(entries: &[Entry]) -> Option<Aabb2d> {
|
|
let mut iter = entries.iter();
|
|
let first = *iter.next()?.mbr();
|
|
Some(iter.fold(first, |acc, entry| acc.merge(entry.mbr())))
|
|
}
|