1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
use std::ffi::c_void;
use std::marker::PhantomData;
use std::sync::Arc;

use geos_sys::*;

use context_handle::PtrWrap;
use ContextHandling;
use {AsRaw, AsRawMut, GResult};
use {ContextHandle, Geom};

pub trait SpatialIndex<'a, I> {
    fn insert<'b, G: Geom<'b>>(&mut self, geometry: &G, item: I);

    fn query<'b, G: Geom<'b>, V: FnMut(&I)>(&self, geometry: &G, visitor: V);
}

pub struct STRtree<'a, I> {
    pub(crate) ptr: PtrWrap<*mut GEOSSTRtree>,
    context: Arc<ContextHandle<'a>>,
    item_type: PhantomData<I>,
}

impl<'a, I> STRtree<'a, I> {
    pub fn with_capacity(node_capacity: usize) -> GResult<STRtree<'a, I>> {
        match ContextHandle::init_e(Some("STRtree::with_capacity")) {
            Ok(context_handle) => unsafe {
                let ptr = GEOSSTRtree_create_r(context_handle.as_raw(), node_capacity);
                Ok(STRtree {
                    ptr: PtrWrap(ptr),
                    context: Arc::new(context_handle),
                    item_type: PhantomData,
                })
            },
            Err(e) => Err(e),
        }
    }

    pub fn iterate<V>(&self, visitor: V)
    where
        V: FnMut(&I),
    {
        unsafe {
            let (closure, callback) = unpack_closure(&visitor);
            GEOSSTRtree_iterate_r(self.get_raw_context(), *self.ptr, Some(callback), closure);
        }
    }
}

impl<'a, I> SpatialIndex<'a, I> for STRtree<'a, I> {
    fn insert<'b, G: Geom<'b>>(&mut self, geometry: &G, item: I) {
        unsafe {
            GEOSSTRtree_insert_r(
                self.get_raw_context(),
                *self.ptr,
                geometry.as_raw(),
                Box::into_raw(Box::new(item)) as *mut c_void,
            );
        }
    }

    fn query<'b, G: Geom<'b>, V: FnMut(&I)>(&self, geometry: &G, visitor: V) {
        unsafe {
            let (closure, callback) = unpack_closure(&visitor);
            GEOSSTRtree_query_r(
                self.get_raw_context(),
                *self.ptr,
                geometry.as_raw(),
                Some(callback),
                closure,
            );
        }
    }
}

impl<'a, I> AsRaw for STRtree<'a, I> {
    type RawType = GEOSSTRtree;

    fn as_raw(&self) -> *const Self::RawType {
        *self.ptr
    }
}

impl<'a, I> AsRawMut for STRtree<'a, I> {
    type RawType = GEOSSTRtree;

    unsafe fn as_raw_mut_override(&self) -> *mut Self::RawType {
        *self.ptr
    }
}

impl<'a, I> ContextHandling for STRtree<'a, I> {
    type Context = Arc<ContextHandle<'a>>;

    fn get_raw_context(&self) -> GEOSContextHandle_t {
        self.context.as_raw()
    }

    fn clone_context(&self) -> Arc<ContextHandle<'a>> {
        Arc::clone(&self.context)
    }
}

impl<I> Drop for STRtree<'_, I> {
    fn drop(&mut self) {
        unsafe extern "C" fn callback<I>(item: *mut c_void, _data: *mut c_void) {
            drop(Box::from_raw(item as *mut I));
        }

        unsafe {
            GEOSSTRtree_iterate_r(
                self.get_raw_context(),
                *self.ptr,
                Some(callback::<I>),
                std::ptr::null_mut(),
            );
            GEOSSTRtree_destroy_r(self.get_raw_context(), *self.ptr);
        }
    }
}

unsafe fn unpack_closure<F, I>(
    closure: &F,
) -> (*mut c_void, extern "C" fn(*mut c_void, *mut c_void))
where
    F: FnMut(&I),
{
    extern "C" fn trampoline<F, I>(item: *mut c_void, data: *mut c_void)
    where
        F: FnMut(&I),
    {
        unsafe {
            let closure: &mut F = &mut *(data as *mut F);
            (*closure)(&mut *(item as *mut I));
        }
    }

    (closure as *const F as *mut c_void, trampoline::<F, I>)
}

#[cfg(test)]
mod test {
    use std::collections::HashSet;

    use {Geometry, STRtree, SpatialIndex};

    #[test]
    fn test_strtree() {
        let mut tree = STRtree::<&str>::with_capacity(10).unwrap();

        let point = Geometry::new_from_wkt("POINT(5 5)").unwrap();
        let line = Geometry::new_from_wkt("LINESTRING (0 0, 10 0)").unwrap();
        let polygon = Geometry::new_from_wkt("POLYGON((2 2, 8 2, 8 8, 2 8, 2 2))").unwrap();

        tree.insert(&point, "Point");
        tree.insert(&line, "Line");
        tree.insert(&polygon, "Polygon");

        // Test iterate

        let mut items = HashSet::<&str>::new();
        tree.iterate(|item| {
            items.insert(*item);
        });

        assert_eq!(
            items,
            vec!["Line", "Point", "Polygon"].into_iter().collect()
        );

        // Test query

        items.clear();
        tree.query(&point, |item| {
            items.insert(*item);
        });

        assert_eq!(items, vec!["Point", "Polygon"].into_iter().collect());
    }
}