Skip to main content

pizarra/
storage.rs

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
50/// Handles delete from the storage by Id
51struct 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
75/// Handles delete from the storage by center and radius
76struct 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/// This struct handles storage and fast retrieval of shapes.
105#[derive(Debug)]
106pub struct Storage {
107    /// used to easily find a shape that will be deleted by id, because the
108    /// rstar tree knows how to find bboxes easily but not ids.
109    ids: HashMap<ShapeId, AABB<Vec2D<WorldUnit>>>,
110
111    /// The rstart tree that backs this storage.
112    shapes: RTree<Shape>,
113
114    /// The next id available.
115    next_id: ShapeId,
116
117    /// the next z-index available.
118    next_z_index: usize,
119}
120
121/// A storage struct that organizes shapes by their zoom level and allows for
122/// fast queries given a zoom and a bbox.
123impl 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    /// Adds a new shape to the storage
146    pub fn add(&mut self, shape: Box<dyn ShapeStored>) -> ShapeId {
147        // The previous shape has been finished, move it to persistent storage
148        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    /// Restores a deleted shape
159    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    /// Deletes a shape from this storage given its id if such id exists
169    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    /// Returns the id of a shape intersecting the circle described by center
178    /// and radius, if any.
179    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    /// Returns the number of shapes in this storage
185    pub fn shape_count(&self) -> usize {
186        self.shapes.size()
187    }
188
189    /// returns the draw commands necessary to paint the portion of the screen
190    /// indicated by bbox.
191    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    /// Returns the shapes stored in this storage ordered by their index. It is
205    /// used for rendering the shapes in the order they where drawn. Creates an
206    /// allocation.
207    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    /// gets all the draw commands necessary to paint the portion of the screen
218    /// indicated by bbox. Commands are cached to improve the performance of
219    /// continuously drawing on the same portion of the screen
220    pub fn draw_commands(&self, bbox: [Vec2D<WorldUnit>; 2]) -> Vec<DrawCommand> {
221        self.query_commands(bbox)
222    }
223
224    /// Gets the bounding box that surrounds every shape drawn or returns None
225    /// if there are no shapes
226    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            // in green the ones that go
412            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            // in red the ones that dont
432            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}