add more SAFETY comments and lint for missing ones in bevy_ecs (b…

…evyengine#4835)

# Objective

`SAFETY` comments are meant to be placed before `unsafe` blocks and should contain the reasoning of why in this case the usage of unsafe is okay. This is useful when reading the code because it makes it clear which assumptions are required for safety, and makes it easier to spot possible unsoundness holes. It also forces the code writer to think of something to write and maybe look at the safety contracts of any called unsafe methods again to double-check their correct usage.

There's a clippy lint called `undocumented_unsafe_blocks` which warns when using a block without such a comment. 

## Solution

- since clippy expects `SAFETY` instead of `SAFE`, rename those
- add `SAFETY` comments in more places
- for the last remaining 3 places, add an `#[allow()]` and `// TODO` since I wasn't comfortable enough with the code to justify their safety
- add ` #![warn(clippy::undocumented_unsafe_blocks)]` to `bevy_ecs`


### Note for reviewers

The first commit only renames `SAFETY` to `SAFE` so it doesn't need a thorough review.
https://github.com/bevyengine/bevy/pull/4835/files/cb042a416ecbe5e7d74797449969e064d8a5f13c..55cef2d6fa3aa634667a60f6d5abc16f43f16298 is the diff for all other changes.

### Safety comments where I'm not too familiar with the code

https://github.com/bevyengine/bevy/blob/774012ece50e4add4fcc8324ec48bbecf5546c3c/crates/bevy_ecs/src/entity/mod.rs#L540-L546

https://github.com/bevyengine/bevy/blob/774012ece50e4add4fcc8324ec48bbecf5546c3c/crates/bevy_ecs/src/world/entity_ref.rs#L249-L252

### Locations left undocumented with a `TODO` comment

https://github.com/bevyengine/bevy/blob/5dde944a3051426ac69fdedc5699f7da97a7e147/crates/bevy_ecs/src/schedule/executor_parallel.rs#L196-L199

https://github.com/bevyengine/bevy/blob/5dde944a3051426ac69fdedc5699f7da97a7e147/crates/bevy_ecs/src/world/entity_ref.rs#L287-L289

https://github.com/bevyengine/bevy/blob/5dde944a3051426ac69fdedc5699f7da97a7e147/crates/bevy_ecs/src/world/entity_ref.rs#L413-L415

Co-authored-by: Jakob Hellermann <[email protected]>

crates/bevy_ecs/macros/src/lib.rs

@@ -144,7 +144,7 @@ pub fn derive_bundle(input: TokenStream) -> TokenStream {
     let struct_name = &ast.ident;
 
     TokenStream::from(quote! {
-        /// SAFE: ComponentId is returned in field-definition-order. [from_components] and [get_components] use field-definition-order
+        /// SAFETY: ComponentId is returned in field-definition-order. [from_components] and [get_components] use field-definition-order
         unsafe impl #impl_generics #ecs_path::bundle::Bundle for #struct_name #ty_generics #where_clause {
             fn component_ids(
                 components: &mut #ecs_path::component::Components,
@@ -192,7 +192,7 @@ pub fn impl_param_set(_input: TokenStream) -> TokenStream {
         let index = Index::from(i);
         param_fn_muts.push(quote! {
             pub fn #fn_name<'a>(&'a mut self) -> <#param::Fetch as SystemParamFetch<'a, 'a>>::Item {
-                // SAFE: systems run without conflicts with other systems.
+                // SAFETY: systems run without conflicts with other systems.
                 // Conflicting params in ParamSet are not accessible at the same time
                 // ParamSets are guaranteed to not conflict with other SystemParams
                 unsafe {
@@ -213,13 +213,13 @@ pub fn impl_param_set(_input: TokenStream) -> TokenStream {
                 type Fetch = ParamSetState<(#(#param::Fetch,)*)>;
             }
 
-            // SAFE: All parameters are constrained to ReadOnlyFetch, so World is only read
+            // SAFETY: All parameters are constrained to ReadOnlyFetch, so World is only read
 
             unsafe impl<#(#param_fetch: for<'w1, 's1> SystemParamFetch<'w1, 's1>,)*> ReadOnlySystemParamFetch for ParamSetState<(#(#param_fetch,)*)>
             where #(#param_fetch: ReadOnlySystemParamFetch,)*
             { }
 
-            // SAFE: Relevant parameter ComponentId and ArchetypeComponentId access is applied to SystemMeta. If any ParamState conflicts
+            // SAFETY: Relevant parameter ComponentId and ArchetypeComponentId access is applied to SystemMeta. If any ParamState conflicts
             // with any prior access, a panic will occur.
 
             unsafe impl<#(#param_fetch: for<'w1, 's1> SystemParamFetch<'w1, 's1>,)*> SystemParamState for ParamSetState<(#(#param_fetch,)*)>

crates/bevy_ecs/src/archetype.rs

@@ -420,13 +420,13 @@ impl Archetypes {
 
     #[inline]
     pub fn empty(&self) -> &Archetype {
-        // SAFE: empty archetype always exists
+        // SAFETY: empty archetype always exists
         unsafe { self.archetypes.get_unchecked(ArchetypeId::EMPTY.index()) }
     }
 
     #[inline]
     pub(crate) fn empty_mut(&mut self) -> &mut Archetype {
-        // SAFE: empty archetype always exists
+        // SAFETY: empty archetype always exists
         unsafe {
             self.archetypes
                 .get_unchecked_mut(ArchetypeId::EMPTY.index())
@@ -435,13 +435,13 @@ impl Archetypes {
 
     #[inline]
     pub fn resource(&self) -> &Archetype {
-        // SAFE: resource archetype always exists
+        // SAFETY: resource archetype always exists
         unsafe { self.archetypes.get_unchecked(ArchetypeId::RESOURCE.index()) }
     }
 
     #[inline]
     pub(crate) fn resource_mut(&mut self) -> &mut Archetype {
-        // SAFE: resource archetype always exists
+        // SAFETY: resource archetype always exists
         unsafe {
             self.archetypes
                 .get_unchecked_mut(ArchetypeId::RESOURCE.index())

crates/bevy_ecs/src/bundle.rs

@@ -99,6 +99,10 @@ pub unsafe trait Bundle: Send + Sync + 'static {
 
 macro_rules! tuple_impl {
     ($($name: ident),*) => {
+        // SAFETY:
+        // - `Bundle::component_ids` returns the `ComponentId`s for each component type in the
+        // bundle, in the exact order that `Bundle::get_components` is called.
+        // - `Bundle::from_components` calls `func` exactly once for each `ComponentId` returned by `Bundle::component_ids`.
         unsafe impl<$($name: Component),*> Bundle for ($($name,)*) {
             #[allow(unused_variables)]
             fn component_ids(components: &mut Components, storages: &mut Storages) -> Vec<ComponentId> {
@@ -325,7 +329,7 @@ impl BundleInfo {
                 bundle_status.push(ComponentStatus::Mutated);
             } else {
                 bundle_status.push(ComponentStatus::Added);
-                // SAFE: component_id exists
+                // SAFETY: component_id exists
                 let component_info = unsafe { components.get_info_unchecked(component_id) };
                 match component_info.storage_type() {
                     StorageType::Table => new_table_components.push(component_id),
@@ -354,7 +358,7 @@ impl BundleInfo {
                     new_table_components.extend(current_archetype.table_components());
                     // sort to ignore order while hashing
                     new_table_components.sort();
-                    // SAFE: all component ids in `new_table_components` exist
+                    // SAFETY: all component ids in `new_table_components` exist
                     table_id = unsafe {
                         storages
                             .tables
@@ -492,7 +496,7 @@ impl<'a, 'b> BundleInserter<'a, 'b> {
                     } else if new_archetype.id() == swapped_location.archetype_id {
                         new_archetype
                     } else {
-                        // SAFE: the only two borrowed archetypes are above and we just did collision checks
+                        // SAFETY: the only two borrowed archetypes are above and we just did collision checks
                         &mut *self
                             .archetypes_ptr
                             .add(swapped_location.archetype_id.index())
@@ -567,7 +571,7 @@ impl<'a, 'b> BundleSpawner<'a, 'b> {
     #[inline]
     pub unsafe fn spawn<T: Bundle>(&mut self, bundle: T) -> Entity {
         let entity = self.entities.alloc();
-        // SAFE: entity is allocated (but non-existent), `T` matches this BundleInfo's type
+        // SAFETY: entity is allocated (but non-existent), `T` matches this BundleInfo's type
         self.spawn_non_existent(entity, bundle);
         entity
     }
@@ -599,14 +603,14 @@ impl Bundles {
         let id = self.bundle_ids.entry(TypeId::of::<T>()).or_insert_with(|| {
             let component_ids = T::component_ids(components, storages);
             let id = BundleId(bundle_infos.len());
-            // SAFE: T::component_id ensures info was created
+            // SAFETY: T::component_id ensures info was created
             let bundle_info = unsafe {
                 initialize_bundle(std::any::type_name::<T>(), component_ids, id, components)
             };
             bundle_infos.push(bundle_info);
             id
         });
-        // SAFE: index either exists, or was initialized
+        // SAFETY: index either exists, or was initialized
         unsafe { self.bundle_infos.get_unchecked(id.0) }
     }
 }
@@ -623,7 +627,7 @@ unsafe fn initialize_bundle(
     let mut storage_types = Vec::new();
 
     for &component_id in &component_ids {
-        // SAFE: component_id exists and is therefore valid
+        // SAFETY: component_id exists and is therefore valid
         let component_info = components.get_info_unchecked(component_id);
         storage_types.push(component_info.storage_type());
     }

crates/bevy_ecs/src/component.rs

@@ -474,7 +474,7 @@ impl Components {
 
     #[inline]
     pub fn init_resource<T: Resource>(&mut self) -> ComponentId {
-        // SAFE: The [`ComponentDescriptor`] matches the [`TypeId`]
+        // SAFETY: The [`ComponentDescriptor`] matches the [`TypeId`]
         unsafe {
             self.get_or_insert_resource_with(TypeId::of::<T>(), || {
                 ComponentDescriptor::new_resource::<T>()
@@ -484,7 +484,7 @@ impl Components {
 
     #[inline]
     pub fn init_non_send<T: Any>(&mut self) -> ComponentId {
-        // SAFE: The [`ComponentDescriptor`] matches the [`TypeId`]
+        // SAFETY: The [`ComponentDescriptor`] matches the [`TypeId`]
         unsafe {
             self.get_or_insert_resource_with(TypeId::of::<T>(), || {
                 ComponentDescriptor::new_non_send::<T>(StorageType::default())

crates/bevy_ecs/src/entity/mod.rs

@@ -591,6 +591,8 @@ impl Entities {
     // Flushes all reserved entities to an "invalid" state. Attempting to retrieve them will return None
     // unless they are later populated with a valid archetype.
     pub fn flush_as_invalid(&mut self) {
+        // SAFETY: as per `flush` safety docs, the archetype id can be set to [`ArchetypeId::INVALID`] if
+        // the [`Entity`] has not been assigned to an [`Archetype`][crate::archetype::Archetype], which is the case here
         unsafe {
             self.flush(|_entity, location| {
                 location.archetype_id = ArchetypeId::INVALID;
@@ -658,6 +660,7 @@ mod tests {
     fn reserve_entity_len() {
         let mut e = Entities::default();
         e.reserve_entity();
+        // SAFETY: entity_location is left invalid
         unsafe { e.flush(|_, _| {}) };
         assert_eq!(e.len(), 1);
     }
@@ -669,6 +672,7 @@ mod tests {
         assert!(entities.contains(e));
         assert!(entities.get(e).is_none());
 
+        // SAFETY: entity_location is left invalid
         unsafe {
             entities.flush(|_entity, _location| {
                 // do nothing ... leaving entity location invalid

crates/bevy_ecs/src/lib.rs

@@ -1,3 +1,4 @@
+#![warn(clippy::undocumented_unsafe_blocks)]
 #![doc = include_str!("../README.md")]
 
 #[cfg(target_pointer_width = "16")]

crates/bevy_ecs/src/query/filter.rs

@@ -480,7 +480,7 @@ macro_rules! impl_query_filter_tuple {
             }
         }
 
-        // SAFE: filters are read only
+        // SAFETY: filters are read only
         unsafe impl<$($filter: ReadOnlyWorldQuery),*> ReadOnlyWorldQuery for Or<($($filter,)*)> {}
     };
 }

crates/bevy_ecs/src/query/iter.rs

@@ -52,6 +52,9 @@ where
 
     #[inline(always)]
     fn next(&mut self) -> Option<Self::Item> {
+        // SAFETY:
+        // `tables` and `archetypes` belong to the same world that the cursor was initialized for.
+        // `query_state` is the state that was passed to `QueryIterationCursor::init`.
         unsafe {
             self.cursor
                 .next(self.tables, self.archetypes, self.query_state)
@@ -150,33 +153,42 @@ where
 
     #[inline(always)]
     fn next(&mut self) -> Option<Self::Item> {
-        unsafe {
-            for entity in self.entity_iter.by_ref() {
-                let location = match self.entities.get(*entity.borrow()) {
-                    Some(location) => location,
-                    None => continue,
-                };
-
-                if !self
-                    .query_state
-                    .matched_archetypes
-                    .contains(location.archetype_id.index())
-                {
-                    continue;
-                }
+        for entity in self.entity_iter.by_ref() {
+            let location = match self.entities.get(*entity.borrow()) {
+                Some(location) => location,
+                None => continue,
+            };
+
+            if !self
+                .query_state
+                .matched_archetypes
+                .contains(location.archetype_id.index())
+            {
+                continue;
+            }
 
-                let archetype = &self.archetypes[location.archetype_id];
+            let archetype = &self.archetypes[location.archetype_id];
 
+            // SAFETY: `archetype` is from the world that `fetch/filter` were created for,
+            // `fetch_state`/`filter_state` are the states that `fetch/filter` were initialized with
+            unsafe {
                 self.fetch
                     .set_archetype(&self.query_state.fetch_state, archetype, self.tables);
+            }
+            // SAFETY: `table` is from the world that `fetch/filter` were created for,
+            // `fetch_state`/`filter_state` are the states that `fetch/filter` were initialized with
+            unsafe {
                 self.filter
                     .set_archetype(&self.query_state.filter_state, archetype, self.tables);
-                if self.filter.archetype_filter_fetch(location.index) {
-                    return Some(self.fetch.archetype_fetch(location.index));
-                }
             }
-            None
+            // SAFETY: set_archetype was called prior.
+            // `location.index` is an archetype index row in range of the current archetype, because if it was not, the match above would have `continue`d
+            if unsafe { self.filter.archetype_filter_fetch(location.index) } {
+                // SAFETY: set_archetype was called prior, `location.index` is an archetype index in range of the current archetype
+                return Some(unsafe { self.fetch.archetype_fetch(location.index) });
+            }
         }
+        None
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
@@ -289,7 +301,7 @@ impl<'w, 's, Q: WorldQuery, F: WorldQuery, const K: usize> QueryCombinationIter<
         for<'a> QueryFetch<'a, Q>: Clone,
         for<'a> QueryFetch<'a, F>: Clone,
     {
-        // safety: we are limiting the returned reference to self,
+        // SAFETY: we are limiting the returned reference to self,
         // making sure this method cannot be called multiple times without getting rid
         // of any previously returned unique references first, thus preventing aliasing.
         unsafe {
@@ -374,7 +386,9 @@ struct QueryIterationCursor<'w, 's, Q: WorldQuery, QF: Fetch<'w, State = Q::Stat
     archetype_id_iter: std::slice::Iter<'s, ArchetypeId>,
     fetch: QF,
     filter: QueryFetch<'w, F>,
+    // length of the table table or length of the archetype, depending on whether both `Q`'s and `F`'s fetches are dense
     current_len: usize,
+    // either table row or archetype index, depending on whether both `Q`'s and `F`'s fetches are dense
     current_index: usize,
     phantom: PhantomData<(&'w (), Q)>,
 }
@@ -461,6 +475,10 @@ where
 
     // NOTE: If you are changing query iteration code, remember to update the following places, where relevant:
     // QueryIterationCursor, QueryState::for_each_unchecked_manual, QueryState::par_for_each_unchecked_manual
+    /// # Safety
+    /// `tables` and `archetypes` must belong to the same world that the [`QueryIterationCursor`]
+    /// was initialized for.
+    /// `query_state` must be the same [`QueryState`] that was passed to `init` or `init_empty`.
     #[inline(always)]
     unsafe fn next(
         &mut self,
@@ -470,21 +488,28 @@ where
     ) -> Option<QF::Item> {
         if Self::IS_DENSE {
             loop {
+                // we are on the beginning of the query, or finished processing a table, so skip to the next
                 if self.current_index == self.current_len {
                     let table_id = self.table_id_iter.next()?;
                     let table = &tables[*table_id];
+                    // SAFETY: `table` is from the world that `fetch/filter` were created for,
+                    // `fetch_state`/`filter_state` are the states that `fetch/filter` were initialized with
                     self.fetch.set_table(&query_state.fetch_state, table);
                     self.filter.set_table(&query_state.filter_state, table);
                     self.current_len = table.len();
                     self.current_index = 0;
                     continue;
                 }
 
+                // SAFETY: set_table was called prior.
+                // `current_index` is a table row in range of the current table, because if it was not, then the if above would have been executed.
                 if !self.filter.table_filter_fetch(self.current_index) {
                     self.current_index += 1;
                     continue;
                 }
 
+                // SAFETY: set_table was called prior.
+                // `current_index` is a table row in range of the current table, because if it was not, then the if above would have been executed.
                 let item = self.fetch.table_fetch(self.current_index);
 
                 self.current_index += 1;
@@ -495,6 +520,8 @@ where
                 if self.current_index == self.current_len {
                     let archetype_id = self.archetype_id_iter.next()?;
                     let archetype = &archetypes[*archetype_id];
+                    // SAFETY: `archetype` and `tables` are from the world that `fetch/filter` were created for,
+                    // `fetch_state`/`filter_state` are the states that `fetch/filter` were initialized with
                     self.fetch
                         .set_archetype(&query_state.fetch_state, archetype, tables);
                     self.filter
@@ -504,11 +531,15 @@ where
                     continue;
                 }
 
+                // SAFETY: set_archetype was called prior.
+                // `current_index` is an archetype index row in range of the current archetype, because if it was not, then the if above would have been executed.
                 if !self.filter.archetype_filter_fetch(self.current_index) {
                     self.current_index += 1;
                     continue;
                 }
 
+                // SAFETY: set_archetype was called prior, `current_index` is an archetype index in range of the current archetype
+                // `current_index` is an archetype index row in range of the current archetype, because if it was not, then the if above would have been executed.
                 let item = self.fetch.archetype_fetch(self.current_index);
                 self.current_index += 1;
                 return Some(item);