Add geo-traits integration (#71)

Cargo.toml

@@ -22,6 +22,7 @@ categories = ["data-structures", "algorithms", "science::geo"]
 [dependencies]
 bytemuck = "1"
 float_next_after = "1"
+geo-traits = "0.2"
 num-traits = "0.2"
 rayon = { version = "1.8.0", optional = true }
 thiserror = "1"

src/kdtree/builder.rs

@@ -2,11 +2,14 @@ use std::cmp;
 use std::marker::PhantomData;
 
 use bytemuck::cast_slice_mut;
+use geo_traits::{CoordTrait, PointTrait};
 
+use crate::error::Result;
 use crate::indices::MutableIndices;
 use crate::kdtree::constants::{KDBUSH_HEADER_SIZE, KDBUSH_MAGIC, KDBUSH_VERSION};
 use crate::kdtree::OwnedKDTree;
 use crate::r#type::IndexableNum;
+use crate::GeoIndexError;
 
 const DEFAULT_NODE_SIZE: u16 = 64;
 
@@ -68,7 +71,10 @@ impl<N: IndexableNum> KDTreeBuilder<N> {
         }
     }
 
-    /// Add a point to the index.
+    /// Add a point to the KDTree.
+    ///
+    /// This returns a positional index that provides a lookup back into the original data.
+    #[inline]
     pub fn add(&mut self, x: N, y: N) -> usize {
         let index = self.pos >> 1;
         let (coords, mut ids) = split_data_borrow(
@@ -88,6 +94,29 @@ impl<N: IndexableNum> KDTreeBuilder<N> {
         index
     }
 
+    /// Add a coord to the KDTree.
+    ///
+    /// This returns a positional index that provides a lookup back into the original data.
+    #[inline]
+    pub fn add_coord(&mut self, coord: &impl CoordTrait<T = N>) -> usize {
+        self.add(coord.x(), coord.y())
+    }
+
+    /// Add a point to the KDTree.
+    ///
+    /// This returns a positional index that provides a lookup back into the original data.
+    ///
+    /// ## Errors
+    ///
+    /// - If the point is empty.
+    #[inline]
+    pub fn add_point(&mut self, point: &impl PointTrait<T = N>) -> Result<usize> {
+        let coord = point.coord().ok_or(GeoIndexError::General(
+            "Unable to add empty point to KDTree".to_string(),
+        ))?;
+        Ok(self.add_coord(&coord))
+    }
+
     /// Consume this builder, perfoming the k-d sort and generating a KDTree ready for queries.
     pub fn finish(mut self) -> OwnedKDTree<N> {
         assert_eq!(

src/kdtree/trait.rs

@@ -1,5 +1,6 @@
 use std::borrow::Cow;
 
+use geo_traits::{CoordTrait, RectTrait};
 use tinyvec::TinyVec;
 
 use crate::indices::Indices;
@@ -80,6 +81,18 @@ pub trait KDTreeIndex<N: IndexableNum> {
         result
     }
 
+    /// Search the index for items within a given bounding box.
+    ///
+    /// Returns indices of found items
+    fn range_rect(&self, rect: &impl RectTrait<T = N>) -> Vec<usize> {
+        self.range(
+            rect.min().x(),
+            rect.min().y(),
+            rect.max().x(),
+            rect.max().y(),
+        )
+    }
+
     /// Search the index for items within a given radius.
     ///
     /// - qx: x value of query point
@@ -144,6 +157,16 @@ pub trait KDTreeIndex<N: IndexableNum> {
         }
         result
     }
+
+    /// Search the index for items within a given radius.
+    ///
+    /// - coord: coordinate of query point
+    /// - r: radius
+    ///
+    /// Returns indices of found items
+    fn within_coord(&self, coord: &impl CoordTrait<T = N>, r: N) -> Vec<usize> {
+        self.within(coord.x(), coord.y(), r)
+    }
 }
 
 impl<N: IndexableNum> KDTreeIndex<N> for KDTreeRef<'_, N> {

src/rtree/builder.rs

@@ -1,4 +1,5 @@
 use bytemuck::cast_slice_mut;
+use geo_traits::{CoordTrait, RectTrait};
 
 use crate::indices::MutableIndices;
 use crate::r#type::IndexableNum;
@@ -73,7 +74,9 @@ impl<N: IndexableNum> RTreeBuilder<N> {
         }
     }
 
-    /// Add a given rectangle to the index.
+    /// Add a given rectangle to the RTree.
+    ///
+    /// This returns a positional index that provides a lookup back into the original data.
     #[inline]
     pub fn add(&mut self, min_x: N, min_y: N, max_x: N, max_y: N) -> usize {
         let index = self.pos >> 2;
@@ -110,6 +113,19 @@ impl<N: IndexableNum> RTreeBuilder<N> {
         index
     }
 
+    /// Add a given rectangle to the RTree.
+    ///
+    /// This returns a positional index that provides a lookup back into the original data.
+    #[inline]
+    pub fn add_rect(&mut self, rect: &impl RectTrait<T = N>) -> usize {
+        self.add(
+            rect.min().x(),
+            rect.min().y(),
+            rect.max().x(),
+            rect.max().y(),
+        )
+    }
+
     /// Consume this builder, perfoming the sort and generating an RTree ready for queries.
     pub fn finish<S: Sort<N>>(mut self) -> OwnedRTree<N> {
         assert_eq!(

src/rtree/index.rs

@@ -23,7 +23,7 @@ pub(crate) struct TreeMetadata<N: IndexableNum> {
 }
 
 impl<N: IndexableNum> TreeMetadata<N> {
-    pub fn try_new(data: &[u8]) -> Result<Self> {
+    fn try_new(data: &[u8]) -> Result<Self> {
         let magic = data[0];
         if magic != 0xfb {
             return Err(GeoIndexError::General(

src/rtree/trait.rs

@@ -1,3 +1,5 @@
+use geo_traits::{CoordTrait, RectTrait};
+
 use crate::error::Result;
 use crate::indices::Indices;
 use crate::r#type::IndexableNum;
@@ -97,6 +99,18 @@ pub trait RTreeIndex<N: IndexableNum>: Sized {
         results
     }
 
+    /// Search an RTree given the provided bounding box.
+    ///
+    /// Results are the indexes of the inserted objects in insertion order.
+    fn search_rect(&self, rect: &impl RectTrait<T = N>) -> Vec<usize> {
+        self.search(
+            rect.min().x(),
+            rect.min().y(),
+            rect.max().x(),
+            rect.max().y(),
+        )
+    }
+
     // #[allow(unused_mut, unused_labels, unused_variables)]
     // fn neighbors(&self, x: N, y: N, max_distance: Option<N>) -> Vec<usize> {
     //     let boxes = self.boxes();