energyml_section/spatial.rs
1//! A static R-tree over 2D boxes, with an overflow list for what is added after
2//! the build (ADR-0010).
3//!
4//! **Written to move as is.** This module imports nothing from the rest of the
5//! crate (nor any dependency): the design lets the heads and the spatial layer
6//! be born in `energyml-section` and move down to the future
7//! `energyml-cache` crate by copying this one file.
8//!
9//! # What it is
10//!
11//! - **Static, packed by STR** (sort-tile-recursive, Leutenegger et al. 1997):
12//! the boxes are sorted by the x of their centre, cut into vertical slices,
13//! each slice sorted by y and cut into nodes of [`FAN_OUT`] boxes; the nodes
14//! are packed the same way, level by level, up to one root. No box is ever
15//! moved after the build, so there is no balancing to get wrong.
16//! - **An overflow list**: [`Index::insert`] after the build appends to a flat
17//! list scanned on every query; once it outgrows a share of the tree
18//! ([`Index::needs_rebuild`]) the next insert folds it in and rebuilds.
19//! - **Exact by construction**: a query returns every item whose box meets the
20//! query box (closed boxes: touching counts), and **only** those - the same
21//! answer as a linear scan, in another order. The tests confront the two on
22//! random data. What the tree saves is the boxes it never opens, not any
23//! precision; a query by a polygon is its bounding box ([`Rect::of_points`]),
24//! then the caller's exact test on what comes back.
25//!
26//! A box with a non-finite bound is not indexable: [`Rect::new`] refuses it
27//! (`None`), so the caller decides what such an item means rather than the tree
28//! losing it silently.
29
30/// Boxes per node. Eight keeps a 5 000-item tree four levels deep.
31pub const FAN_OUT: usize = 8;
32
33/// The overflow is folded into the tree when it holds more than this many
34/// items **and** more than a quarter of the tree.
35pub const OVERFLOW_FLOOR: usize = 64;
36
37/// A closed 2D box: `min <= max` on both axes, both finite.
38#[derive(Clone, Copy, Debug, PartialEq)]
39pub struct Rect {
40 /// The lower corner, `[x, y]`.
41 pub min: [f64; 2],
42 /// The upper corner, `[x, y]`.
43 pub max: [f64; 2],
44}
45
46impl Rect {
47 /// The box of two corners, in any order; `None` if a coordinate is not finite.
48 pub fn new(a: [f64; 2], b: [f64; 2]) -> Option<Self> {
49 if !a.iter().chain(b.iter()).all(|v| v.is_finite()) {
50 return None;
51 }
52 Some(Self {
53 min: [a[0].min(b[0]), a[1].min(b[1])],
54 max: [a[0].max(b[0]), a[1].max(b[1])],
55 })
56 }
57
58 /// A point, as a degenerate box; `None` if not finite.
59 pub fn point(p: [f64; 2]) -> Option<Self> {
60 Self::new(p, p)
61 }
62
63 /// The bounding box of some points (a polygon's, for a query by polygon);
64 /// `None` if there are none or one is not finite.
65 pub fn of_points(points: &[[f64; 2]]) -> Option<Self> {
66 let first = *points.first()?;
67 let mut rect = Self::point(first)?;
68 for p in &points[1..] {
69 rect = rect.union(&Self::point(*p)?);
70 }
71 Some(rect)
72 }
73
74 /// This box grown by `margin` on every side (a negative margin shrinks it,
75 /// and an emptied box stays a box at its centre).
76 pub fn grown(&self, margin: f64) -> Self {
77 let mut r = *self;
78 for axis in 0..2 {
79 r.min[axis] -= margin;
80 r.max[axis] += margin;
81 if r.min[axis] > r.max[axis] {
82 let mid = (self.min[axis] + self.max[axis]) / 2.0;
83 r.min[axis] = mid;
84 r.max[axis] = mid;
85 }
86 }
87 r
88 }
89
90 /// The smallest box holding both.
91 pub fn union(&self, other: &Self) -> Self {
92 Self {
93 min: [self.min[0].min(other.min[0]), self.min[1].min(other.min[1])],
94 max: [self.max[0].max(other.max[0]), self.max[1].max(other.max[1])],
95 }
96 }
97
98 /// The two boxes share a point (touching counts).
99 pub fn intersects(&self, other: &Self) -> bool {
100 self.min[0] <= other.max[0]
101 && other.min[0] <= self.max[0]
102 && self.min[1] <= other.max[1]
103 && other.min[1] <= self.max[1]
104 }
105
106 fn centre(&self, axis: usize) -> f64 {
107 (self.min[axis] + self.max[axis]) / 2.0
108 }
109}
110
111#[derive(Clone, Debug)]
112struct Node {
113 rect: Rect,
114 /// Indices into the level below: entries for a leaf, nodes otherwise.
115 kids: Vec<usize>,
116 leaf: bool,
117}
118
119/// The index: items of type `T` filed under their box.
120#[derive(Clone, Debug)]
121pub struct Index<T> {
122 /// Every item, built and overflowed alike; `entries[..built]` are in the tree.
123 entries: Vec<(Rect, T)>,
124 built: usize,
125 nodes: Vec<Node>,
126 root: Option<usize>,
127 rebuilds: usize,
128}
129
130impl<T> Default for Index<T> {
131 fn default() -> Self {
132 Self {
133 entries: Vec::new(),
134 built: 0,
135 nodes: Vec::new(),
136 root: None,
137 rebuilds: 0,
138 }
139 }
140}
141
142impl<T> Index<T> {
143 /// An empty index.
144 pub fn new() -> Self {
145 Self::default()
146 }
147
148 /// An index built at once over `items`.
149 pub fn build(items: Vec<(Rect, T)>) -> Self {
150 let mut index = Self {
151 entries: items,
152 ..Self::default()
153 };
154 index.rebuild();
155 index
156 }
157
158 /// How many items it holds, in the tree and in the overflow.
159 pub fn len(&self) -> usize {
160 self.entries.len()
161 }
162
163 /// No item at all.
164 pub fn is_empty(&self) -> bool {
165 self.entries.is_empty()
166 }
167
168 /// How many items wait in the overflow list.
169 pub fn overflow(&self) -> usize {
170 self.entries.len() - self.built
171 }
172
173 /// How many times the tree was packed (the build counts).
174 pub fn rebuilds(&self) -> usize {
175 self.rebuilds
176 }
177
178 /// The overflow has outgrown its share: more than [`OVERFLOW_FLOOR`] items
179 /// and more than a quarter of the tree.
180 pub fn needs_rebuild(&self) -> bool {
181 self.overflow() > OVERFLOW_FLOOR && self.overflow() * 4 > self.built
182 }
183
184 /// Add an item. It is found by the very next query; the tree is packed
185 /// again when the overflow has outgrown its share ([`Self::needs_rebuild`]).
186 pub fn insert(&mut self, rect: Rect, item: T) {
187 self.entries.push((rect, item));
188 if self.needs_rebuild() {
189 self.rebuild();
190 }
191 }
192
193 /// Pack every item into a fresh tree and empty the overflow.
194 pub fn rebuild(&mut self) {
195 self.rebuilds += 1;
196 self.built = self.entries.len();
197 self.nodes.clear();
198 self.root = None;
199 if self.entries.is_empty() {
200 return;
201 }
202 // Level 0: the entries, tiled.
203 let rects: Vec<Rect> = self.entries.iter().map(|(r, _)| *r).collect();
204 let mut level: Vec<usize> = (0..rects.len()).collect();
205 let mut leaf = true;
206 let mut rects_of_level = rects;
207 loop {
208 let order = tile(&rects_of_level, FAN_OUT);
209 let mut next_nodes: Vec<usize> = Vec::new();
210 let mut next_rects: Vec<Rect> = Vec::new();
211 for group in order.chunks(FAN_OUT) {
212 let kids: Vec<usize> = group.iter().map(|&i| level[i]).collect();
213 let rect = group
214 .iter()
215 .map(|&i| rects_of_level[i])
216 .reduce(|a, b| a.union(&b))
217 .expect("non-empty chunk");
218 next_nodes.push(self.nodes.len());
219 next_rects.push(rect);
220 self.nodes.push(Node { rect, kids, leaf });
221 }
222 if next_nodes.len() == 1 {
223 self.root = Some(next_nodes[0]);
224 return;
225 }
226 level = next_nodes;
227 rects_of_level = next_rects;
228 leaf = false;
229 }
230 }
231
232 /// Every item whose box meets `window`, the tree's and the overflow's - the
233 /// set a linear scan gives, in no particular order.
234 pub fn query(&self, window: &Rect) -> Vec<&T> {
235 let mut out = Vec::new();
236 self.visit(window, &mut |item| out.push(item));
237 out
238 }
239
240 /// [`Self::query`], calling `found` for each instead of collecting.
241 pub fn visit<'a>(&'a self, window: &Rect, found: &mut dyn FnMut(&'a T)) {
242 if let Some(root) = self.root {
243 let mut stack = vec![root];
244 while let Some(at) = stack.pop() {
245 let node = &self.nodes[at];
246 if !node.rect.intersects(window) {
247 continue;
248 }
249 if node.leaf {
250 for &k in &node.kids {
251 let (rect, item) = &self.entries[k];
252 if rect.intersects(window) {
253 found(item);
254 }
255 }
256 } else {
257 stack.extend(node.kids.iter().copied());
258 }
259 }
260 }
261 for (rect, item) in &self.entries[self.built..] {
262 if rect.intersects(window) {
263 found(item);
264 }
265 }
266 }
267}
268
269/// The STR order of some boxes: by the x of the centre, cut into vertical
270/// slices of `ceil(sqrt(n / fan)) * fan` boxes, each slice by the y of the
271/// centre. Returns indices into `rects`.
272fn tile(rects: &[Rect], fan: usize) -> Vec<usize> {
273 let n = rects.len();
274 let mut order: Vec<usize> = (0..n).collect();
275 order.sort_by(|&a, &b| rects[a].centre(0).total_cmp(&rects[b].centre(0)));
276 let nodes = n.div_ceil(fan);
277 let slices = (nodes as f64).sqrt().ceil() as usize;
278 let per_slice = slices.max(1) * fan;
279 for slice in order.chunks_mut(per_slice) {
280 slice.sort_by(|&a, &b| rects[a].centre(1).total_cmp(&rects[b].centre(1)));
281 }
282 order
283}