moirai/query/
effects.rs

1//! Side effects available during query traversal callbacks.
2//!
3//! [`QueryCommands`] enqueues deferred spawn, despawn, insert, remove, and bundle operations.
4//! [`QueryEffects`] also exposes declared event emission when the active stage permits it.
5
6use alloc::string::String;
7use core::any::type_name;
8
9use crate::command::{CommandOp, CommandQueue};
10use crate::entity::{AllocatorError, EntityAllocator, EntityId};
11use crate::operation::StageOperation;
12use crate::query::QueryError;
13use crate::world::guard::RunGuard;
14use crate::world::{Bundle, BundleWriter, WorldError, WorldEvents, WorldOwner};
15
16/// Deferred structural commands available while a query callback runs.
17pub struct QueryCommands<'w> {
18    allocator: &'w mut EntityAllocator,
19    queue: &'w mut CommandQueue,
20}
21
22impl<'w> QueryCommands<'w> {
23    /// Reserves a new entity id and queues a deferred spawn.
24    pub fn spawn(&mut self) -> Result<EntityId, QueryError> {
25        let entity = self
26            .allocator
27            .reserve()
28            .map_err(map_allocator_error_query)?;
29        self.queue.push(CommandOp::SpawnReserved { entity });
30        Ok(entity)
31    }
32
33    /// Queues a deferred despawn when the target is live or reserved.
34    pub fn despawn(&mut self, entity: EntityId) -> Result<(), QueryError> {
35        self.ensure_target(entity)?;
36        self.queue.push(CommandOp::Despawn { entity });
37        Ok(())
38    }
39
40    /// Queues a deferred component insert for a live or reserved entity.
41    pub fn insert<T: 'static>(&mut self, entity: EntityId, value: T) -> Result<(), QueryError> {
42        self.ensure_target(entity)?;
43        self.queue
44            .enqueue_insert(entity, value)
45            .map_err(map_command_error)
46    }
47
48    /// Queues a deferred component removal for a live or reserved entity.
49    pub fn remove<T: 'static>(&mut self, entity: EntityId) -> Result<(), QueryError> {
50        self.ensure_target(entity)?;
51        self.queue
52            .enqueue_remove::<T>(entity)
53            .map_err(map_command_error)
54    }
55
56    /// Queues a deferred bundle write for a live or reserved entity.
57    pub fn insert_bundle<B: Bundle>(
58        &mut self,
59        entity: EntityId,
60        bundle: B,
61    ) -> Result<(), QueryError> {
62        self.ensure_target(entity)?;
63        let queue_len = self.queue.len();
64        match bundle.write(&mut BundleWriter::query(self.allocator, self.queue, entity)) {
65            Ok(()) => Ok(()),
66            Err(error) => {
67                self.queue.truncate(queue_len);
68                Err(map_command_error(error))
69            }
70        }
71    }
72
73    fn ensure_target(&self, entity: EntityId) -> Result<(), QueryError> {
74        if self.allocator.is_alive(entity) || self.allocator.is_reserved(entity) {
75            Ok(())
76        } else {
77            Err(QueryError::CommandRejected {
78                detail: alloc::format!("stale command target {entity:?}"),
79            })
80        }
81    }
82}
83
84/// Borrow-checked side-effect surface for prepared and eager query callbacks.
85pub struct QueryEffects<'w> {
86    owner: WorldOwner,
87    run_guard: RunGuard,
88    command_queue: &'w mut CommandQueue,
89    allocator: &'w mut EntityAllocator,
90    events: &'w mut WorldEvents,
91}
92
93impl<'w> QueryEffects<'w> {
94    pub(crate) fn from_parts(
95        command_queue: &'w mut CommandQueue,
96        allocator: &'w mut EntityAllocator,
97        events: &'w mut WorldEvents,
98        run_guard: RunGuard,
99        owner: WorldOwner,
100    ) -> Self {
101        Self {
102            owner,
103            run_guard,
104            command_queue,
105            allocator,
106            events,
107        }
108    }
109
110    /// Structural commands when the active stage operation is `Update`.
111    pub fn commands(&mut self) -> Result<QueryCommands<'_>, QueryError> {
112        match self.run_guard.operation() {
113            Some(StageOperation::Update) => {}
114            Some(StageOperation::Render) => {
115                return Err(QueryError::BorrowConflict {
116                    detail: String::from("structural commands are unavailable during Render"),
117                });
118            }
119            None => {
120                return Err(QueryError::BorrowConflict {
121                    detail: String::from(
122                        "structural commands require an active Update operation context",
123                    ),
124                });
125            }
126        }
127        Ok(QueryCommands {
128            allocator: self.allocator,
129            queue: self.command_queue,
130        })
131    }
132
133    /// Emits a declared event when the active run guard permits that event id.
134    pub fn send<E: Clone + 'static>(&mut self, event: E) -> Result<(), QueryError> {
135        let event_id = self
136            .events
137            .registry
138            .id_of::<E>(&self.owner)
139            .ok_or_else(|| QueryError::WrongQuery {
140                detail: alloc::format!("unregistered event {}", type_name::<E>()),
141            })?;
142        if !self.run_guard.permits_emit(&event_id) {
143            return Err(QueryError::WrongQuery {
144                detail: alloc::format!("undeclared event {}", event_id.index()),
145            });
146        }
147        self.events
148            .storage
149            .send(&event_id, event)
150            .map_err(|error| QueryError::WrongQuery {
151                detail: alloc::format!("{error:?}"),
152            })
153    }
154}
155
156fn map_allocator_error_query(error: AllocatorError) -> QueryError {
157    let detail = match error {
158        AllocatorError::GenerationOverflow => String::from("allocator generation overflow"),
159        AllocatorError::SlotRetired => String::from("allocator slot retired"),
160        AllocatorError::StaleEntity | AllocatorError::DoubleFree | AllocatorError::NotLive => {
161            String::from("allocator rejected entity")
162        }
163    };
164    QueryError::CommandRejected { detail }
165}
166
167fn map_command_error(error: WorldError) -> QueryError {
168    QueryError::CommandRejected {
169        detail: alloc::format!("{error:?}"),
170    }
171}
172
173#[cfg(test)]
174mod tests {
175    use super::*;
176    #[test]
177    fn send_ok_path_propagates_success() {
178        use crate::component::ComponentOptions;
179        use crate::event::{EventOptions, EventReaderStart};
180        use crate::operation::StageOperation;
181        use crate::world::WorldBuilder;
182
183        #[derive(Clone, Copy, Debug, PartialEq)]
184        struct Ping(u8);
185
186        let mut builder = WorldBuilder::new();
187        builder
188            .register_component::<Ping>(ComponentOptions::sparse())
189            .expect("component");
190        builder
191            .add_event::<Ping>(EventOptions::frame(StageOperation::Update))
192            .expect("event");
193        let mut world = builder.build().expect("world");
194        let entity = world.spawn().expect("spawn");
195        world.insert(entity, Ping(1)).expect("insert");
196        let mut reader = world
197            .event_reader::<Ping>(EventReaderStart::OldestRetained)
198            .expect("reader");
199        world.begin_run(StageOperation::Update).expect("begin");
200        world
201            .for_each_mut_with_effects::<Ping>(
202                &crate::query::QuerySpec::new(),
203                crate::query::QueryParams::new(),
204                |_, _, effects| effects.send(Ping(2)).map(|_| ()),
205            )
206            .expect("send");
207        world.end_run();
208        assert_eq!(
209            world.read_event(&mut reader).expect("read").map(|p| p.0),
210            Some(2)
211        );
212        assert!(world.read_event(&mut reader).expect("drain").is_none());
213    }
214
215    #[test]
216    fn send_maps_closed_channel_errors() {
217        use crate::component::ComponentOptions;
218        use crate::event::EventOptions;
219        use crate::operation::StageOperation;
220        use crate::world::WorldBuilder;
221
222        #[derive(Clone, Copy)]
223        struct Ping(#[allow(dead_code)] u8);
224
225        let mut builder = WorldBuilder::new();
226        builder
227            .register_component::<Ping>(ComponentOptions::sparse())
228            .expect("component");
229        builder
230            .add_event::<Ping>(EventOptions::frame(StageOperation::Update))
231            .expect("event");
232        let mut world = builder.build().expect("world");
233        let entity = world.spawn().expect("spawn");
234        world.insert(entity, Ping(1)).expect("insert");
235        world.set_event_sequence_for_test(2, 0, true);
236        world.begin_run(StageOperation::Update).expect("begin");
237        let err = world
238            .for_each_mut_with_effects::<Ping>(
239                &crate::query::QuerySpec::new(),
240                crate::query::QueryParams::new(),
241                |_, _, effects| effects.send(Ping(2)).map(|_| ()),
242            )
243            .expect_err("closed");
244        world.end_run();
245        assert!(matches!(err, QueryError::WrongQuery { .. }));
246    }
247
248    #[test]
249    fn map_allocator_error_query_covers_all_variants() {
250        assert!(matches!(
251            map_allocator_error_query(AllocatorError::GenerationOverflow),
252            QueryError::CommandRejected { .. }
253        ));
254        assert!(matches!(
255            map_allocator_error_query(AllocatorError::SlotRetired),
256            QueryError::CommandRejected { .. }
257        ));
258        assert!(matches!(
259            map_allocator_error_query(AllocatorError::StaleEntity),
260            QueryError::CommandRejected { .. }
261        ));
262    }
263
264    #[test]
265    fn query_commands_reject_stale_targets() {
266        use crate::component::ComponentOptions;
267        use crate::operation::StageOperation;
268        use crate::query::{QueryPolicy, QuerySpec, QueryWindow};
269        use crate::world::WorldBuilder;
270
271        #[derive(Clone, Copy)]
272        struct Marker;
273
274        let mut builder = WorldBuilder::new();
275        builder
276            .register_component::<Marker>(ComponentOptions::sparse())
277            .expect("marker");
278        let mut world = builder.build().expect("world");
279        let live = world.spawn().expect("live");
280        world.insert(live, Marker).expect("marker");
281        let stale = world.spawn().expect("stale");
282        world.despawn(stale).expect("despawn");
283        let mut query = world
284            .prepare_query1::<Marker>(QuerySpec::new(), QueryPolicy::Prepared)
285            .expect("prepare");
286
287        world.begin_run(StageOperation::Update).expect("begin");
288        let error = query
289            .for_each_mut_with_effects(&mut world, QueryWindow::All, |_, _, effects| {
290                effects.commands()?.despawn(stale)
291            })
292            .expect_err("stale command target");
293        world.end_run();
294        assert!(matches!(error, QueryError::CommandRejected { .. }));
295    }
296}
moirai/query/effects.rs

moirai/query/
effects.rs
