steel_core/entity/manager/
spatial_query.rs

1//! Ordered query hits gathered from already sorted spatial cells.
2
3use std::{slice, vec::IntoIter};
4
5use rustc_hash::FxHashSet;
6use smallvec::SmallVec;
7use steel_utils::WorldAabb;
8
9use super::{EntityEntry, EntityQueryOrder, EntitySlot, ManagerState};
10
11const MAX_MERGED_QUERY_CELLS: usize = 8;
12
13#[expect(
14    clippy::large_enum_variant,
15    reason = "Common query cursors stay inline to avoid an allocation per spatial lookup"
16)]
17pub(super) enum EntityQueryEntries<'a> {
18    Single {
19        state: &'a ManagerState,
20        aabb: Option<WorldAabb>,
21        remaining: slice::Iter<'a, EntitySlot>,
22    },
23    Merged {
24        state: &'a ManagerState,
25        aabb: Option<WorldAabb>,
26        cursors: SmallVec<[EntityQueryCursor<'a>; MAX_MERGED_QUERY_CELLS]>,
27    },
28    Sorted(IntoIter<&'a EntityEntry>),
29}
30
31pub(super) struct EntityQueryCursor<'a> {
32    entry: &'a EntityEntry,
33    order: EntityQueryOrder,
34    remaining: slice::Iter<'a, EntitySlot>,
35}
36
37impl<'a> EntityQueryCursor<'a> {
38    fn new(
39        state: &'a ManagerState,
40        aabb: Option<&WorldAabb>,
41        cell: &'a [EntitySlot],
42    ) -> Option<Self> {
43        let mut remaining = cell.iter();
44        let entry = next_candidate(state, aabb, &mut remaining)?;
45
46        Some(Self {
47            entry,
48            order: entry.query_order,
49            remaining,
50        })
51    }
52
53    fn advance(&mut self, state: &'a ManagerState, aabb: Option<&WorldAabb>) -> bool {
54        let Some(entry) = next_candidate(state, aabb, &mut self.remaining) else {
55            return false;
56        };
57
58        self.entry = entry;
59        self.order = entry.query_order;
60
61        true
62    }
63}
64
65fn next_candidate<'a>(
66    state: &'a ManagerState,
67    aabb: Option<&WorldAabb>,
68    remaining: &mut slice::Iter<'a, EntitySlot>,
69) -> Option<&'a EntityEntry> {
70    remaining.find_map(|slot| {
71        let entry = state.live_by_id.entry(*slot);
72
73        aabb.is_none_or(|aabb| entry.bounding_box.intersects(*aabb))
74            .then_some(entry)
75    })
76}
77
78impl<'a> EntityQueryEntries<'a> {
79    pub(super) fn new(
80        state: &'a ManagerState,
81        aabb: Option<&WorldAabb>,
82        cells: SmallVec<[&'a [EntitySlot]; MAX_MERGED_QUERY_CELLS]>,
83    ) -> Self {
84        if let [cell] = cells.as_slice() {
85            return Self::Single {
86                state,
87                aabb: aabb.copied(),
88                remaining: cell.iter(),
89            };
90        }
91
92        if cells.len() <= MAX_MERGED_QUERY_CELLS {
93            return Self::Merged {
94                state,
95                aabb: aabb.copied(),
96                cursors: cells
97                    .into_iter()
98                    .filter_map(|cell| EntityQueryCursor::new(state, aabb, cell))
99                    .collect(),
100            };
101        }
102
103        let mut entity_ids = FxHashSet::default();
104
105        for cell in cells {
106            entity_ids.extend(cell.iter().copied());
107        }
108
109        let mut entries = entity_ids
110            .into_iter()
111            .map(|slot| state.live_by_id.entry(slot))
112            .filter(|entry| aabb.is_none_or(|aabb| entry.bounding_box.intersects(*aabb)))
113            .collect::<Vec<_>>();
114
115        entries.sort_unstable_by_key(|entry| entry.query_order);
116        Self::Sorted(entries.into_iter())
117    }
118}
119
120impl<'a> Iterator for EntityQueryEntries<'a> {
121    type Item = &'a EntityEntry;
122
123    fn fold<B, F>(self, initial: B, mut visit: F) -> B
124    where
125        F: FnMut(B, Self::Item) -> B,
126    {
127        match self {
128            Self::Single {
129                state,
130                aabb,
131                remaining,
132            } => remaining
133                .map(|slot| state.live_by_id.entry(*slot))
134                .filter(|entry| aabb.is_none_or(|aabb| entry.bounding_box.intersects(aabb)))
135                .fold(initial, visit),
136            Self::Sorted(entries) => entries.fold(initial, visit),
137            entries @ Self::Merged { .. } => {
138                let mut result = initial;
139
140                for entry in entries {
141                    result = visit(result, entry);
142                }
143
144                result
145            }
146        }
147    }
148
149    fn next(&mut self) -> Option<Self::Item> {
150        match self {
151            Self::Single {
152                state,
153                aabb,
154                remaining,
155            } => next_candidate(state, aabb.as_ref(), remaining),
156            Self::Merged {
157                state,
158                aabb,
159                cursors,
160            } => {
161                let cursor = cursors.iter().min_by_key(|cursor| cursor.order)?;
162                let entry = cursor.entry;
163                let order = cursor.order;
164
165                let mut index = 0;
166
167                while index < cursors.len() {
168                    if cursors[index].order != order || cursors[index].advance(state, aabb.as_ref())
169                    {
170                        index += 1;
171                    } else {
172                        cursors.swap_remove(index);
173                    }
174                }
175
176                Some(entry)
177            }
178            Self::Sorted(entries) => entries.next(),
179        }
180    }
181}
steel_core/entity/manager/spatial_query.rs

steel_core/entity/manager/
spatial_query.rs
