1use std::collections::HashMap;
2use std::f64;
3
4use rstar::{RTree, RTreeObject, AABB, SelectionFunction, Envelope, Point};
5
6use crate::point::{Vec2D, WorldUnit};
7
8use super::shape::{Shape, ShapeId, ShapeStored};
9use super::draw_commands::DrawCommand;
10
11impl Point for Vec2D<WorldUnit> {
12 type Scalar = WorldUnit;
13
14 const DIMENSIONS: usize = 2;
15
16 fn generate(mut generator: impl FnMut(usize) -> Self::Scalar) -> Self {
17 Vec2D {
18 x: generator(0),
19 y: generator(1)
20 }
21 }
22
23 fn nth(&self, index: usize) -> Self::Scalar {
24 match index {
25 0 => self.x,
26 1 => self.y,
27 _ => unreachable!()
28 }
29 }
30
31 fn nth_mut(&mut self, index: usize) -> &mut Self::Scalar {
32 match index {
33 0 => &mut self.x,
34 1 => &mut self.y,
35 _ => unreachable!()
36 }
37 }
38}
39
40impl RTreeObject for Shape {
41 type Envelope = AABB<Vec2D<WorldUnit>>;
42
43 fn envelope(&self) -> Self::Envelope {
44 let [p1, p2] = self.bbox();
45
46 AABB::from_corners(p1, p2)
47 }
48}
49
50struct BBoxIdSelection {
52 bbox: AABB<Vec2D<WorldUnit>>,
53 id: ShapeId,
54}
55
56impl BBoxIdSelection {
57 fn new(bbox: AABB<Vec2D<WorldUnit>>, id: ShapeId) -> BBoxIdSelection {
58 BBoxIdSelection {
59 bbox,
60 id,
61 }
62 }
63}
64
65impl SelectionFunction<Shape> for BBoxIdSelection {
66 fn should_unpack_parent(&self, envelope: &<Shape as RTreeObject>::Envelope) -> bool {
67 envelope.contains_envelope(&self.bbox)
68 }
69
70 fn should_unpack_leaf(&self, leaf: &Shape) -> bool {
71 leaf.id() == self.id
72 }
73}
74
75struct CenterAndRadiusSelection {
77 center: Vec2D<WorldUnit>,
78 radius: WorldUnit,
79}
80
81impl CenterAndRadiusSelection {
82 fn new(center: Vec2D<WorldUnit>, radius: WorldUnit) -> CenterAndRadiusSelection {
83 CenterAndRadiusSelection {
84 center, radius,
85 }
86 }
87}
88
89impl SelectionFunction<Shape> for CenterAndRadiusSelection {
90 fn should_unpack_parent(&self, envelope: &<Shape as RTreeObject>::Envelope) -> bool {
91 let bbox = AABB::from_corners(
92 self.center - Vec2D::new(self.radius, self.radius),
93 self.center + Vec2D::new(self.radius, self.radius),
94 );
95
96 envelope.intersects(&bbox)
97 }
98
99 fn should_unpack_leaf(&self, leaf: &Shape) -> bool {
100 leaf.intersects_circle(self.center, self.radius)
101 }
102}
103
104#[derive(Debug)]
106pub struct Storage {
107 ids: HashMap<ShapeId, AABB<Vec2D<WorldUnit>>>,
110
111 shapes: RTree<Shape>,
113
114 next_id: ShapeId,
116
117 next_z_index: usize,
119}
120
121impl Storage {
124 pub fn new() -> Storage {
125 Storage {
126 shapes: RTree::new(),
127 ids: HashMap::new(),
128 next_id: ShapeId::from(1),
129 next_z_index: 1,
130 }
131 }
132
133 fn nex_id(&mut self) -> ShapeId {
134 let id = self.next_id;
135 self.next_id = id.next();
136 id
137 }
138
139 fn next_z_index(&mut self) -> usize {
140 let index = self.next_z_index;
141 self.next_z_index = index + 1;
142 index
143 }
144
145 pub fn add(&mut self, shape: Box<dyn ShapeStored>) -> ShapeId {
147 let id = self.nex_id();
149 let z_index = self.next_z_index();
150 let shape = Shape::from_shape(shape, id, z_index);
151
152 self.ids.insert(id, shape.envelope());
153 self.shapes.insert(shape);
154
155 id
156 }
157
158 pub fn restore(&mut self, shape: Shape) -> ShapeId {
160 let shape_id = shape.id();
161
162 self.ids.insert(shape_id, shape.envelope());
163 self.shapes.insert(shape);
164
165 shape_id
166 }
167
168 pub fn remove(&mut self, id: ShapeId) -> Option<Shape> {
170 if let Some(bbox) = self.ids.get(&id) {
171 self.shapes.remove_with_selection_function(BBoxIdSelection::new(*bbox, id))
172 } else {
173 None
174 }
175 }
176
177 pub fn remove_circle(&mut self, center: Vec2D<WorldUnit>, radius: WorldUnit) -> impl Iterator<Item=Shape> + '_ {
180 self.shapes
181 .drain_with_selection_function(CenterAndRadiusSelection::new(center, radius))
182 }
183
184 pub fn shape_count(&self) -> usize {
186 self.shapes.size()
187 }
188
189 fn query_commands(&self, bbox: [Vec2D<WorldUnit>; 2]) -> Vec<DrawCommand> {
192 let bbox = AABB::from_corners(bbox[0], bbox[1]);
193 let mut commands: Vec<_> = self
194 .shapes
195 .locate_in_envelope_intersecting(&bbox)
196 .map(|shape| (shape.z_index(), shape.draw_commands()))
197 .collect();
198
199 commands.sort_unstable_by_key(|(i, _dc)| *i);
200
201 commands.into_iter().map(|(_i, dc)| dc).collect()
202 }
203
204 pub fn shapes_by_index(&self) -> impl Iterator<Item=&dyn ShapeStored> {
208 let mut shapes: Vec<_> = self
209 .shapes.iter()
210 .collect();
211
212 shapes.sort_unstable_by_key(|s| s.z_index());
213
214 shapes.into_iter().map(|s| s.shape_impl())
215 }
216
217 pub fn draw_commands(&self, bbox: [Vec2D<WorldUnit>; 2]) -> Vec<DrawCommand> {
221 self.query_commands(bbox)
222 }
223
224 pub fn get_bounds(&self) -> Option<[Vec2D<WorldUnit>; 2]> {
227 if self.shape_count() == 0 {
228 return None;
229 }
230
231 Some(self.shapes.iter().map(|shape| {
232 shape.envelope()
233 }).fold([Vec2D::new_world(f64::MAX, f64::MAX), Vec2D::new_world(f64::MIN, f64::MIN)], |acc, cur| {
234 let lower = cur.lower();
235 let upper = cur.upper();
236
237 [
238 Vec2D::new(
239 if lower.x < acc[0].x { lower.x } else { acc[0].x },
240 if lower.y < acc[0].y { lower.y } else { acc[0].y },
241 ),
242 Vec2D::new(
243 if upper.x > acc[1].x { upper.x } else { acc[1].x },
244 if upper.y > acc[1].y { upper.y } else { acc[1].y },
245 ),
246 ]
247 }))
248 }
249}
250
251impl Default for Storage {
252 fn default() -> Storage {
253 Storage::new()
254 }
255}
256
257pub struct ShapeIterator<'a> {
258 iterator: std::slice::Iter<'a, Shape>,
259}
260
261impl <'a> Iterator for ShapeIterator<'a> {
262 type Item = &'a Shape;
263
264 fn next(&mut self) -> Option<Self::Item> {
265 self.iterator.next()
266 }
267}
268
269#[cfg(test)]
270mod tests {
271 use pretty_assertions::assert_eq;
272
273 use crate::shape::{ShapeId, stored::{path::Path, ellipse::Ellipse}};
274 use crate::path_command::PathCommand;
275 use crate::color::Color;
276 use crate::draw_commands::DrawCommand;
277 use crate::geom::Angle;
278 use crate::style::{Style, Stroke};
279
280 use super::*;
281
282 #[test]
283 fn add_shapes_at_zoom() {
284 let mut storage = Storage::new();
285 let shapes: Vec<Box<dyn ShapeStored>> = vec![
286 Box::new(Ellipse::from_parts(Vec2D::new_world(0.0, 0.0), 1.0.into(), 1.0.into(), Angle::from_radians(0.0), Default::default())),
287 Box::new(Path::from_parts(vec![
288 PathCommand::MoveTo(Vec2D::new_world(0.0, 0.0)),
289 PathCommand::LineTo(Vec2D::new_world(1.0, 1.0)),
290 ], Default::default())),
291 ];
292 let ids: Vec<_> = shapes.into_iter().map(|shape| {
293 storage.add(shape)
294 }).collect();
295
296 assert_eq!(storage.shape_count(), 2);
297
298 assert_eq!(ids[0], ShapeId::from(1));
299 assert_eq!(ids[1], ShapeId::from(2));
300 }
301
302 #[test]
303 fn remove_line() {
304 let mut storage = Storage::new();
305 let line = Path::from_parts(vec![
306 PathCommand::MoveTo(Vec2D::new_world(0.0, -1.0)),
307 PathCommand::LineTo(Vec2D::new_world(0.0, 0.0)),
308 PathCommand::LineTo(Vec2D::new_world(1.0, 0.0)),
309 PathCommand::LineTo(Vec2D::new_world(1.0, 1.0)),
310 PathCommand::LineTo(Vec2D::new_world(0.0, 1.0)),
311 ], Style {
312 stroke: Some(Stroke {
313 size: 0.0.into(),
314 color: Default::default(),
315 }),
316 fill: None,
317 });
318
319 storage.add(Box::new(line));
320
321 assert_eq!(storage.remove_circle(Vec2D::new_world(0.5, -0.5), 0.5.into()).collect::<Vec<_>>().len(), 0);
322 assert_eq!(storage.remove_circle(Vec2D::new_world(0.5, 0.5), 0.5.into()).collect::<Vec<_>>().len(), 0);
323
324 assert!(storage.remove_circle(Vec2D::new_world(-0.25, -1.25), 0.5.into()).next().is_some());
325 }
326
327 #[test]
328 fn remove_by_id() {
329 let mut storage = Storage::new();
330
331 assert_eq!(storage.shape_count(), 0);
332
333 let obj_to_remove = Box::new(Path::from_parts(vec![
334 PathCommand::MoveTo(Vec2D::new_world(0.0, 0.0)),
335 PathCommand::LineTo(Vec2D::new_world(1.0, 0.0)),
336 ], Default::default()));
337 let obj_to_interfere = Box::new(Path::from_parts(vec![
338 PathCommand::MoveTo(Vec2D::new_world(0.0, 0.0)),
339 PathCommand::LineTo(Vec2D::new_world(1.0, 0.0)),
340 ], Default::default()));
341
342 let id_to_remove = storage.add(obj_to_remove);
343 let _id_to_interfere = storage.add(obj_to_interfere);
344
345 assert_eq!(storage.shape_count(), 2);
346
347 let data = storage.remove(id_to_remove).unwrap();
348
349 assert_eq!(storage.shape_count(), 1);
350
351 assert_eq!(data.id(), id_to_remove);
352 }
353
354 #[test]
355 fn remove_by_id_2() {
356 let mut storage = Storage::new();
357
358 assert_eq!(storage.shape_count(), 0);
359
360 let obj_to_remove = Box::new(Path::from_parts(vec![
361 PathCommand::MoveTo(Vec2D::new_world(0.0, 0.0)),
362 PathCommand::LineTo(Vec2D::new_world(1.0, 0.0)),
363 ], Default::default()));
364 let obj_to_interfere = Box::new(Path::from_parts(vec![
365 PathCommand::MoveTo(Vec2D::new_world(0.0, 0.0)),
366 PathCommand::LineTo(Vec2D::new_world(1.0, 0.0)),
367 ], Default::default()));
368
369 let _id_to_interfere = storage.add(obj_to_interfere);
370 let id_to_remove = storage.add(obj_to_remove);
371
372 assert_eq!(storage.shape_count(), 2);
373
374 let data = storage.remove(id_to_remove).unwrap();
375
376 assert_eq!(storage.shape_count(), 1);
377
378 assert_eq!(data.id(), id_to_remove);
379 }
380
381 #[test]
382 fn erase_invalidates_cache() {
383 let mut storage = Storage::new();
384 let bbox = [Vec2D::new_world(-40.0, -40.0), Vec2D::new_world(40.0, 40.0)];
385
386 storage.add(Box::new(Path::from_parts(vec![
387 PathCommand::MoveTo(Vec2D::new_world(0.0, 0.0)),
388 PathCommand::LineTo(Vec2D::new_world(1.0, 0.0)),
389 ], Default::default())));
390 let id = storage.add(Box::new(Path::from_parts(vec![
391 PathCommand::MoveTo(Vec2D::new_world(0.0, 1.0)),
392 PathCommand::LineTo(Vec2D::new_world(1.0, 1.0)),
393 ], Default::default())));
394
395 storage.remove(id);
396
397 assert_eq!(storage.draw_commands(bbox), vec![DrawCommand::Path {
398 style: Default::default(),
399 commands: vec![
400 PathCommand::MoveTo(Vec2D::new_world(0.0, 0.0)),
401 PathCommand::LineTo(Vec2D::new_world(1.0, 0.0)),
402 ],
403 }]);
404 }
405
406 #[test]
407 fn iter_by_bounds() {
408 let mut storage = Storage::new();
409
410 let shapes = vec![
411 Box::new(Path::from_parts(vec![
413 PathCommand::MoveTo(Vec2D::new_world(-1.1, -1.1)),
414 PathCommand::LineTo(Vec2D::new_world(-1.1, 1.1)),
415 PathCommand::LineTo(Vec2D::new_world(1.1, 1.1)),
416 PathCommand::LineTo(Vec2D::new_world(1.1, -1.1)),
417 ], Default::default())),
418 Box::new(Path::from_parts(vec![
419 PathCommand::MoveTo(Vec2D::new_world(-1.5, 0.5)),
420 PathCommand::LineTo(Vec2D::new_world(-1.5, 0.5)),
421 PathCommand::LineTo(Vec2D::new_world(-0.5, 1.5)),
422 PathCommand::LineTo(Vec2D::new_world(-0.5, 1.5)),
423 ], Default::default())),
424 Box::new(Path::from_parts(vec![
425 PathCommand::MoveTo(Vec2D::new_world(0.25, -0.5)),
426 PathCommand::LineTo(Vec2D::new_world(0.25, -0.5)),
427 PathCommand::LineTo(Vec2D::new_world(0.5, -0.25)),
428 PathCommand::LineTo(Vec2D::new_world(0.5, -0.25)),
429 ], Default::default())),
430
431 Box::new(Path::from_parts(vec![
433 PathCommand::MoveTo(Vec2D::new_world(20.0, 20.0)),
434 PathCommand::LineTo(Vec2D::new_world(20.0, 20.0)),
435 PathCommand::LineTo(Vec2D::new_world(20.0, 20.0)),
436 PathCommand::LineTo(Vec2D::new_world(20.0, 20.0)),
437 ], Style {
438 stroke: Some(Stroke {
439 color: Color::red(),
440 size: 1.0.into(),
441 }),
442 fill: None,
443 })),
444 ];
445
446 for shape in shapes.into_iter() {
447 storage.add(shape);
448 }
449
450 let buffer = storage.draw_commands([Vec2D::new_world(-1.0, -1.0), Vec2D::new_world(1.0, 1.0)]);
451
452 assert_eq!(buffer.len(), 3);
453
454 buffer.iter().for_each(|x| assert_eq!(x.color(), Some(Color::default())));
455 }
456
457 #[test]
458 fn get_bounds() {
459 let mut storage = Storage::new();
460
461 storage.add(Box::new(Path::from_parts(vec![
462 PathCommand::MoveTo(Vec2D::new_world(-5.0, -5.0)),
463 PathCommand::LineTo(Vec2D::new_world(1.0, 1.0)),
464 PathCommand::LineTo(Vec2D::new_world(1.0, 1.0)),
465 PathCommand::LineTo(Vec2D::new_world(1.0, 1.0)),
466 ], Default::default())));
467 storage.add(Box::new(Path::from_parts(vec![
468 PathCommand::MoveTo(Vec2D::new_world(5.0, 5.0)),
469 PathCommand::LineTo(Vec2D::new_world(1.0, 1.0)),
470 PathCommand::LineTo(Vec2D::new_world(1.0, 1.0)),
471 PathCommand::LineTo(Vec2D::new_world(1.0, 1.0)),
472 ], Default::default())));
473
474 assert_eq!(storage.get_bounds(), Some([Vec2D::new_world(-7.0, -7.0), Vec2D::new_world(7.0, 7.0)]));
475 }
476
477 #[test]
478 fn draw_oder_is_respected() {
479 let mut storage = Storage::new();
480
481 storage.add(Box::new(Path::from_parts(vec![
482 PathCommand::MoveTo(Vec2D::new_world(0.0, 0.0)),
483 PathCommand::LineTo(Vec2D::new_world(0.0, 2.0)),
484 ], Style::red_line())));
485 storage.add(Box::new(Path::from_parts(vec![
486 PathCommand::MoveTo(Vec2D::new_world(0.0, 0.0)),
487 PathCommand::LineTo(Vec2D::new_world(0.0, 1.0)),
488 ], Default::default())));
489
490 let commands = storage.draw_commands([Vec2D::new_world(-1.0, -1.0), Vec2D::new_world(1.0, 3.0)]);
491
492 assert_eq!(commands.len(), 2);
493
494 assert_eq!(commands[0].color(), Some(Color::red()));
495 assert_eq!(commands[1].color(), Some(Default::default()));
496 }
497}