Pass target location through System::run. (#14)

Prior to this change, receivers used the event's declared target entity to fetch their query data. This is unsound because the event's target could be changed via mutable receivers. Bad implementations of `Event` were also possible.
rj00a · Feb 13, 2024 · b452eb5 · b452eb5
1 parent f25d44f
commit b452eb5
Show file tree

Hide file tree

Showing 8 changed files with 158 additions and 61 deletions.
diff --git a/evenio_macros/src/system_param.rs b/evenio_macros/src/system_param.rs
@@ -61,6 +61,7 @@ pub(crate) fn derive_system_param(input: TokenStream) -> Result<TokenStream> {
                             state,
                             info,
                             event_ptr,
+                            target_location,
                             world
                         );
 
@@ -81,6 +82,7 @@ pub(crate) fn derive_system_param(input: TokenStream) -> Result<TokenStream> {
                             state,
                             info,
                             event_ptr,
+                            target_location,
                             world
                         );
 
@@ -130,6 +132,7 @@ pub(crate) fn derive_system_param(input: TokenStream) -> Result<TokenStream> {
                 state: &'__a mut Self::State,
                 info: &'__a ::evenio::system::SystemInfo,
                 event_ptr: ::evenio::event::EventPtr<'__a>,
+                target_location: ::evenio::entity::EntityLocation,
                 world: ::evenio::world::UnsafeWorldCell<'__a>,
             ) -> Self::Item<'__a> {
                 #get_body

diff --git a/src/archetype.rs b/src/archetype.rs
@@ -489,6 +489,7 @@ unsafe impl SystemParam for &'_ Archetypes {
         _state: &'a mut Self::State,
         _info: &'a SystemInfo,
         _event_ptr: EventPtr<'a>,
+        _target_location: EntityLocation,
         world: UnsafeWorldCell<'a>,
     ) -> Self::Item<'a> {
         world.archetypes()
@@ -529,6 +530,11 @@ unsafe impl SparseIndex for ArchetypeIdx {
 #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
 pub struct ArchetypeRow(pub u32);
 
+impl ArchetypeRow {
+    /// The archetype row that is always invalid.
+    pub const NULL: Self = Self(u32::MAX);
+}
+
 /// A container for all entities with a particular set of components.
 ///
 /// Each component has a corresponding columnm of contiguous data in

diff --git a/src/component.rs b/src/component.rs
@@ -11,6 +11,7 @@ pub use evenio_macros::Component;
 use crate::archetype::Archetype;
 use crate::assert::UnwrapDebugChecked;
 use crate::drop::DropFn;
+use crate::entity::EntityLocation;
 use crate::event::{Event, EventId, EventPtr};
 use crate::map::{Entry, TypeIdMap};
 use crate::prelude::World;
@@ -179,6 +180,7 @@ unsafe impl SystemParam for &'_ Components {
         _state: &'a mut Self::State,
         _info: &'a SystemInfo,
         _event_ptr: EventPtr<'a>,
+        _target_location: EntityLocation,
         world: UnsafeWorldCell<'a>,
     ) -> Self::Item<'a> {
         world.components()

diff --git a/src/entity.rs b/src/entity.rs
@@ -118,6 +118,7 @@ unsafe impl SystemParam for &'_ Entities {
         _state: &'a mut Self::State,
         _info: &'a SystemInfo,
         _event_ptr: EventPtr<'a>,
+        _target_location: EntityLocation,
         world: UnsafeWorldCell<'a>,
     ) -> Self::Item<'a> {
         world.entities()
@@ -137,6 +138,14 @@ pub struct EntityLocation {
     pub row: ArchetypeRow,
 }
 
+impl EntityLocation {
+    /// A location which is always invalid.
+    pub(crate) const NULL: Self = Self {
+        archetype: ArchetypeIdx::NULL,
+        row: ArchetypeRow::NULL,
+    };
+}
+
 /// Lightweight identifier for an entity.
 ///
 /// Entity identifiers are implemented using an [index] and a generation count.

diff --git a/src/event.rs b/src/event.rs
@@ -24,7 +24,7 @@ use crate::assert::{
 };
 use crate::component::ComponentIdx;
 use crate::drop::DropFn;
