KDTree traversal

src/kdtree/traversal.rs

@@ -17,15 +17,22 @@ pub struct Node<'a, N: IndexableNum, T: KDTreeIndex<N>> {
     /// TODO: switch to bool
     axis: usize,
 
+    /// The index of the right child.
     right_child: usize,
-    left_child: usize,
 
-    phantom: PhantomData<N>,
+    /// 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<N>,
 }
 
 impl<'a, N: IndexableNum, T: KDTreeIndex<N>> Node<'a, N, T> {
@@ -35,11 +42,11 @@ impl<'a, N: IndexableNum, T: KDTreeIndex<N>> 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<N>> 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<Node<'_, N, T>> {
+        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<N>> 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<Node<'_, N, T>> {
+        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<N>> 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<u32> {
+        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<Vec<u32>> {
+        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<u32> {
+        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<N: IndexableNum, T: KDTreeIndex<N>> RectTrait for Node<'_, N, T> {

src/rtree/traversal.rs

@@ -108,9 +108,20 @@ impl<'a, N: IndexableNum, T: RTreeIndex<N>> 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<Item = Node<'_, N, T>> {
+    /// Returns an iterator over the child nodes of this node.
+    ///
+    /// Returns `None` if [`Self::is_parent`] is `false`.
+    pub fn children(&self) -> Option<impl Iterator<Item = Node<'_, N, T>>> {
+        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<Item = Node<'_, N, T>> {
         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<N>> 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<u32> {
+        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::<Vec<_>>();
+        let level_1 = root_node.children_unchecked().collect::<Vec<_>>();
         assert_eq!(level_1.len(), level_1_boxes.len() / 4);
     }
 }