anput_spatial/
lib.rs

1use anput::{
2    entity::Entity,
3    query::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 locate_contained_query<'a, const LOCKING: bool, Fetch: TypedLookupFetch<'a, LOCKING>>(
92        &'a self,
93        world: &'a World,
94        envelope: &<Extractor::SpatialObject as RTreeObject>::Envelope,
95    ) -> impl Iterator<Item = Fetch::Value> {
96        world.lookup::<LOCKING, Fetch>(self.locate_contained_entities(envelope))
97    }
98
99    pub fn locate_intersecting_query<
100        'a,
101        const LOCKING: bool,
102        Fetch: TypedLookupFetch<'a, LOCKING>,
103    >(
104        &'a self,
105        world: &'a World,
106        envelope: &<Extractor::SpatialObject as RTreeObject>::Envelope,
107    ) -> impl Iterator<Item = Fetch::Value> {
108        world.lookup::<LOCKING, Fetch>(self.locate_intersecting_entities(envelope))
109    }
110}
111
112pub fn spatial_partitioning<const LOCKING: bool, Extractor: SpatialExtractor>(
113    context: SystemContext,
114) -> Result<(), Box<dyn Error>> {
115    let (world, mut partitioning) =
116        context.fetch::<(&World, Res<LOCKING, &mut SpatialPartitioning<Extractor>>)>()?;
117
118    partitioning.rebuild::<LOCKING>(world);
119
120    Ok(())
121}
122
123pub trait SpatialExtractor: 'static
124where
125    <<Self as SpatialExtractor>::SpatialObject as RTreeObject>::Envelope: Send + Sync,
126{
127    type SpatialObject: RTreeObject + PointDistance + Send + Sync;
128
129    fn extract<const LOCKING: bool>(
130        world: &World,
131    ) -> impl Iterator<Item = (Entity, Self::SpatialObject)>;
132}
anput_spatial/lib.rs

anput_spatial/
lib.rs
