geo_index/rtree/
builder.rs

1use bytemuck::cast_slice_mut;
2use geo_traits::{CoordTrait, RectTrait};
3
4use crate::indices::MutableIndices;
5use crate::r#type::IndexableNum;
6use crate::rtree::constants::VERSION;
7use crate::rtree::index::{RTree, RTreeMetadata};
8use crate::rtree::sort::{Sort, SortParams};
9
10/// The default node size used by [`RTreeBuilder::new`]
11pub const DEFAULT_RTREE_NODE_SIZE: u16 = 16;
12
13/// A builder to create an [`RTree`].
14///
15/// ```
16/// use geo_index::rtree::RTreeBuilder;
17/// use geo_index::rtree::sort::HilbertSort;
18///
19/// let mut builder = RTreeBuilder::<f64>::new(3);
20/// builder.add(0., 0., 2., 2.);
21/// builder.add(1., 1., 3., 3.);
22/// builder.add(2., 2., 4., 4.);
23/// let tree = builder.finish::<HilbertSort>();
24/// ```
25pub struct RTreeBuilder<N: IndexableNum> {
26    /// data buffer
27    data: Vec<u8>,
28    metadata: RTreeMetadata<N>,
29    pos: usize,
30    min_x: N,
31    min_y: N,
32    max_x: N,
33    max_y: N,
34}
35
36impl<N: IndexableNum> RTreeBuilder<N> {
37    /// Create a new builder with the provided number of items and the default node size.
38    pub fn new(num_items: u32) -> Self {
39        Self::new_with_node_size(num_items, DEFAULT_RTREE_NODE_SIZE)
40    }
41
42    /// Create a new builder with the provided number of items and node size.
43    pub fn new_with_node_size(num_items: u32, node_size: u16) -> Self {
44        let metadata = RTreeMetadata::new(num_items, node_size);
45        Self::from_metadata(metadata)
46    }
47
48    /// Create a new builder with the provided metadata
49    pub fn from_metadata(metadata: RTreeMetadata<N>) -> Self {
50        let mut data = vec![0; metadata.data_buffer_length()];
51
52        // Set data header
53        data[0] = 0xfb;
54        data[1] = (VERSION << 4) + N::TYPE_INDEX;
55        cast_slice_mut(&mut data[2..4])[0] = metadata.node_size();
56        cast_slice_mut(&mut data[4..8])[0] = metadata.num_items();
57
58        Self {
59            data,
60            metadata,
61            pos: 0,
62            min_x: N::max_value(),
63            min_y: N::max_value(),
64            max_x: N::min_value(),
65            max_y: N::min_value(),
66        }
67    }
68
69    /// Access the underlying [RTreeMetadata] of this instance.
70    pub fn metadata(&self) -> &RTreeMetadata<N> {
71        &self.metadata
72    }
73
74    /// Add a given rectangle to the RTree.
75    ///
76    /// This returns the insertion index, which provides a lookup back into the original data.
77    ///
78    /// `RTreeIndex::search` will return this same insertion index, which allows you to reference
79    /// your original collection.
80    #[inline]
81    pub fn add(&mut self, min_x: N, min_y: N, max_x: N, max_y: N) -> u32 {
82        let index = self.pos >> 2;
83        let (boxes, mut indices) = split_data_borrow(&mut self.data, &self.metadata);
84
85        indices.set(index, index);
86        boxes[self.pos] = min_x;
87        self.pos += 1;
88        boxes[self.pos] = min_y;
89        self.pos += 1;
90        boxes[self.pos] = max_x;
91        self.pos += 1;
92        boxes[self.pos] = max_y;
93        self.pos += 1;
94
95        if min_x < self.min_x {
96            self.min_x = min_x
97        };
98        if min_y < self.min_y {
99            self.min_y = min_y
100        };
101        if max_x > self.max_x {
102            self.max_x = max_x
103        };
104        if max_y > self.max_y {
105            self.max_y = max_y
106        };
107
108        index.try_into().unwrap()
109    }
110
111    /// Add a given rectangle to the RTree.
112    ///
113    /// This returns the insertion index, which provides a lookup back into the original data.
114    ///
115    /// `RTreeIndex::search` will return this same insertion index, which allows you to reference
116    /// your original collection.
117    #[inline]
118    pub fn add_rect(&mut self, rect: &impl RectTrait<T = N>) -> u32 {
119        self.add(
120            rect.min().x(),
121            rect.min().y(),
122            rect.max().x(),
123            rect.max().y(),
124        )
125    }
126
127    /// Consume this builder, perfoming the sort and generating an RTree ready for queries.
128    ///
129    /// [`HilbertSort`] and [`STRSort`] both implement [`Sort`], allowing you to choose the method
130    /// used.
131    ///
132    /// [`HilbertSort`]: crate::rtree::sort::HilbertSort
133    /// [`STRSort`]: crate::rtree::sort::STRSort
134    pub fn finish<S: Sort<N>>(mut self) -> RTree<N> {
135        assert_eq!(
136            self.pos >> 2,
137            self.metadata.num_items() as usize,
138            "Added {} items when expected {}.",
139            self.pos >> 2,
140            self.metadata.num_items()
141        );
142
143        let (boxes, mut indices) = split_data_borrow(&mut self.data, &self.metadata);
144
145        if self.metadata.num_items() <= 1 {
146            // Empty or only one item, we don't even have a root node to fill
147            return RTree {
148                buffer: self.data,
149                metadata: self.metadata,
150            };
151        }
152
153        if self.metadata.num_items() as usize <= self.metadata.node_size() as usize {
154            // only one node, skip sorting and just fill the root box
155            boxes[self.pos] = self.min_x;
156            self.pos += 1;
157            boxes[self.pos] = self.min_y;
158            self.pos += 1;
159            boxes[self.pos] = self.max_x;
160            self.pos += 1;
161            boxes[self.pos] = self.max_y;
162            self.pos += 1;
163
164            return RTree {
165                buffer: self.data,
166                metadata: self.metadata,
167            };
168        }
169
170        let mut sort_params = SortParams {
171            num_items: self.metadata.num_items() as usize,
172            node_size: self.metadata.node_size() as usize,
173            min_x: self.min_x,
174            min_y: self.min_y,
175            max_x: self.max_x,
176            max_y: self.max_y,
177        };
178        S::sort(&mut sort_params, boxes, &mut indices);
179
180        {
181            // generate nodes at each tree level, bottom-up
182            let mut pos = 0;
183            for end in self.metadata.level_bounds()[..self.metadata.level_bounds().len() - 1].iter()
184            {
185                while pos < *end {
186                    let node_index = pos;
187
188                    // calculate bbox for the new node
189                    let mut node_min_x = boxes[pos];
190                    pos += 1;
191                    let mut node_min_y = boxes[pos];
192                    pos += 1;
193                    let mut node_max_x = boxes[pos];
194                    pos += 1;
195                    let mut node_max_y = boxes[pos];
196                    pos += 1;
197                    for _ in 1..self.metadata.node_size() {
198                        if pos >= *end {
199                            break;
200                        }
201
202                        if boxes[pos] < node_min_x {
203                            node_min_x = boxes[pos];
204                        }
205                        pos += 1;
206                        if boxes[pos] < node_min_y {
207                            node_min_y = boxes[pos];
208                        }
209                        pos += 1;
210                        if boxes[pos] > node_max_x {
211                            node_max_x = boxes[pos]
212                        }
213                        pos += 1;
214                        if boxes[pos] > node_max_y {
215                            node_max_y = boxes[pos]
216                        }
217                        pos += 1;
218                    }
219
220                    // add the new node to the tree data
221                    indices.set(self.pos >> 2, node_index);
222                    boxes[self.pos] = node_min_x;
223                    self.pos += 1;
224                    boxes[self.pos] = node_min_y;
225                    self.pos += 1;
226                    boxes[self.pos] = node_max_x;
227                    self.pos += 1;
228                    boxes[self.pos] = node_max_y;
229                    self.pos += 1;
230                }
231            }
232        }
233
234        RTree {
235            buffer: self.data,
236            metadata: self.metadata,
237        }
238    }
239}
240
241/// Mutable borrow of boxes and indices
242fn split_data_borrow<'a, N: IndexableNum>(
243    data: &'a mut [u8],
244    metadata: &'a RTreeMetadata<N>,
245) -> (&'a mut [N], MutableIndices<'a>) {
246    let (boxes_buf, indices_buf) = data[8..].split_at_mut(metadata.nodes_byte_length);
247    debug_assert_eq!(indices_buf.len(), metadata.indices_byte_length);
248
249    let boxes = cast_slice_mut(boxes_buf);
250    let indices = MutableIndices::new(indices_buf, metadata.num_nodes());
251    (boxes, indices)
252}
253
254#[cfg(test)]
255mod test {
256    use crate::rtree::sort::{HilbertSort, STRSort};
257    use crate::rtree::RTreeIndex;
258
259    use super::*;
260
261    #[test]
262    fn does_not_panic_length_1_tree() {
263        let mut builder = RTreeBuilder::<f64>::new(1);
264        builder.add(-20., -20., 1020., 1020.);
265        let tree = builder.finish::<HilbertSort>();
266        let result = tree.search(0., 0., 0., 0.);
267        assert_eq!(result, vec![0]);
268    }
269
270    #[test]
271    fn build_str_index_with_various_items() {
272        let num_items_arr = [0, 1, 4, 8, 16, 20, 40, 80];
273        for num_items in num_items_arr {
274            build_index_with_various_items::<STRSort>(num_items);
275        }
276    }
277
278    #[test]
279    fn build_hilbert_index_with_various_items() {
280        let num_items_arr = [0, 1, 4, 8, 16, 20, 40, 80];
281        for num_items in num_items_arr {
282            build_index_with_various_items::<HilbertSort>(num_items);
283        }
284    }
285
286    fn build_index_with_various_items<S: Sort<f64>>(num_items: u32) {
287        let mut builder = RTreeBuilder::<f64>::new(num_items);
288        for i in 0..num_items {
289            builder.add(i as f64, i as f64, i as f64, i as f64);
290        }
291        let tree = builder.finish::<S>();
292        assert_eq!(tree.num_items(), num_items);
293        if num_items == 0 {
294            assert!(tree.search(0., 0., 0., 0.).is_empty());
295        } else {
296            for i in 0..num_items {
297                let results = tree.search(i as f64, i as f64, i as f64, i as f64);
298                assert_eq!(results, vec![i]);
299            }
300        }
301    }
302}
geo_index/rtree/builder.rs

geo_index/rtree/
builder.rs
