From 2253102f2488ce9aa21f17512dd61d971d65dc70 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Sun, 5 Jan 2025 11:59:42 -0500 Subject: [PATCH] KDTree traversal --- src/kdtree/traversal.rs | 96 ++++++++++++++++++++++++++++++++++++++--- src/rtree/traversal.rs | 50 +++++++++++++++------ 2 files changed, 129 insertions(+), 17 deletions(-) diff --git a/src/kdtree/traversal.rs b/src/kdtree/traversal.rs index e293af1..e1382f4 100644 --- a/src/kdtree/traversal.rs +++ b/src/kdtree/traversal.rs @@ -17,15 +17,22 @@ pub struct Node<'a, N: IndexableNum, T: KDTreeIndex> { /// TODO: switch to bool axis: usize, + /// The index of the right child. right_child: usize, - left_child: usize, - phantom: PhantomData, + /// The index of the left child. + left_child: usize, + /// The min_x of this node. min_x: N, + /// The min_y of this node. min_y: N, + /// The max_x of this node. max_x: N, + /// The max_y of this node. max_y: N, + + phantom: PhantomData, } impl<'a, N: IndexableNum, T: KDTreeIndex> Node<'a, N, T> { @@ -35,11 +42,11 @@ impl<'a, N: IndexableNum, T: KDTreeIndex> Node<'a, N, T> { axis: 0, right_child: tree.indices().len() - 1, left_child: 0, - phantom: PhantomData, min_x: N::max_value(), min_y: N::max_value(), max_x: N::min_value(), max_y: N::min_value(), + phantom: PhantomData, } } @@ -60,10 +67,23 @@ impl<'a, N: IndexableNum, T: KDTreeIndex> Node<'a, N, T> { (x, y) } + /// The child node representing the "left" half. + /// + /// Returns `None` if [`Self::is_parent`] is `false`. + pub fn left_child(&self) -> Option> { + if self.is_parent() { + Some(self.left_child_unchecked()) + } else { + None + } + } + /// The child node representing the "left" half. /// /// Note that this **does not include** the middle index of the current node. - pub fn left_child(&self) -> Node<'_, N, T> { + pub fn left_child_unchecked(&self) -> Node<'_, N, T> { + debug_assert!(self.is_parent()); + let m = self.middle_index(); let (x, y) = self.middle_xy(m); @@ -88,10 +108,23 @@ impl<'a, N: IndexableNum, T: KDTreeIndex> Node<'a, N, T> { } } + /// The child node representing the "right" half. + /// + /// Returns `None` if [`Self::is_parent`] is `false`. + pub fn right_child(&self) -> Option> { + if self.is_parent() { + Some(self.right_child_unchecked()) + } else { + None + } + } + /// The child node representing the "right" half. /// /// Note that this **does not include** the middle index of the current node. - pub fn right_child(&self) -> Node<'_, N, T> { + pub fn right_child_unchecked(&self) -> Node<'_, N, T> { + debug_assert!(self.is_parent()); + let m = self.middle_index(); let (x, y) = self.middle_xy(m); @@ -127,6 +160,59 @@ impl<'a, N: IndexableNum, T: KDTreeIndex> Node<'a, N, T> { pub fn is_parent(&self) -> bool { !self.is_leaf() } + + /// The original insertion index of the "middle child" of this node. This is only valid when + /// this is a parent node, which you can check with `Self::is_parent`. + /// + /// Returns `None` if [`Self::is_parent`] is `false`. + #[inline] + pub fn middle_insertion_index(&self) -> Option { + if self.is_parent() { + Some(self.middle_insertion_index_unchecked()) + } else { + None + } + } + + /// The original insertion index of the "middle child" of this node. This is only valid when + /// this is a parent node, which you can check with `Self::is_parent`. + #[inline] + pub fn middle_insertion_index_unchecked(&self) -> u32 { + debug_assert!(self.is_parent()); + + let m = self.middle_index(); + let indices = self.tree.indices(); + indices.get(m) as u32 + } + + /// The original insertion indices. This is only valid when this is a leaf node, which you can + /// check with `Self::is_leaf`. + /// + /// Returns `None` if [`Self::is_leaf`] is `false`. + #[inline] + pub fn leaf_insertion_indices(&self) -> Option> { + if self.is_leaf() { + Some(self.leaf_insertion_indices_unchecked()) + } else { + None + } + } + + /// The original insertion indices. This is only valid when this is a leaf node, which you can + /// check with `Self::is_leaf`. + #[inline] + pub fn leaf_insertion_indices_unchecked(&self) -> Vec { + debug_assert!(self.is_leaf()); + + let mut result = Vec::with_capacity(self.tree.node_size() as _); + + let indices = self.tree.indices(); + for i in self.left_child..self.right_child + 1 { + result.push(indices.get(i) as u32); + } + + result + } } impl> RectTrait for Node<'_, N, T> { diff --git a/src/rtree/traversal.rs b/src/rtree/traversal.rs index d0501b9..45822ab 100644 --- a/src/rtree/traversal.rs +++ b/src/rtree/traversal.rs @@ -108,9 +108,20 @@ impl<'a, N: IndexableNum, T: RTreeIndex> Node<'a, N, T> { true } - /// Returns an iterator over the child nodes of this node. This must only be called if - /// `is_parent` is `true`. - pub fn children(&self) -> impl Iterator> { + /// Returns an iterator over the child nodes of this node. + /// + /// Returns `None` if [`Self::is_parent`] is `false`. + pub fn children(&self) -> Option>> { + if self.is_parent() { + Some(self.children_unchecked()) + } else { + None + } + } + + /// Returns an iterator over the child nodes of this node. This is only valid when + /// [`Self::is_parent`] is `true`. + pub fn children_unchecked(&self) -> impl Iterator> { debug_assert!(self.is_parent()); // find the start and end indexes of the children of this node @@ -123,12 +134,25 @@ impl<'a, N: IndexableNum, T: RTreeIndex> Node<'a, N, T> { .map(|pos| Node::new(self.tree, pos)) } + /// The original insertion index. This is only valid when this is a leaf node, which you can + /// check with `Self::is_leaf`. + /// + /// Returns `None` if [`Self::is_leaf`] is `false`. + #[inline] + pub fn insertion_index(&self) -> Option { + if self.is_leaf() { + Some(self.insertion_index_unchecked()) + } else { + None + } + } + /// The original insertion index. This is only valid when this is a leaf node, which you can /// check with `Self::is_leaf`. #[inline] - pub fn index(&self) -> usize { + pub fn insertion_index_unchecked(&self) -> u32 { debug_assert!(self.is_leaf()); - self.tree.indices().get(self.pos >> 2) + self.tree.indices().get(self.pos >> 2) as u32 } } @@ -215,12 +239,14 @@ where return; } - let children1 = parent1.children().filter(|c1| c1.intersects(parent2)); + let children1 = parent1 + .children_unchecked() + .filter(|c1| c1.intersects(parent2)); let mut children2 = take(&mut self.candidates); children2.extend( parent2 - .children() + .children_unchecked() .filter(|c2| c2.intersects(parent1)) .map(|c| c.pos), ); @@ -251,15 +277,15 @@ where match (left.is_leaf(), right.is_leaf()) { (true, true) => { return Some(( - left.index().try_into().unwrap(), - right.index().try_into().unwrap(), + left.insertion_index_unchecked(), + right.insertion_index_unchecked(), )) } (true, false) => right - .children() + .children_unchecked() .for_each(|c| self.push_if_intersecting(&left, &c)), (false, true) => left - .children() + .children_unchecked() .for_each(|c| self.push_if_intersecting(&c, &right)), (false, false) => self.add_intersecting_children(&left, &right), } @@ -292,7 +318,7 @@ mod test { assert!(root_node.is_parent()); let level_1_boxes = tree.boxes_at_level(1).unwrap(); - let level_1 = root_node.children().collect::>(); + let level_1 = root_node.children_unchecked().collect::>(); assert_eq!(level_1.len(), level_1_boxes.len() / 4); } }