anput_spatial/
lib.rs

1use anput::{
2    entity::Entity,
3    query::{DynamicQueryFilter, DynamicQueryItem, TypedLookupFetch},
4    scheduler::GraphSchedulerQuickPlugin,
5    systems::SystemContext,
6    universe::{Plugin, Res},
7    world::World,
8};
9use rstar::{primitives::GeomWithData, Envelope, PointDistance, RTree, RTreeObject};
10use std::{error::Error, marker::PhantomData};
11
12pub mod third_party {
13    pub use rstar;
14}
15
16pub struct SpatialPartitioningPlugin<const LOCKING: bool, Extractor: SpatialExtractor>(
17    PhantomData<fn() -> Extractor>,
18);
19
20impl<const LOCKING: bool, Extractor: SpatialExtractor>
21    SpatialPartitioningPlugin<LOCKING, Extractor>
22{
23    pub fn make() -> impl Plugin {
24        GraphSchedulerQuickPlugin::<LOCKING, Self>::default()
25            .resource(SpatialPartitioning::<Extractor>::default())
26            .system(
27                spatial_partitioning::<LOCKING, Extractor>,
28                "spatial_partitioning",
29                (),
30            )
31            .commit()
32    }
33}
34
35pub struct SpatialPartitioning<Extractor: SpatialExtractor> {
36    tree: RTree<GeomWithData<Extractor::SpatialObject, Entity>>,
37}
38
39impl<Extractor: SpatialExtractor> Default for SpatialPartitioning<Extractor> {
40    fn default() -> Self {
41        Self {
42            tree: RTree::default(),
43        }
44    }
45}
46
47impl<Extractor: SpatialExtractor> SpatialPartitioning<Extractor> {
48    pub fn rebuild<const LOCKING: bool>(&mut self, world: &World) {
49        self.tree = RTree::bulk_load(
50            Extractor::extract::<LOCKING>(world)
51                .map(|(entity, object)| GeomWithData::new(object, entity))
52                .collect::<Vec<_>>(),
53        );
54    }
55
56    pub fn tree(&self) -> &RTree<GeomWithData<Extractor::SpatialObject, Entity>> {
57        &self.tree
58    }
59
60    pub fn nearest_entities(
61        &self,
62        point: &<<Extractor::SpatialObject as RTreeObject>::Envelope as Envelope>::Point,
63    ) -> impl Iterator<Item = Entity> + '_ {
64        self.tree.nearest_neighbor_iter(point).map(|geom| geom.data)
65    }
66
67    pub fn locate_contained_entities(
68        &self,
69        envelope: &<Extractor::SpatialObject as RTreeObject>::Envelope,
70    ) -> impl Iterator<Item = Entity> + '_ {
71        self.tree.locate_in_envelope(envelope).map(|geom| geom.data)
72    }
73
74    pub fn locate_intersecting_entities(
75        &self,
76        envelope: &<Extractor::SpatialObject as RTreeObject>::Envelope,
77    ) -> impl Iterator<Item = Entity> + '_ {
78        self.tree
79            .locate_in_envelope_intersecting(envelope)
80            .map(|geom| geom.data)
81    }
82
83    pub fn nearest_query<'a, const LOCKING: bool, Fetch: TypedLookupFetch<'a, LOCKING>>(
84        &'a self,
85        world: &'a World,
86        point: &<<Extractor::SpatialObject as RTreeObject>::Envelope as Envelope>::Point,
87    ) -> impl Iterator<Item = Fetch::Value> {
88        world.lookup::<LOCKING, Fetch>(self.nearest_entities(point))
89    }
90
91    pub fn nearest_dynamic_query<'a, const LOCKING: bool>(
92        &'a self,
93        world: &'a World,
94        point: &<<Extractor::SpatialObject as RTreeObject>::Envelope as Envelope>::Point,
95        filter: &DynamicQueryFilter,
96    ) -> impl Iterator<Item = DynamicQueryItem<'a>> {
97        world.dynamic_lookup::<LOCKING>(filter, self.nearest_entities(point))
98    }
99
100    pub fn locate_contained_query<'a, const LOCKING: bool, Fetch: TypedLookupFetch<'a, LOCKING>>(
101        &'a self,
102        world: &'a World,
103        envelope: &<Extractor::SpatialObject as RTreeObject>::Envelope,
104    ) -> impl Iterator<Item = Fetch::Value> {
105        world.lookup::<LOCKING, Fetch>(self.locate_contained_entities(envelope))
106    }
107
108    pub fn locate_contained_dynamic_query<'a, const LOCKING: bool>(
109        &'a self,
110        world: &'a World,
111        envelope: &<Extractor::SpatialObject as RTreeObject>::Envelope,
112        filter: &DynamicQueryFilter,
113    ) -> impl Iterator<Item = DynamicQueryItem<'a>> {
114        world.dynamic_lookup::<LOCKING>(filter, self.locate_contained_entities(envelope))
115    }
116
117    pub fn locate_intersecting_query<
118        'a,
119        const LOCKING: bool,
120        Fetch: TypedLookupFetch<'a, LOCKING>,
121    >(
122        &'a self,
123        world: &'a World,
124        envelope: &<Extractor::SpatialObject as RTreeObject>::Envelope,
125    ) -> impl Iterator<Item = Fetch::Value> {
126        world.lookup::<LOCKING, Fetch>(self.locate_intersecting_entities(envelope))
127    }
128
129    pub fn locate_intersecting_dynamic_query<'a, const LOCKING: bool>(
130        &'a self,
131        world: &'a World,
132        envelope: &<Extractor::SpatialObject as RTreeObject>::Envelope,
133        filter: &DynamicQueryFilter,
134    ) -> impl Iterator<Item = DynamicQueryItem<'a>> {
135        world.dynamic_lookup::<LOCKING>(filter, self.locate_intersecting_entities(envelope))
136    }
137}
138
139pub fn spatial_partitioning<const LOCKING: bool, Extractor: SpatialExtractor>(
140    context: SystemContext,
141) -> Result<(), Box<dyn Error>> {
142    let (world, mut partitioning) =
143        context.fetch::<(&World, Res<LOCKING, &mut SpatialPartitioning<Extractor>>)>()?;
144
145    partitioning.rebuild::<LOCKING>(world);
146
147    Ok(())
148}
149
150pub trait SpatialExtractor: 'static
151where
152    <<Self as SpatialExtractor>::SpatialObject as RTreeObject>::Envelope: Send + Sync,
153{
154    type SpatialObject: RTreeObject + PointDistance + Send + Sync;
155
156    fn extract<const LOCKING: bool>(
157        world: &World,
158    ) -> impl Iterator<Item = (Entity, Self::SpatialObject)>;
159}
anput_spatial/lib.rs

anput_spatial/
lib.rs
