geos/
spatial_index.rs

1use std::ffi::c_void;
2use std::marker::PhantomData;
3
4use geos_sys::*;
5
6use crate::context_handle::{with_context, PtrWrap};
7use crate::Geom;
8use crate::{AsRaw, AsRawMut, GResult};
9
10pub trait SpatialIndex<I> {
11    fn insert<G: Geom>(&mut self, geometry: &G, item: I);
12
13    fn query<G: Geom, V: FnMut(&I)>(&self, geometry: &G, visitor: V);
14}
15
16pub struct STRtree<I> {
17    pub(crate) ptr: PtrWrap<*mut GEOSSTRtree>,
18    item_type: PhantomData<I>,
19}
20
21impl<I> STRtree<I> {
22    pub fn with_capacity(node_capacity: usize) -> GResult<STRtree<I>> {
23        with_context(|ctx| unsafe {
24            let ptr = GEOSSTRtree_create_r(ctx.as_raw(), node_capacity);
25            Ok(STRtree {
26                ptr: PtrWrap(ptr),
27                item_type: PhantomData,
28            })
29        })
30    }
31
32    pub fn iterate<V>(&self, visitor: V)
33    where
34        V: FnMut(&I),
35    {
36        with_context(|ctx| unsafe {
37            let (closure, callback) = unpack_closure(&visitor);
38            GEOSSTRtree_iterate_r(ctx.as_raw(), *self.ptr, Some(callback), closure);
39        })
40    }
41}
42
43impl<I> SpatialIndex<I> for STRtree<I> {
44    fn insert<G: Geom>(&mut self, geometry: &G, item: I) {
45        with_context(|ctx| unsafe {
46            GEOSSTRtree_insert_r(
47                ctx.as_raw(),
48                *self.ptr,
49                geometry.as_raw(),
50                Box::into_raw(Box::new(item)) as *mut c_void,
51            );
52        })
53    }
54
55    fn query<'b, G: Geom, V: FnMut(&I)>(&self, geometry: &G, visitor: V) {
56        with_context(|ctx| unsafe {
57            let (closure, callback) = unpack_closure(&visitor);
58            GEOSSTRtree_query_r(
59                ctx.as_raw(),
60                *self.ptr,
61                geometry.as_raw(),
62                Some(callback),
63                closure,
64            );
65        })
66    }
67}
68
69impl<I> AsRaw for STRtree<I> {
70    type RawType = GEOSSTRtree;
71
72    fn as_raw(&self) -> *const Self::RawType {
73        *self.ptr
74    }
75}
76
77impl<I> AsRawMut for STRtree<I> {
78    type RawType = GEOSSTRtree;
79
80    unsafe fn as_raw_mut_override(&self) -> *mut Self::RawType {
81        *self.ptr
82    }
83}
84
85impl<I> Drop for STRtree<I> {
86    fn drop(&mut self) {
87        unsafe extern "C" fn callback<I>(item: *mut c_void, _data: *mut c_void) {
88            drop(Box::from_raw(item as *mut I));
89        }
90
91        with_context(|ctx| unsafe {
92            GEOSSTRtree_iterate_r(
93                ctx.as_raw(),
94                *self.ptr,
95                Some(callback::<I>),
96                std::ptr::null_mut(),
97            );
98            GEOSSTRtree_destroy_r(ctx.as_raw(), *self.ptr);
99        })
100    }
101}
102
103unsafe fn unpack_closure<F, I>(
104    closure: &F,
105) -> (*mut c_void, extern "C" fn(*mut c_void, *mut c_void))
106where
107    F: FnMut(&I),
108{
109    extern "C" fn trampoline<F, I>(item: *mut c_void, data: *mut c_void)
110    where
111        F: FnMut(&I),
112    {
113        unsafe {
114            let closure: &mut F = &mut *(data as *mut F);
115            (*closure)(&mut *(item as *mut I));
116        }
117    }
118
119    (closure as *const F as *mut c_void, trampoline::<F, I>)
120}
121
122#[cfg(test)]
123mod test {
124    use std::collections::HashSet;
125
126    use crate::{Geometry, STRtree, SpatialIndex};
127
128    #[test]
129    fn test_strtree() {
130        let mut tree = STRtree::<&str>::with_capacity(10).unwrap();
131
132        let point = Geometry::new_from_wkt("POINT(5 5)").unwrap();
133        let line = Geometry::new_from_wkt("LINESTRING (0 0, 10 0)").unwrap();
134        let polygon = Geometry::new_from_wkt("POLYGON((2 2, 8 2, 8 8, 2 8, 2 2))").unwrap();
135
136        tree.insert(&point, "Point");
137        tree.insert(&line, "Line");
138        tree.insert(&polygon, "Polygon");
139
140        // Test iterate
141
142        let mut items = HashSet::<&str>::new();
143        tree.iterate(|item| {
144            items.insert(*item);
145        });
146
147        assert_eq!(
148            items,
149            vec!["Line", "Point", "Polygon"].into_iter().collect()
150        );
151
152        // Test query
153
154        items.clear();
155        tree.query(&point, |item| {
156            items.insert(*item);
157        });
158
159        assert_eq!(items, vec!["Point", "Polygon"].into_iter().collect());
160    }
161}
geos/spatial_index.rs

geos/
spatial_index.rs
