Fix panic when building STR tree with certain number of items; fix hanging when building empty index (#129)

* Fix building STR tree with certain number of items; fix hanging when building empty index

* Add changelog

Author: Kontinuation

CHANGELOG.md

@@ -2,6 +2,13 @@
 
 **This is the changelog for the core Rust library**. There's a [separate changelog](./python/CHANGELOG.md) for the Python bindings.
 
+## [0.3.1] - <RELEASE_DATE>
+
+## Bug fixes
+
+- Fix hanging when building RTree with 0 items by @kontinuation in https://github.com/kylebarron/geo-index/pull/129
+- Fix panic when building STR tree with certain number of items by @kontinuation in https://github.com/kylebarron/geo-index/pull/129
+
 ## [0.3.0] - 2025-06-13
 
 ### Breaking

src/indices.rs

@@ -71,6 +71,7 @@ impl MutableIndices<'_> {
         }
     }
 
+    #[allow(dead_code)]
     pub(crate) fn chunks_mut(&mut self, chunk_size: usize) -> Vec<MutableIndices> {
         match self {
             Self::U16(arr) => arr

src/rtree/builder.rs

@@ -142,8 +142,8 @@ impl<N: IndexableNum> RTreeBuilder<N> {
 
         let (boxes, mut indices) = split_data_borrow(&mut self.data, &self.metadata);
 
-        if self.metadata.num_items() == 1 {
-            // Only one item, we don't even have a root node to fill
+        if self.metadata.num_items() <= 1 {
+            // Empty or only one item, we don't even have a root node to fill
             return RTree {
                 buffer: self.data,
                 metadata: self.metadata,
@@ -253,7 +253,7 @@ fn split_data_borrow<'a, N: IndexableNum>(
 
 #[cfg(test)]
 mod test {
-    use crate::rtree::sort::HilbertSort;
+    use crate::rtree::sort::{HilbertSort, STRSort};
     use crate::rtree::RTreeIndex;
 
     use super::*;
@@ -266,4 +266,37 @@ mod test {
         let result = tree.search(0., 0., 0., 0.);
         assert_eq!(result, vec![0]);
     }
+
+    #[test]
+    fn build_str_index_with_various_items() {
+        let num_items_arr = [0, 1, 4, 8, 16, 20, 40, 80];
+        for num_items in num_items_arr {
+            build_index_with_various_items::<STRSort>(num_items);
+        }
+    }
+
+    #[test]
+    fn build_hilbert_index_with_various_items() {
+        let num_items_arr = [0, 1, 4, 8, 16, 20, 40, 80];
+        for num_items in num_items_arr {
+            build_index_with_various_items::<HilbertSort>(num_items);
+        }
+    }
+
+    fn build_index_with_various_items<S: Sort<f64>>(num_items: u32) {
+        let mut builder = RTreeBuilder::<f64>::new(num_items);
+        for i in 0..num_items {
+            builder.add(i as f64, i as f64, i as f64, i as f64);
+        }
+        let tree = builder.finish::<S>();
+        assert_eq!(tree.num_items(), num_items);
+        if num_items == 0 {
+            assert!(tree.search(0., 0., 0., 0.).is_empty());
+        } else {
+            for i in 0..num_items {
+                let results = tree.search(i as f64, i as f64, i as f64, i as f64);
+                assert_eq!(results, vec![i]);
+            }
+        }
+    }
 }

src/rtree/sort/str.rs

@@ -47,7 +47,8 @@ impl<N: IndexableNum> Sort<N> for STRSort {
 
         let num_leaf_nodes = (params.num_items as f64 / params.node_size as f64).ceil();
         let num_vertical_slices = num_leaf_nodes.sqrt().ceil() as usize;
-        let num_items_per_slice = num_vertical_slices * params.node_size;
+        let num_items_per_slice =
+            (params.num_items as f64 / num_vertical_slices as f64).ceil() as usize;
 
         #[cfg(feature = "rayon")]
         {

src/rtree/trait.rs

@@ -74,8 +74,11 @@ pub trait RTreeIndex<N: IndexableNum>: Sized {
     fn search(&self, min_x: N, min_y: N, max_x: N, max_y: N) -> Vec<u32> {
         let boxes = self.boxes();
         let indices = self.indices();
+        if boxes.is_empty() {
+            return vec![];
+        }
 
-        let mut outer_node_index = Some(boxes.len() - 4);
+        let mut outer_node_index = boxes.len().checked_sub(4);
 
         let mut queue = vec![];
         let mut results = vec![];

src/rtree/util.rs

@@ -12,7 +12,7 @@ pub(crate) fn compute_num_nodes(num_items: u32, node_size: u16) -> (usize, Vec<u
     let mut n = num_items;
     let mut num_nodes = n;
     let mut level_bounds = vec![n * 4];
-    while n != 1 {
+    while n > 1 {
         n = (n as f64 / node_size as f64).ceil() as usize;
         num_nodes += n;
         level_bounds.push(num_nodes * 4);