-use crate::entity::EntityId;
+use crate::entity::{EntityId, EntityLocation};
 use crate::fetch::FetcherState;
 use crate::map::{Entry, TypeIdMap};
 use crate::prelude::Component;
@@ -222,6 +222,7 @@ unsafe impl SystemParam for &'_ Events {
         _state: &'a mut Self::State,
         _info: &'a SystemInfo,
         _event_ptr: EventPtr<'a>,
+        _target_location: EntityLocation,
         world: UnsafeWorldCell<'a>,
     ) -> Self::Item<'a> {
         world.events()
@@ -869,6 +870,7 @@ unsafe impl<E: Event> SystemParam for Receiver<'_, E> {
         _state: &'a mut Self::State,
         _info: &'a SystemInfo,
         event_ptr: EventPtr<'a>,
+        _target_location: EntityLocation,
         _world: UnsafeWorldCell<'a>,
     ) -> Self::Item<'a> {
         Receiver {
@@ -921,16 +923,15 @@ unsafe impl<E: Event, Q: Query + 'static> SystemParam for Receiver<'_, E, Q> {
         state: &'a mut Self::State,
         _info: &'a SystemInfo,
         event_ptr: EventPtr<'a>,
-        world: UnsafeWorldCell<'a>,
+        target_location: EntityLocation,
+        _world: UnsafeWorldCell<'a>,
     ) -> Self::Item<'a> {
         assert!(E::IS_TARGETED);
+
         let event = event_ptr.as_ptr().cast::<E>().as_ref();
-        let target = event.target();
 
-        // SAFETY: The target entity is guaranteed to match the query.
-        let query = state
-            .get_mut(world.entities(), target)
-            .unwrap_debug_checked();
+        // SAFETY: Caller guarantees the target entity matches the query.
+        let query = state.get_by_location_mut(target_location);
 
         Receiver { event, query }
     }
@@ -988,6 +989,7 @@ unsafe impl<E: Event> SystemParam for ReceiverMut<'_, E> {
         _state: &'a mut Self::State,
         _info: &'a SystemInfo,
         event_ptr: EventPtr<'a>,
+        _target_location: EntityLocation,
         _world: UnsafeWorldCell<'a>,
     ) -> Self::Item<'a> {
         ReceiverMut {
@@ -1038,17 +1040,15 @@ unsafe impl<E: Event, Q: Query + 'static> SystemParam for ReceiverMut<'_, E, Q>
         state: &'a mut Self::State,
         _info: &'a SystemInfo,
         event_ptr: EventPtr<'a>,
-        world: UnsafeWorldCell<'a>,
+        target_location: EntityLocation,
+        _world: UnsafeWorldCell<'a>,
     ) -> Self::Item<'a> {
         assert!(E::IS_TARGETED);
 
         let event = EventMut::<E>::new(event_ptr);
-        let target = event.target();
 
-        // SAFETY: The target entity is guaranteed to match the query.
-        let query = state
-            .get_mut(world.entities(), target)
-            .unwrap_debug_checked();
+        // SAFETY: Caller guarantees the target entity matches the query.
+        let query = state.get_by_location_mut(target_location);
 
         ReceiverMut { event, query }
     }
@@ -1305,6 +1305,7 @@ unsafe impl<T: EventSet> SystemParam for Sender<'_, T> {
         state: &'a mut Self::State,
         _info: &'a SystemInfo,
         _event_ptr: EventPtr<'a>,
+        _target_location: EntityLocation,
         world: UnsafeWorldCell<'a>,
     ) -> Self::Item<'a> {
         Sender { state, world }
