From 1b89195d53691bdf0cf50a3e58d22e8949e373be Mon Sep 17 00:00:00 2001
From: Victoron <59878206+Victoronz@users.noreply.github.com>
Date: Fri, 20 Dec 2024 20:34:52 +0100
Subject: [PATCH] more QueryIter EntitySetIterator impls

---
 crates/bevy_ecs/src/query/iter.rs | 36 ++++++++++++++++++++++++++++++-
 1 file changed, 35 insertions(+), 1 deletion(-)

diff --git a/crates/bevy_ecs/src/query/iter.rs b/crates/bevy_ecs/src/query/iter.rs
index 816c92e76ca14..c7303382f68d8 100644
--- a/crates/bevy_ecs/src/query/iter.rs
+++ b/crates/bevy_ecs/src/query/iter.rs
@@ -1,11 +1,15 @@
 use super::{QueryData, QueryFilter, ReadOnlyQueryData};
 use crate::{
     archetype::{Archetype, ArchetypeEntity, Archetypes},
+    bundle::Bundle,
     component::Tick,
     entity::{Entities, Entity, EntityBorrow, EntitySet, EntitySetIterator},
     query::{ArchetypeFilter, DebugCheckedUnwrap, QueryState, StorageId},
     storage::{Table, TableRow, Tables},
-    world::unsafe_world_cell::UnsafeWorldCell,
+    world::{
+        unsafe_world_cell::UnsafeWorldCell, EntityMut, EntityMutExcept, EntityRef, EntityRefExcept,
+        FilteredEntityMut, FilteredEntityRef,
+    },
 };
 use alloc::vec::Vec;
 use core::{
@@ -1105,6 +1109,36 @@ impl<'w, 's, D: QueryData, F: QueryFilter> FusedIterator for QueryIter<'w, 's, D
 // SAFETY: [`QueryIter`] is guaranteed to return every matching entity once and only once.
 unsafe impl<'w, 's, F: QueryFilter> EntitySetIterator for QueryIter<'w, 's, Entity, F> {}
 
+// SAFETY: [`QueryIter`] is guaranteed to return every matching entity once and only once.
+unsafe impl<'w, 's, F: QueryFilter> EntitySetIterator for QueryIter<'w, 's, EntityRef<'_>, F> {}
+
+// SAFETY: [`QueryIter`] is guaranteed to return every matching entity once and only once.
+unsafe impl<'w, 's, F: QueryFilter> EntitySetIterator for QueryIter<'w, 's, EntityMut<'_>, F> {}
+
+// SAFETY: [`QueryIter`] is guaranteed to return every matching entity once and only once.
+unsafe impl<'w, 's, F: QueryFilter> EntitySetIterator
+    for QueryIter<'w, 's, FilteredEntityRef<'_>, F>
+{
+}
+
+// SAFETY: [`QueryIter`] is guaranteed to return every matching entity once and only once.
+unsafe impl<'w, 's, F: QueryFilter> EntitySetIterator
+    for QueryIter<'w, 's, FilteredEntityMut<'_>, F>
+{
+}
+
+// SAFETY: [`QueryIter`] is guaranteed to return every matching entity once and only once.
+unsafe impl<'w, 's, F: QueryFilter, B: Bundle> EntitySetIterator
+    for QueryIter<'w, 's, EntityRefExcept<'_, B>, F>
+{
+}
+
+// SAFETY: [`QueryIter`] is guaranteed to return every matching entity once and only once.
+unsafe impl<'w, 's, F: QueryFilter, B: Bundle> EntitySetIterator
+    for QueryIter<'w, 's, EntityMutExcept<'_, B>, F>
+{
+}
+
 impl<'w, 's, D: QueryData, F: QueryFilter> Debug for QueryIter<'w, 's, D, F> {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         f.debug_struct("QueryIter").finish()