fix: Fix stack overflow in degenerate cases (#137)

* fix: Fix Stack Overflow in degenerate cases

* Refactor swap impl

* Apply upstream fix

* Bump flatbush version in tests

Author: kylebarron

benches/README.md

@@ -1,3 +1,5 @@
+If the download link is dead, find the most recent download link here: https://github.com/microsoft/USBuildingFootprints
+
 ```
 wget https://usbuildingdata.blob.core.windows.net/usbuildings-v2/Utah.geojson.zip
 pip install -U geopandas pyogrio

fixtures/README.md

@@ -0,0 +1,3 @@
+Fixtures data for test
+
+- Run `npm run generate_fixtures` to generate data from the upstream `flatbush` library.

fixtures/data1_flatbush_js.raw

@@ -1,11 +1,13 @@
 {
   "type": "module",
-  "scripts": {},
+  "scripts": {
+    "generate_fixtures": "node fixtures/flatbush_generate.js"
+  },
   "volta": {
     "node": "20.9.0"
   },
   "devDependencies": {
     "@tsconfig/node20": "^20.1.2",
-    "flatbush": "^4.2.0"
+    "flatbush": "^4.5.0"
   }
 }

package-lock.json

@@ -223,11 +223,47 @@ impl<'a, N: IndexableNum> RTreeRef<'a, N> {
 
 #[cfg(test)]
 mod tests {
+    use crate::rtree::sort::{HilbertSort, STRSort, Sort};
+    use crate::rtree::{RTreeBuilder, RTreeIndex};
+
     use super::*;
 
     #[test]
     fn rejects_short_buffers() {
         assert!(RTreeMetadata::<f64>::from_slice(&[]).is_err());
         assert!(RTreeMetadata::<f64>::from_slice(&[0; 7]).is_err());
     }
+
+    fn linspace(start: usize, stop: usize, num: usize, endpoint: bool) -> Vec<f64> {
+        let div = if endpoint { num - 1 } else { num };
+        let step = (stop - start) as f64 / div as f64;
+        (0..num).map(|i| start as f64 + step * i as f64).collect()
+    }
+
+    #[test]
+    // https://github.com/mourner/flatbush/pull/65
+    fn quicksort_should_work_with_an_inbalanced_dataset() {
+        _quicksort_should_work_with_an_inbalanced_dataset::<HilbertSort>();
+        _quicksort_should_work_with_an_inbalanced_dataset::<STRSort>();
+    }
+
+    fn _quicksort_should_work_with_an_inbalanced_dataset<S: Sort<f64>>() {
+        let n = 15000;
+        let mut builder = RTreeBuilder::new(2 * n);
+
+        let items = linspace(0, 1000, n as usize, true);
+        let items2 = linspace(0, 1000, n as usize, true);
+
+        for item in items {
+            builder.add(item, 0., item, 0.);
+        }
+
+        for item in items2 {
+            builder.add(item, 0., item, 0.);
+        }
+
+        let index = builder.finish::<S>();
+
+        index.search(-100., -1., 15000., 1.);
+    }
 }

package.json

@@ -1,5 +1,6 @@
 use crate::indices::MutableIndices;
 use crate::r#type::IndexableNum;
+use crate::rtree::sort::util::swap;
 use crate::rtree::sort::{Sort, SortParams};
 
 /// An implementation of hilbert sorting.
@@ -70,8 +71,22 @@ fn sort<N: IndexableNum>(
         return;
     }
 
-    let midpoint = (left + right) / 2;
-    let pivot = values[midpoint];
+    // apply median of three method
+    let start = values[left];
+    let mid = values[(left + right) >> 1];
+    let end = values[right];
+
+    let x = start.max(mid);
+    let pivot = if end > x {
+        x
+    } else if x == start {
+        mid.max(end)
+    } else if x == mid {
+        start.max(end)
+    } else {
+        end
+    };
+
     let mut i = left.wrapping_sub(1);
     let mut j = right.wrapping_add(1);
 
@@ -101,27 +116,6 @@ fn sort<N: IndexableNum>(
     sort(values, boxes, indices, j.wrapping_add(1), right, node_size);
 }
 
-/// Swap two values and two corresponding boxes.
-#[inline]
-fn swap<N: IndexableNum>(
-    values: &mut [u32],
-    boxes: &mut [N],
-    indices: &mut MutableIndices,
-    i: usize,
-    j: usize,
-) {
-    values.swap(i, j);
-
-    let k = 4 * i;
-    let m = 4 * j;
-    boxes.swap(k, m);
-    boxes.swap(k + 1, m + 1);
-    boxes.swap(k + 2, m + 2);
-    boxes.swap(k + 3, m + 3);
-
-    indices.swap(i, j);
-}
-
 // Taken from static_aabb2d_index under the mit/apache license
 // https://github.com/jbuckmccready/static_aabb2d_index/blob/9e6add59d77b74d4de0ac32159db47fbcb3acc28/src/static_aabb2d_index.rs#L486C1-L544C2
 #[inline]

src/rtree/index.rs

@@ -3,6 +3,7 @@
 mod hilbert;
 mod str;
 mod r#trait;
+mod util;
 
 pub use hilbert::HilbertSort;
 pub use r#str::STRSort;

src/rtree/sort/hilbert.rs

@@ -3,6 +3,7 @@ use rayon::iter::{IndexedParallelIterator, IntoParallelIterator, ParallelIterato
 
 use crate::indices::MutableIndices;
 use crate::r#type::IndexableNum;
+use crate::rtree::sort::util::swap;
 use crate::rtree::sort::{Sort, SortParams};
 
 /// An implementation of sort-tile-recursive (STR) sorting.
@@ -98,6 +99,17 @@ impl<N: IndexableNum> Sort<N> for STRSort {
     }
 }
 
+/// Max is only implemented for `Ord` types, but float types do not implement `Ord`.
+///
+/// So we use this as a hack to get the maximum of two PartialOrd values.
+fn partial_ord_max<N: IndexableNum>(a: N, b: N) -> N {
+    if a > b {
+        a
+    } else {
+        b
+    }
+}
+
 /// Custom quicksort that partially sorts bbox data alongside their sort values.
 // Partially taken from static_aabb2d_index under the MIT/Apache license
 fn sort<N: IndexableNum>(
@@ -114,8 +126,22 @@ fn sort<N: IndexableNum>(
         return;
     }
 
-    let midpoint = (left + right) / 2;
-    let pivot = values[midpoint];
+    // apply median of three method
+    let start = values[left];
+    let mid = values[(left + right) >> 1];
+    let end = values[right];
+
+    let x = partial_ord_max(start, mid);
+    let pivot = if end > x {
+        x
+    } else if x == start {
+        partial_ord_max(mid, end)
+    } else if x == mid {
+        partial_ord_max(start, end)
+    } else {
+        end
+    };
+
     let mut i = left.wrapping_sub(1);
     let mut j = right.wrapping_add(1);
 
@@ -144,24 +170,3 @@ fn sort<N: IndexableNum>(
     sort(values, boxes, indices, left, j, node_size);
     sort(values, boxes, indices, j.wrapping_add(1), right, node_size);
 }
-
-/// Swap two values and two corresponding boxes.
-#[inline]
-fn swap<N: IndexableNum>(
-    values: &mut [N],
-    boxes: &mut [N],
-    indices: &mut MutableIndices,
-    i: usize,
-    j: usize,
-) {
-    values.swap(i, j);
-
-    let k = 4 * i;
-    let m = 4 * j;
-    boxes.swap(k, m);
-    boxes.swap(k + 1, m + 1);
-    boxes.swap(k + 2, m + 2);
-    boxes.swap(k + 3, m + 3);
-
-    indices.swap(i, j);
-}

src/rtree/sort/mod.rs

@@ -0,0 +1,23 @@
+use crate::indices::MutableIndices;
+use crate::IndexableNum;
+
+/// Swap two values and two corresponding boxes.
+#[inline]
+pub(super) fn swap<V: IndexableNum, N: IndexableNum>(
+    values: &mut [V],
+    boxes: &mut [N],
+    indices: &mut MutableIndices,
+    i: usize,
+    j: usize,
+) {
+    values.swap(i, j);
+
+    let k = 4 * i;
+    let m = 4 * j;
+    boxes.swap(k, m);
+    boxes.swap(k + 1, m + 1);
+    boxes.swap(k + 2, m + 2);
+    boxes.swap(k + 3, m + 3);
+
+    indices.swap(i, j);
+}