From 34d831f5fd2a71fe9fad866a8c05fe27bfdba352 Mon Sep 17 00:00:00 2001 From: Ryan Date: Mon, 6 Jun 2022 23:20:34 -0700 Subject: [PATCH] Add R-Tree --- Cargo.toml | 2 + src/rtree.rs | 708 +++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 710 insertions(+) create mode 100644 src/rtree.rs diff --git a/Cargo.toml b/Cargo.toml index 1733f53..227512b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,6 +10,7 @@ build = "build/main.rs" [dependencies] aes = "0.7" anyhow = "1" +approx = "0.5.1" arrayvec = "0.7" ascii = "1" async-trait = "0.1" @@ -25,6 +26,7 @@ futures = "0.3" hematite-nbt = "0.5" log = "0.4" num = "0.4" +ordered-float = "3.0.0" parking_lot = "0.12" paste = "1" rand = "0.8" diff --git a/src/rtree.rs b/src/rtree.rs new file mode 100644 index 0000000..4033289 --- /dev/null +++ b/src/rtree.rs @@ -0,0 +1,708 @@ +use std::mem; + +#[cfg(test)] +use approx::relative_eq; +use arrayvec::ArrayVec; +use ordered_float::OrderedFloat; +use vek::Aabb; + +pub struct RTree { + root: Node, + // The bufs are put here to reuse their allocations. + internal_split_buf: InternalBuf, + leaf_split_buf: LeafBuf, + reinsert_buf: LeafBuf, +} + +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub enum QueryAction { + Continue, + Break, +} + +type InternalBuf = Vec<(Box>, Aabb)>; +type LeafBuf = Vec<(T, Aabb)>; + +enum Node { + Internal(ArrayVec<(Box>, Aabb), MAX>), + Leaf(ArrayVec<(T, Aabb), MAX>), +} + +impl RTree { + pub fn new() -> Self { + assert!( + MIN >= 2 && MIN <= MAX / 2 && MAX >= 2, + "invalid R-Tree configuration" + ); + + Self { + root: Node::Leaf(ArrayVec::new()), + internal_split_buf: Vec::new(), + leaf_split_buf: Vec::new(), + reinsert_buf: Vec::new(), + } + } + + pub fn insert(&mut self, data: T, data_aabb: Aabb) { + if let InsertResult::Split(new_node) = self.root.insert( + data, + data_aabb, + &mut self.internal_split_buf, + &mut self.leaf_split_buf, + ) { + let root_aabb = self.root.bounds(); + let new_node_aabb = new_node.bounds(); + + let old_root = mem::replace(&mut self.root, Node::Internal(ArrayVec::new())); + + match &mut self.root { + Node::Internal(children) => { + children.push((Box::new(old_root), root_aabb)); + children.push((new_node, new_node_aabb)); + } + Node::Leaf(_) => unreachable!(), + } + } + } + + pub fn retain( + &mut self, + mut collides: impl FnMut(Aabb) -> bool, + mut retain: impl FnMut(&mut T, &mut Aabb) -> bool, + ) { + self.root + .retain(None, &mut collides, &mut retain, &mut self.reinsert_buf); + + if let Node::Internal(children) = &mut self.root { + if children.len() == 1 { + let new_root = *children.drain(..).next().unwrap().0; + self.root = new_root; + } else if children.is_empty() { + self.root = Node::Leaf(ArrayVec::new()); + } + } + + let mut reinsert_buf = mem::take(&mut self.reinsert_buf); + + for (data, data_aabb) in reinsert_buf.drain(..) { + self.insert(data, data_aabb); + } + + debug_assert!(self.reinsert_buf.capacity() == 0); + self.reinsert_buf = reinsert_buf; + + // Don't waste too much memory after a large restructuring. + self.reinsert_buf.shrink_to(16); + } + + pub fn query( + &self, + mut collides: impl FnMut(Aabb) -> bool, + mut callback: impl FnMut(&T, Aabb) -> QueryAction, + ) { + self.root.query(&mut collides, &mut callback); + } + + pub fn clear(&mut self) { + self.root = Node::Leaf(ArrayVec::new()); + } + + pub fn depth(&self) -> usize { + self.root.depth(0) + } + + /// For the purposes of rendering the R-Tree. + pub fn visit(&self, mut f: impl FnMut(Aabb, usize)) { + self.root.visit(&mut f, 1); + if self.root.children_count() != 0 { + f(self.root.bounds(), 0); + } + } + + #[cfg(test)] + fn check_invariants(&self, expected_len: usize) { + assert!(self.internal_split_buf.is_empty()); + assert!(self.leaf_split_buf.is_empty()); + assert!(self.reinsert_buf.is_empty()); + + if let Node::Internal(children) = &self.root { + assert!( + children.len() != 1, + "internal root with a single entry should become the child" + ); + assert!(!children.is_empty(), "empty internal root should be a leaf"); + } + + let mut len_counter = 0; + + self.root.check_invariants(None, 0, &mut len_counter); + + assert_eq!( + len_counter, expected_len, + "unexpected number of entries in rtree" + ) + } +} + +impl Node { + fn bounds(&self) -> Aabb { + match self { + Node::Internal(children) => children + .iter() + .map(|(_, aabb)| *aabb) + .reduce(Aabb::union) + .unwrap(), + Node::Leaf(children) => children + .iter() + .map(|(_, aabb)| *aabb) + .reduce(Aabb::union) + .unwrap(), + } + } + + fn children_count(&self) -> usize { + match self { + Node::Internal(children) => children.len(), + Node::Leaf(children) => children.len(), + } + } + + fn insert( + &mut self, + data: T, + data_aabb: Aabb, + internal_split_buf: &mut InternalBuf, + leaf_split_buf: &mut LeafBuf, + ) -> InsertResult { + match self { + Self::Internal(children) => { + let children_is_full = children.is_full(); + + let (best_child, best_child_aabb) = { + let best = area_insertion_heuristic(data_aabb, children); + &mut children[best] + }; + + match best_child.insert(data, data_aabb, internal_split_buf, leaf_split_buf) { + InsertResult::Ok => { + best_child_aabb.expand_to_contain(data_aabb); + InsertResult::Ok + } + InsertResult::Split(new_node) => { + let new_node_aabb = new_node.bounds(); + *best_child_aabb = best_child.bounds(); + + if children_is_full { + let other = split_node::<_, MIN, MAX>( + internal_split_buf, + children, + (new_node, new_node_aabb), + ); + InsertResult::Split(Box::new(Node::Internal(other))) + } else { + children.push((new_node, new_node_aabb)); + InsertResult::Ok + } + } + } + } + Self::Leaf(children) => { + if children.is_full() { + let other = + split_node::<_, MIN, MAX>(leaf_split_buf, children, (data, data_aabb)); + debug_assert!(other.len() >= MIN); + + InsertResult::Split(Box::new(Node::Leaf(other))) + } else { + children.push((data, data_aabb)); + InsertResult::Ok + } + } + } + } + + fn query( + &self, + collides: &mut impl FnMut(Aabb) -> bool, + callback: &mut impl FnMut(&T, Aabb) -> QueryAction, + ) -> QueryAction { + match self { + Node::Internal(children) => { + for child in children { + if collides(child.1) { + if let QueryAction::Break = child.0.query(collides, callback) { + return QueryAction::Break; + } + } + } + } + Node::Leaf(children) => { + for (child, child_aabb) in children { + if collides(*child_aabb) { + if let QueryAction::Break = callback(child, *child_aabb) { + return QueryAction::Break; + } + } + } + } + } + QueryAction::Continue + } + + fn retain( + &mut self, + bounds: Option>, // `None` when self is root. + collides: &mut impl FnMut(Aabb) -> bool, + retain: &mut impl FnMut(&mut T, &mut Aabb) -> bool, + reinsert_buf: &mut LeafBuf, + ) -> RetainResult { + match self { + Node::Internal(children) => { + let mut recalculate_bounds = false; + + children.retain(|(child, child_aabb)| { + if collides(*child_aabb) { + match child.retain(Some(*child_aabb), collides, retain, reinsert_buf) { + RetainResult::Ok => true, + RetainResult::Deleted => { + recalculate_bounds = true; + false + } + RetainResult::ShrunkAabb(new_aabb) => { + *child_aabb = new_aabb; + recalculate_bounds = true; + true + } + } + } else { + true + } + }); + + if let Some(bounds) = bounds { + if children.len() < MIN { + for (child, _) in children.drain(..) { + child.collect_orphans(reinsert_buf); + } + RetainResult::Deleted + } else if recalculate_bounds { + let new_bounds = self.bounds(); + debug_assert!(bounds.contains_aabb(new_bounds)); + + if bounds != new_bounds { + RetainResult::ShrunkAabb(new_bounds) + } else { + RetainResult::Ok + } + } else { + RetainResult::Ok + } + } else { + RetainResult::Ok + } + } + Node::Leaf(children) => { + let mut recalculate_bounds = false; + + let mut i = 0; + while i < children.len() { + let (child, child_aabb) = &mut children[i]; + let before = *child_aabb; + if collides(before) { + if retain(child, child_aabb) { + let after = *child_aabb; + if before != after { + if let Some(bounds) = bounds { + recalculate_bounds = true; + // A child can move within a leaf node without reinsertion + // as long as it does not increase the bounds of the leaf. + if !bounds.contains_aabb(after) { + reinsert_buf.push(children.swap_remove(i)); + } else { + i += 1; + } + } else { + i += 1; + } + } else { + i += 1; + } + } else { + recalculate_bounds = true; + children.swap_remove(i); + } + } else { + i += 1; + } + } + + if let Some(bounds) = bounds { + if children.len() < MIN { + reinsert_buf.extend(children.drain(..)); + RetainResult::Deleted + } else if recalculate_bounds { + let new_bounds = self.bounds(); + debug_assert!(bounds.contains_aabb(new_bounds)); + + if bounds != new_bounds { + RetainResult::ShrunkAabb(new_bounds) + } else { + RetainResult::Ok + } + } else { + RetainResult::Ok + } + } else { + RetainResult::Ok + } + } + } + } + + fn collect_orphans(self, reinsert_buf: &mut LeafBuf) { + match self { + Node::Internal(children) => { + for (child, _) in children { + child.collect_orphans(reinsert_buf); + } + } + Node::Leaf(children) => reinsert_buf.extend(children), + } + } + + fn depth(&self, level: usize) -> usize { + match self { + Node::Internal(children) => children[0].0.depth(level + 1), + Node::Leaf(_) => level, + } + } + + fn visit(&self, f: &mut impl FnMut(Aabb, usize), level: usize) { + match self { + Node::Internal(children) => { + for (child, child_aabb) in children { + child.visit(f, level + 1); + f(*child_aabb, level); + } + } + Node::Leaf(children) => { + for (_, child_aabb) in children { + f(*child_aabb, level); + } + } + } + } + + #[cfg(test)] + fn check_invariants( + &self, + bounds: Option>, + depth: usize, + len_counter: &mut usize, + ) -> usize { + let mut child_depth = None; + + match self { + Node::Internal(children) => { + assert!(!children.is_empty()); + + if let Some(bounds) = bounds { + let tight_bounds = self.bounds(); + assert!( + relative_eq!(tight_bounds.min, bounds.min) + && relative_eq!(tight_bounds.max, bounds.max), + "bounding rectangle for internal node is not tight" + ); + } + + for (child, child_aabb) in children { + let d = child.check_invariants(Some(*child_aabb), depth + 1, len_counter); + if let Some(child_depth) = &mut child_depth { + assert_eq!(*child_depth, d, "rtree is not balanced"); + } else { + child_depth = Some(d); + } + } + } + Node::Leaf(children) => { + if let Some(bounds) = bounds { + let tight_bounds = self.bounds(); + assert!( + relative_eq!(tight_bounds.min, bounds.min) + && relative_eq!(tight_bounds.max, bounds.max), + "bounding rectangle for leaf node is not tight" + ); + } + + *len_counter += children.len(); + child_depth = Some(depth); + } + } + + if let Some(bounds) = bounds { + assert!(bounds == self.bounds()); + } + + child_depth.unwrap() + } +} +enum InsertResult { + /// No split occurred. + Ok, + /// Contains the new node that was split off. + Split(Box>), +} + +enum RetainResult { + /// Nothing changed. + Ok, + /// This node must be deleted from its parent. + Deleted, + /// This node was not deleted but its AABR was shrunk. + /// Contains the new AABR. + ShrunkAabb(Aabb), +} + +fn area_insertion_heuristic(data_aabb: Aabb, children: &[(T, Aabb)]) -> usize { + debug_assert!( + !children.is_empty(), + "internal node must have at least one child" + ); + + let mut best = 0; + let mut best_area_increase = f64::INFINITY; + let mut best_aabb = Aabb::default(); + + for (idx, (_, child_aabb)) in children.iter().enumerate() { + let area_increase = volume(child_aabb.union(data_aabb)) - volume(*child_aabb); + if area_increase < best_area_increase { + best = idx; + best_area_increase = area_increase; + best_aabb = *child_aabb; + } else if area_increase == best_area_increase && volume(*child_aabb) < volume(best_aabb) { + best = idx; + best_aabb = *child_aabb; + } + } + + best +} + +/// Splits a node with `children` being the children of the node being split. +/// +/// After returning, `children` contains half the data while the returned +/// `ArrayVec` contains the other half for the new node. +fn split_node( + split_buf: &mut Vec<(T, Aabb)>, + children: &mut ArrayVec<(T, Aabb), MAX>, + data: (T, Aabb), +) -> ArrayVec<(T, Aabb), MAX> { + split_buf.extend(children.take()); + split_buf.push(data); + + let dists = MIN..MAX - MIN + 2; + + let bb = |es: &[(T, Aabb)]| es.iter().map(|e| e.1).reduce(Aabb::union).unwrap(); + + let mut sum_x = 0.0; + split_buf.sort_unstable_by_key(|e| OrderedFloat(e.1.min.x / 2.0 + e.1.max.x / 2.0)); + + for split in dists.clone() { + sum_x += surface_area(bb(&split_buf[..split])) + surface_area(bb(&split_buf[split..])); + } + + let mut sum_y = 0.0; + split_buf.sort_unstable_by_key(|e| OrderedFloat(e.1.min.y / 2.0 + e.1.max.y / 2.0)); + + for split in dists.clone() { + sum_y += surface_area(bb(&split_buf[..split])) + surface_area(bb(&split_buf[split..])); + } + + let mut sum_z = 0.0; + split_buf.sort_unstable_by_key(|e| OrderedFloat(e.1.min.z / 2.0 + e.1.max.z / 2.0)); + + for split in dists.clone() { + sum_z += surface_area(bb(&split_buf[..split])) + surface_area(bb(&split_buf[split..])); + } + + // Sort by the winning axis + split_buf.sort_unstable_by_key(|e| { + let (min, max) = if sum_x <= sum_y && sum_x <= sum_z { + (e.1.min.x, e.1.max.x) + } else if sum_y <= sum_x && sum_y <= sum_z { + (e.1.min.y, e.1.max.y) + } else { + (e.1.min.z, e.1.max.z) + }; + + OrderedFloat(min / 2.0 + max / 2.0) + }); + + let mut best_dist = 0; + let mut best_overlap_value = f64::INFINITY; + let mut best_area_value = f64::INFINITY; + + for split in dists { + let group_1 = bb(&split_buf[..split]); + let group_2 = bb(&split_buf[split..]); + let overlap_value = { + let int = group_1.intersection(group_2); + if int.is_valid() { + volume(int) + } else { + 0.0 + } + }; + let area_value = volume(group_1) + volume(group_2); + + if overlap_value < best_overlap_value { + best_overlap_value = overlap_value; + best_area_value = area_value; + best_dist = split; + } else if overlap_value == best_overlap_value && area_value < best_area_value { + best_area_value = area_value; + best_dist = split; + } + } + + debug_assert!(children.is_empty()); + debug_assert_eq!(split_buf.len(), MAX + 1); + + let mut other = ArrayVec::new(); + other.extend(split_buf.drain(best_dist..)); + + children.extend(split_buf.drain(..)); + + other +} + +fn volume(aabb: Aabb) -> f64 { + (aabb.max - aabb.min).product() +} + +fn surface_area(aabb: Aabb) -> f64 { + let d = aabb.max - aabb.min; + (d.x * d.y + d.x * d.z + d.y * d.z) * 2.0 +} + +#[cfg(test)] +mod tests { + use std::f64::consts::TAU; + use std::sync::atomic::{AtomicU64, Ordering}; + + use rand::Rng; + use vek::Vec3; + + use super::*; + + fn insert_rand( + rtree: &mut RTree, + ) -> (u64, Aabb) { + static NEXT_UNIQUE_ID: AtomicU64 = AtomicU64::new(0); + + let id = NEXT_UNIQUE_ID.fetch_add(1, Ordering::SeqCst); + + let mut rng = rand::thread_rng(); + + let min = Vec3::new(rng.gen(), rng.gen(), rng.gen()); + let max = Vec3::new( + min.x + rng.gen_range(0.003..=0.01), + min.y + rng.gen_range(0.003..=0.01), + min.z + rng.gen_range(0.003..=0.01), + ); + + let aabb = Aabb { min, max }; + + rtree.insert(id, aabb); + + (id, aabb) + } + + #[test] + fn insert_delete_interleaved() { + let mut rtree: RTree = RTree::new(); + + for i in 0..5_000 { + insert_rand(&mut rtree); + let (id_0, aabb_0) = insert_rand(&mut rtree); + + let mut found = false; + rtree.retain( + |aabb| aabb.collides_with_aabb(aabb_0), + |&mut id, _| { + if id == id_0 { + assert!(!found); + found = true; + false + } else { + true + } + }, + ); + assert!(found); + + rtree.check_invariants(i + 1); + } + } + + #[test] + fn node_underfill() { + let mut rtree: RTree = RTree::new(); + + for i in 0..5_000 { + insert_rand(&mut rtree); + rtree.check_invariants(i + 1); + } + + let mut delete_count = 0; + + rtree.retain( + |_| true, + |_, _| { + if rand::random() { + delete_count += 1; + false + } else { + true + } + }, + ); + rtree.check_invariants(5_000 - delete_count); + + rtree.clear(); + rtree.check_invariants(0); + } + + #[test] + fn movement() { + let mut rtree: RTree = RTree::new(); + + for _ in 0..5_000 { + insert_rand(&mut rtree); + } + + let mut rng = rand::thread_rng(); + for _ in 0..100 { + rtree.retain( + |_| true, + |_, aabb| { + let angle = rng.gen_range(0.0..TAU); + let z: f64 = rng.gen_range(-1.0..=1.0); + + let v = Vec3::new( + (1.0 - z * z).sqrt() * angle.cos(), + (1.0 - z * z).sqrt() * angle.sin(), + z, + ) * 0.03; + + aabb.min += v; + aabb.max += v; + assert!(aabb.is_valid()); + + true + }, + ); + rtree.check_invariants(5_000); + } + } +}