spatial_table/
lib.rs

1use entity_table::ComponentTable;
2#[cfg(feature = "serialize")]
3use entity_table::ComponentTableEntries;
4pub use entity_table::Entity; // public so it can be referenced in macro body
5use grid_2d::Grid;
6pub use grid_2d::{ICoord, UCoord};
7#[cfg(feature = "serialize")]
8pub use serde; // public so it can be referenced in macro body
9#[cfg(feature = "serialize")]
10use serde::{Deserialize, Serialize};
11
12pub trait Layers: Default {
13    type Layer: Copy + PartialEq + Eq;
14    fn select_field_mut(&mut self, layer: Self::Layer) -> &mut Option<Entity>;
15}
16
17#[cfg(not(feature = "serialize"))]
18#[macro_export]
19macro_rules! declare_layers_module {
20    { $module_name:ident { $($field_name:ident: $variant_name:ident,)* } } => {
21        mod $module_name {
22            #[derive(Debug, Clone, Copy, PartialEq, Eq)]
23            pub struct LayerTable<T> {
24                $(pub $field_name: T,)*
25            }
26
27            pub type Layers = LayerTable<Option<$crate::Entity>>;
28
29            #[derive(Debug, Clone, Copy, PartialEq, Eq)]
30            pub enum Layer {
31                $($variant_name,)*
32            }
33
34            impl<T> Default for LayerTable<Option<T>> {
35                fn default() -> Self {
36                    Self {
37                        $($field_name: None,)*
38                    }
39                }
40            }
41
42            impl $crate::Layers for Layers {
43                type Layer = Layer;
44                fn select_field_mut(&mut self, layer: Self::Layer) -> &mut Option<$crate::Entity> {
45                    match layer {
46                        $(Layer::$variant_name => &mut self.$field_name,)*
47                    }
48                }
49            }
50
51            impl<T> LayerTable<T> {
52                #[allow(unused)]
53                pub fn map<U, F: FnMut(&T) -> U>(&self, mut f: F) -> LayerTable<U> {
54                    LayerTable {
55                        $($field_name: f(&self.$field_name),)*
56                    }
57                }
58
59                #[allow(unused)]
60                pub fn for_each<F: FnMut(&T)>(&self, mut f: F) {
61                    $(f(&self.$field_name);)*
62                }
63
64                #[allow(unused)]
65                pub fn for_each_enumerate<F: FnMut(&T, Layer)>(&self, mut f: F) {
66                    $(f(&self.$field_name, Layer::$variant_name);)*
67                }
68            }
69
70            impl<T> LayerTable<Option<T>> {
71                #[allow(unused)]
72                pub fn option_map<U, F: FnMut(&T) -> U>(&self, mut f: F) -> LayerTable<Option<U>> {
73                    self.map(|ot| ot.as_ref().map(|t| f(t)))
74                }
75
76                #[allow(unused)]
77                pub fn option_and_then<U, F: FnMut(&T) -> Option<U>>(&self, mut f: F) -> LayerTable<Option<U>> {
78                    self.map(|ot| ot.as_ref().and_then(|t| f(t)))
79                }
80
81                #[allow(unused)]
82                pub fn option_for_each<F: FnMut(&T)>(&self, mut f: F) {
83                    $(if let Some(t) = self.$field_name.as_ref() { f(t); })*
84                }
85
86                #[allow(unused)]
87                pub fn option_for_each_enumerate<F: FnMut(&T, Layer)>(&self, mut f: F) {
88                    $(if let Some(t) = self.$field_name.as_ref() { f(t, Layer::$variant_name); })*
89                }
90            }
91        }
92    }
93}
94
95#[cfg(feature = "serialize")]
96#[macro_export]
97macro_rules! declare_layers_module {
98    { $module_name:ident { $($field_name:ident: $variant_name:ident,)* } } => {
99        mod $module_name {
100            #[derive(Debug, Clone, Copy, PartialEq, Eq, $crate::serde::Serialize, $crate::serde::Deserialize)]
101            pub struct LayerTable<T> {
102                $(pub $field_name: T,)*
103            }
104
105            pub type Layers = LayerTable<Option<$crate::Entity>>;
106
107            #[derive(Debug, Clone, Copy, PartialEq, Eq, $crate::serde::Serialize, $crate::serde::Deserialize)]
108            pub enum Layer {
109                $($variant_name,)*
110            }
111
112            impl<T> Default for LayerTable<Option<T>> {
113                fn default() -> Self {
114                    Self {
115                        $($field_name: None,)*
116                    }
117                }
118            }
119
120            impl $crate::Layers for Layers {
121                type Layer = Layer;
122                fn select_field_mut(&mut self, layer: Self::Layer) -> &mut Option<$crate::Entity> {
123                    match layer {
124                        $(Layer::$variant_name => &mut self.$field_name,)*
125                    }
126                }
127            }
128
129            impl<T> LayerTable<T> {
130                #[allow(unused)]
131                pub fn map<U, F: FnMut(&T) -> U>(&self, mut f: F) -> LayerTable<U> {
132                    LayerTable {
133                        $($field_name: f(&self.$field_name),)*
134                    }
135                }
136
137                #[allow(unused)]
138                pub fn for_each<F: FnMut(&T)>(&self, mut f: F) {
139                    $(f(&self.$field_name);)*
140                }
141
142                #[allow(unused)]
143                pub fn for_each_enumerate<F: FnMut(&T, Layer)>(&self, mut f: F) {
144                    $(f(&self.$field_name, Layer::$variant_name);)*
145                }
146            }
147
148            impl<T> LayerTable<Option<T>> {
149                #[allow(unused)]
150                pub fn option_map<U, F: FnMut(&T) -> U>(&self, mut f: F) -> LayerTable<Option<U>> {
151                    self.map(|ot| ot.as_ref().map(|t| f(t)))
152                }
153
154                #[allow(unused)]
155                pub fn option_and_then<U, F: FnMut(&T) -> Option<U>>(&self, mut f: F) -> LayerTable<Option<U>> {
156                    self.map(|ot| ot.as_ref().and_then(|t| f(t)))
157                }
158
159                #[allow(unused)]
160                pub fn option_for_each<F: FnMut(&T)>(&self, mut f: F) {
161                    $(if let Some(t) = self.$field_name.as_ref() { f(t); })*
162                }
163
164                #[allow(unused)]
165                pub fn option_for_each_enumerate<F: FnMut(&T, Layer)>(&self, mut f: F) {
166                    $(if let Some(t) = self.$field_name.as_ref() { f(t, Layer::$variant_name); })*
167                }
168            }
169        }
170    }
171}
172
173#[cfg_attr(feature = "serialize", derive(Serialize, Deserialize))]
174#[derive(Debug, Clone, Copy, PartialEq, Eq)]
175pub struct Location<L> {
176    pub coord: ICoord,
177    pub layer: Option<L>,
178}
179
180impl<L> From<(ICoord, L)> for Location<L> {
181    fn from((coord, layer): (ICoord, L)) -> Self {
182        Self {
183            coord,
184            layer: Some(layer),
185        }
186    }
187}
188
189#[derive(Debug)]
190pub struct SpatialTable<L: Layers> {
191    location_component: ComponentTable<Location<L::Layer>>,
192    spatial_grid: Grid<L>,
193}
194
195pub type Enumerate<'a, L> = grid_2d::GridEnumerate<'a, L>;
196
197impl<L: Layers> SpatialTable<L> {
198    pub fn new(size: UCoord) -> Self {
199        let location_component = ComponentTable::default();
200        let spatial_grid = Grid::new_default(size);
201        Self {
202            location_component,
203            spatial_grid,
204        }
205    }
206    pub fn clear(&mut self) {
207        self.location_component.clear();
208        for cell in self.spatial_grid.iter_mut() {
209            *cell = Default::default();
210        }
211    }
212    pub fn enumerate(&self) -> Enumerate<'_, L> {
213        self.spatial_grid.enumerate()
214    }
215    pub fn grid_size(&self) -> UCoord {
216        self.spatial_grid.size()
217    }
218    pub fn layers_at(&self, coord: ICoord) -> Option<&L> {
219        self.spatial_grid.get(coord)
220    }
221    pub fn layers_at_checked(&self, coord: ICoord) -> &L {
222        self.spatial_grid.get_checked(coord)
223    }
224    pub fn location_of(&self, entity: Entity) -> Option<&Location<L::Layer>> {
225        self.location_component.get(entity)
226    }
227    pub fn coord_of(&self, entity: Entity) -> Option<ICoord> {
228        self.location_of(entity).map(|l| l.coord)
229    }
230    pub fn layer_of(&self, entity: Entity) -> Option<L::Layer> {
231        self.location_of(entity).and_then(|l| l.layer)
232    }
233    pub fn update(
234        &mut self,
235        entity: Entity,
236        location: Location<L::Layer>,
237    ) -> Result<(), UpdateError> {
238        if let Some(layer) = location.layer {
239            let cell = self
240                .spatial_grid
241                .get_mut(location.coord)
242                .ok_or(UpdateError::DestinationOutOfBounds)?;
243            insert_layer(cell, entity, layer)?;
244        }
245        if let Some(original_location) = self.location_component.insert(entity, location) {
246            let original_cell = self.spatial_grid.get_checked_mut(original_location.coord);
247            if let Some(original_layer) = original_location.layer {
248                let should_match_entity = clear_layer(original_cell, original_layer);
249                debug_assert_eq!(
250                    should_match_entity,
251                    Some(entity),
252                    "Current location of entity doesn't contain entity in spatial grid"
253                );
254            }
255        }
256        Ok(())
257    }
258    pub fn update_coord(&mut self, entity: Entity, coord: ICoord) -> Result<(), UpdateError> {
259        if let Some(location) = self.location_component.get_mut(entity) {
260            if coord != location.coord {
261                if let Some(layer) = location.layer {
262                    let cell = self
263                        .spatial_grid
264                        .get_mut(coord)
265                        .ok_or(UpdateError::DestinationOutOfBounds)?;
266                    insert_layer(cell, entity, layer)?;
267                    let original_cell = self.spatial_grid.get_checked_mut(location.coord);
268                    let should_match_entity = clear_layer(original_cell, layer);
269                    debug_assert_eq!(
270                        should_match_entity,
271                        Some(entity),
272                        "Current location of entity doesn't contain entity in spatial grid"
273                    );
274                }
275                location.coord = coord;
276            }
277            Ok(())
278        } else {
279            self.update(entity, Location { coord, layer: None })
280        }
281    }
282    pub fn update_layer(
283        &mut self,
284        entity: Entity,
285        layer: L::Layer,
286    ) -> Result<(), UpdateLayerError> {
287        if let Some(location) = self.location_component.get_mut(entity) {
288            if Some(layer) != location.layer {
289                debug_assert!(
290                    location.coord.is_valid(self.spatial_grid.size()),
291                    "Current location is outside the bounds of spatial grid"
292                );
293                let cell = self.spatial_grid.get_mut(location.coord).unwrap();
294                let dest_entity_slot = cell.select_field_mut(layer);
295                if let Some(dest_entity) = dest_entity_slot {
296                    return Err(UpdateLayerError::OccupiedBy(*dest_entity));
297                }
298                *dest_entity_slot = Some(entity);
299                if let Some(current_layer) = location.layer {
300                    let source_entity_slot = cell.select_field_mut(current_layer);
301                    debug_assert_eq!(*source_entity_slot, Some(entity));
302                    *source_entity_slot = None;
303                }
304                location.layer = Some(layer);
305            }
306            Ok(())
307        } else {
308            Err(UpdateLayerError::EntityHasNoICoord)
309        }
310    }
311    pub fn clear_layer(&mut self, entity: Entity) -> Result<(), EntityHasNoICoord> {
312        if let Some(location) = self.location_component.get_mut(entity) {
313            if let Some(layer) = location.layer {
314                debug_assert!(
315                    location.coord.is_valid(self.spatial_grid.size()),
316                    "Current location is outside the bounds of spatial grid"
317                );
318                let cell = self.spatial_grid.get_mut(location.coord).unwrap();
319                let source_entity_slot = cell.select_field_mut(layer);
320                debug_assert_eq!(*source_entity_slot, Some(entity));
321                *source_entity_slot = None;
322                location.layer = None;
323            }
324            Ok(())
325        } else {
326            Err(EntityHasNoICoord)
327        }
328    }
329    pub fn remove(&mut self, entity: Entity) {
330        if let Some(location) = self.location_component.remove(entity)
331            && let Some(layer) = location.layer
332        {
333            clear_layer(self.spatial_grid.get_checked_mut(location.coord), layer);
334        }
335    }
336    #[cfg(feature = "serialize")]
337    fn to_serialize(&self) -> SpatialSerialize<L::Layer> {
338        SpatialSerialize {
339            entries: self.location_component.entries().clone(),
340            size: self.spatial_grid.size(),
341        }
342    }
343    #[cfg(feature = "serialize")]
344    fn from_serialize(SpatialSerialize { entries, size }: SpatialSerialize<L::Layer>) -> Self {
345        let location_component = entries.into_component_table();
346        let mut spatial_grid: Grid<L> = Grid::new_default(size);
347        for (entity, location) in location_component.iter() {
348            if let Some(layer) = location.layer {
349                let cell = spatial_grid.get_checked_mut(location.coord);
350                let slot = cell.select_field_mut(layer);
351                assert!(slot.is_none());
352                *slot = Some(entity);
353            }
354        }
355        Self {
356            location_component,
357            spatial_grid,
358        }
359    }
360}
361
362struct OccupiedBy(pub Entity);
363
364impl From<OccupiedBy> for UpdateError {
365    fn from(occupied_by: OccupiedBy) -> UpdateError {
366        UpdateError::OccupiedBy(occupied_by.0)
367    }
368}
369
370#[derive(Debug, Clone, Copy, PartialEq, Eq)]
371pub enum UpdateError {
372    OccupiedBy(Entity),
373    DestinationOutOfBounds,
374}
375
376impl UpdateError {
377    pub fn unwrap_occupied_by(self) -> Entity {
378        match self {
379            Self::OccupiedBy(entity) => entity,
380            _ => panic!("unexpected {:?} (expected OccupiedBy(_))", self),
381        }
382    }
383}
384
385#[derive(Debug, Clone, Copy, PartialEq, Eq)]
386pub enum UpdateLayerError {
387    OccupiedBy(Entity),
388    EntityHasNoICoord,
389}
390
391impl UpdateLayerError {
392    pub fn unwrap_occupied_by(self) -> Entity {
393        match self {
394            Self::OccupiedBy(entity) => entity,
395            _ => panic!("unexpected {:?} (expected OccupiedBy(_))", self),
396        }
397    }
398}
399
400#[derive(Debug, Clone, Copy, PartialEq, Eq)]
401pub struct EntityHasNoICoord;
402
403fn insert_layer<L: Layers>(
404    layers: &mut L,
405    entity: Entity,
406    layer: L::Layer,
407) -> Result<(), OccupiedBy> {
408    let layer_field = layers.select_field_mut(layer);
409    if let Some(&occupant) = layer_field.as_ref() {
410        Err(OccupiedBy(occupant))
411    } else {
412        *layer_field = Some(entity);
413        Ok(())
414    }
415}
416fn clear_layer<L: Layers>(layers: &mut L, layer: L::Layer) -> Option<Entity> {
417    layers.select_field_mut(layer).take()
418}
419
420#[cfg(feature = "serialize")]
421#[derive(Serialize, Deserialize)]
422struct SpatialSerialize<L> {
423    entries: ComponentTableEntries<Location<L>>,
424    size: UCoord,
425}
426
427#[cfg(feature = "serialize")]
428impl<L: Layers> Serialize for SpatialTable<L>
429where
430    L::Layer: Serialize,
431{
432    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
433        self.to_serialize().serialize(s)
434    }
435}
436
437#[cfg(feature = "serialize")]
438impl<'a, L: Layers> Deserialize<'a> for SpatialTable<L>
439where
440    L::Layer: Deserialize<'a>,
441{
442    fn deserialize<D: serde::Deserializer<'a>>(d: D) -> Result<Self, D::Error> {
443        Deserialize::deserialize(d).map(Self::from_serialize)
444    }
445}
446
447#[cfg(test)]
448mod test {
449    declare_layers_module! {
450        layers {
451            feature: Feature,
452            character: Character,
453        }
454    }
455    use layers::{Layer, Layers};
456    type SpatialTable = super::SpatialTable<Layers>;
457    use super::{ICoord, Location, UCoord, UpdateError, UpdateLayerError};
458    use entity_table::EntityAllocator;
459
460    #[test]
461    fn test() {
462        let mut entity_allocator = EntityAllocator::default();
463        let mut spatial_table = SpatialTable::new(UCoord::new(10, 10));
464        let entity_a = entity_allocator.alloc();
465        let entity_b = entity_allocator.alloc();
466        let entity_c = entity_allocator.alloc();
467
468        assert_eq!(spatial_table.location_of(entity_a), None);
469
470        // try to place a feature out of bounds - should fail
471        assert_eq!(
472            spatial_table.update(
473                entity_a,
474                Location {
475                    coord: ICoord::new(-1, 10),
476                    layer: Some(Layer::Feature),
477                },
478            ),
479            Err(UpdateError::DestinationOutOfBounds),
480        );
481
482        // entity should not have been added
483        assert_eq!(spatial_table.location_of(entity_a), None);
484
485        // try to place a feature at valid coord
486        assert_eq!(
487            spatial_table.update(
488                entity_a,
489                Location {
490                    coord: ICoord::new(4, 2),
491                    layer: Some(Layer::Feature),
492                },
493            ),
494            Ok(()),
495        );
496        assert_eq!(
497            spatial_table.location_of(entity_a).cloned(),
498            Some(Location {
499                coord: ICoord::new(4, 2),
500                layer: Some(Layer::Feature)
501            })
502        );
503
504        // move feature to new coord
505        assert_eq!(
506            spatial_table.update_coord(entity_a, ICoord::new(6, 7)),
507            Ok(()),
508        );
509        assert_eq!(
510            spatial_table.location_of(entity_a).cloned(),
511            Some(Location {
512                coord: ICoord::new(6, 7),
513                layer: Some(Layer::Feature)
514            })
515        );
516
517        assert_eq!(spatial_table.location_of(entity_b), None);
518
519        // try to add a new feature on top - should fail
520        assert_eq!(
521            spatial_table.update(
522                entity_b,
523                Location {
524                    coord: ICoord::new(6, 7),
525                    layer: Some(Layer::Feature),
526                },
527            ),
528            Err(UpdateError::OccupiedBy(entity_a)),
529        );
530
531        assert_eq!(spatial_table.location_of(entity_b), None);
532
533        // add new feature in different coord
534        assert_eq!(
535            spatial_table.update(
536                entity_b,
537                Location {
538                    coord: ICoord::new(6, 8),
539                    layer: Some(Layer::Feature),
540                },
541            ),
542            Ok(()),
543        );
544
545        // try to move it to coord with existing feature
546        assert_eq!(
547            spatial_table.update_coord(entity_b, ICoord::new(6, 7)),
548            Err(UpdateError::OccupiedBy(entity_a)),
549        );
550
551        assert_eq!(spatial_table.coord_of(entity_b), Some(ICoord::new(6, 8)));
552
553        assert_eq!(spatial_table.location_of(entity_c), None);
554
555        // add a character on top of an entity_a
556        assert_eq!(
557            spatial_table.update(
558                entity_c,
559                Location {
560                    coord: ICoord::new(6, 7),
561                    layer: Some(Layer::Character),
562                },
563            ),
564            Ok(()),
565        );
566        assert_eq!(spatial_table.coord_of(entity_c), Some(ICoord::new(6, 7)));
567
568        assert_eq!(
569            *spatial_table.layers_at_checked(ICoord::new(6, 7)),
570            Layers {
571                feature: Some(entity_a),
572                character: Some(entity_c),
573            },
574        );
575        assert_eq!(
576            *spatial_table.layers_at_checked(ICoord::new(6, 8)),
577            Layers {
578                feature: Some(entity_b),
579                character: None,
580            },
581        );
582
583        spatial_table
584            .update_layer(entity_a, Layer::Feature)
585            .unwrap();
586        assert_eq!(
587            *spatial_table.layers_at_checked(ICoord::new(6, 7)),
588            Layers {
589                feature: Some(entity_a),
590                character: Some(entity_c),
591            },
592        );
593
594        assert_eq!(
595            spatial_table.update_layer(entity_a, Layer::Character),
596            Err(UpdateLayerError::OccupiedBy(entity_c))
597        );
598
599        spatial_table
600            .update_layer(entity_b, Layer::Character)
601            .unwrap();
602        assert_eq!(
603            *spatial_table.layers_at_checked(ICoord::new(6, 8)),
604            Layers {
605                feature: None,
606                character: Some(entity_b),
607            },
608        );
609        assert_eq!(spatial_table.layer_of(entity_b), Some(Layer::Character));
610        spatial_table
611            .update_layer(entity_b, Layer::Feature)
612            .unwrap();
613        assert_eq!(spatial_table.layer_of(entity_b), Some(Layer::Feature));
614
615        spatial_table.remove(entity_a);
616        assert_eq!(
617            *spatial_table.layers_at_checked(ICoord::new(6, 7)),
618            Layers {
619                feature: None,
620                character: Some(entity_c),
621            },
622        );
623        assert_eq!(
624            spatial_table.update_coord(entity_b, ICoord::new(6, 7)),
625            Ok(()),
626        );
627        assert_eq!(
628            *spatial_table.layers_at_checked(ICoord::new(6, 7)),
629            Layers {
630                feature: Some(entity_b),
631                character: Some(entity_c),
632            },
633        );
634        assert_eq!(
635            *spatial_table.layers_at_checked(ICoord::new(6, 8)),
636            Layers {
637                feature: None,
638                character: None,
639            },
640        );
641
642        spatial_table.clear_layer(entity_b).unwrap();
643        assert_eq!(
644            *spatial_table.layers_at_checked(ICoord::new(6, 7)),
645            Layers {
646                feature: None,
647                character: Some(entity_c),
648            },
649        );
650        assert_eq!(spatial_table.coord_of(entity_b), Some(ICoord::new(6, 7)));
651        assert_eq!(spatial_table.layer_of(entity_b), None);
652    }
653}