@@ -1378,14 +1379,14 @@ unsafe impl<E: Event> EventSet for E {
 
 macro_rules! impl_event_set_tuple {
     ($(($E:ident, $e:ident)),*) => {
-        #[allow(clippy::unused_unit)]
+        #[allow(unused_variables, clippy::unused_unit)]
         unsafe impl<$($E: EventSet),*> EventSet for ($($E,)*) {
             type EventIndices = ($($E::EventIndices,)*);
 
-            fn new_state(_world: &mut World) -> Self::EventIndices {
+            fn new_state(world: &mut World) -> Self::EventIndices {
                 (
                     $(
-                        $E::new_state(_world),
+                        $E::new_state(world),
                     )*
                 )
             }
@@ -1401,9 +1402,9 @@ macro_rules! impl_event_set_tuple {
                 None
             }
 
-            fn for_each_idx<F: FnMut(EventIdx)>(($($e,)*): &Self::EventIndices, mut _f: F) {
+            fn for_each_idx<F: FnMut(EventIdx)>(($($e,)*): &Self::EventIndices, #[allow(unused_mut)] mut f: F) {
                 $(
-                    $E::for_each_idx($e, &mut _f);
+                    $E::for_each_idx($e, &mut f);
                 )*
             }
         }
@@ -1656,4 +1657,26 @@ mod tests {
         let Result(values) = world.get_component_mut::<Result>(e).unwrap();
         assert_eq!(values, &[1, 2, 3]);
     }
+
+    #[test]
+    fn change_event_target() {
+        let mut world = World::new();
+
+        #[derive(Event)]
+        struct E(#[event(target)] EntityId);
+
+        #[derive(Component)]
+        struct C;
+
+        world.add_system(|mut r: ReceiverMut<E, &C>| {
+            r.event.0 = EntityId::NULL;
+        });
+
+        world.add_system(|_: Receiver<E, &C>| {});
+
+        let e = world.spawn();
+        world.insert(e, C);
+
+        world.send(E(e));
+    }
 }
diff --git a/src/fetch.rs b/src/fetch.rs
@@ -7,7 +7,7 @@ use core::{any, fmt};
 
 use crate::archetype::{Archetype, ArchetypeIdx, ArchetypeRow, Archetypes};
 use crate::assert::{assume_debug_checked, UnwrapDebugChecked};
-use crate::entity::{Entities, EntityId};
+use crate::entity::{Entities, EntityId, EntityLocation};
 use crate::event::EventPtr;
 use crate::query::{Query, ReadOnlyQuery};
 use crate::sparse_map::SparseMap;
@@ -64,7 +64,7 @@ impl<Q: Query> FetcherState<Q> {
             return Err(GetError::NoSuchEntity);
         };
 
-        // Eliminate a panicking branch.
+        // Eliminate a branch in `SparseMap::get`.
         assume_debug_checked(loc.archetype != ArchetypeIdx::NULL);
 
         let Some(state) = self.map.get(loc.archetype) else {
@@ -97,6 +97,15 @@ impl<Q: Query> FetcherState<Q> {
         Ok(Q::get(state, loc.row))
     }
 
+    #[inline]
+    pub(crate) unsafe fn get_by_location_mut(&mut self, loc: EntityLocation) -> Q::Item<'_> {
+        let state = self
+            .map
+            .get(loc.archetype)
+            .expect_debug_checked("invalid entity location");
+        Q::get(state, loc.row)
+    }
+
     // TODO: get_many_mut
 
     #[inline]
@@ -316,6 +325,7 @@ where
         state: &'a mut Self::State,
         _info: &'a SystemInfo,
         _event_ptr: EventPtr<'a>,
+        _target_location: EntityLocation,
         world: UnsafeWorldCell<'a>,
     ) -> Self::Item<'a> {
         Fetcher { state, world }
@@ -375,9 +385,10 @@ unsafe impl<Q: Query + 'static> SystemParam for Single<'_, Q> {
         state: &'a mut Self::State,
         info: &'a SystemInfo,
         event_ptr: EventPtr<'a>,
+        target_location: EntityLocation,
         world: UnsafeWorldCell<'a>,
     ) -> Self::Item<'a> {
-        match TrySingle::get(state, info, event_ptr, world) {
+        match TrySingle::get(state, info, event_ptr, target_location, world) {
             TrySingle(Ok(item)) => Single(item),
             TrySingle(Err(e)) => {
                 panic!(
@@ -417,6 +428,7 @@ unsafe impl<Q: Query + 'static> SystemParam for TrySingle<'_, Q> {
         state: &'a mut Self::State,
         _info: &'a SystemInfo,
         _event_ptr: EventPtr<'a>,
+        _target_location: EntityLocation,
         world: UnsafeWorldCell<'a>,
     ) -> Self::Item<'a> {
         let mut it = state.iter_mut(world.archetypes());