anput_spatial/
lib.rs

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

anput_spatial/
lib.rs
