Use QueryIterator in SequentiallyIterable implementation

martinjonsson01 · Apr 1, 2023 · 9c4fe47 · 9c4fe47
1 parent 8f12e7c
commit 9c4fe47
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 37 deletions.
diff --git a/crates/ecs/src/systems/iteration.rs b/crates/ecs/src/systems/iteration.rs
@@ -1,7 +1,7 @@
 //! Different ways to iterate over all queried entities in a [`System`].
 
 use super::*;
-use crate::{intersection_of_multiple_sets, World};
+use crate::World;
 
 /// Execution of a single [`System`] in a sequential order.
 pub trait SequentiallyIterable {
@@ -31,40 +31,16 @@ macro_rules! impl_sequentially_iterable_system {
             {
 
                 fn run(&self, world: &World) -> SystemResult<()> {
-                    let base_signature: Vec<TypeId> = [$([<P$parameter>]::base_signature(),)*]
-                        .into_iter()
-                        .flatten()
-                        .collect();
-
-                    let universe = world.get_archetype_indices(&base_signature);
-
-                    let archetypes_indices: Vec<_> = intersection_of_multiple_sets(&[
-                        universe.clone(),
-                        $([<P$parameter>]::filter(&universe, world),)*
-                    ])
-                    .into_iter()
-                    .collect();
-
-                    $(let mut [<borrowed_$parameter>] = [<P$parameter>]::borrow(world, &archetypes_indices).map_err(SystemError::MissingParameter)?;)*
-
-                    // SAFETY: This is safe because the result from fetch_parameter will not outlive borrowed
-                    unsafe {
-                        if $([<P$parameter>]::iterates_over_entities() ||)* false {
-                            while let ($(Some([<parameter_$parameter>]),)*) = (
-                                $([<P$parameter>]::fetch_parameter(&mut [<borrowed_$parameter>]),)*
-                            ) {
-                                if let ($(Some([<parameter_$parameter>]),)*) = (
-                                    $([<parameter_$parameter>],)*
-                                ) {
-                                    (self.function)($([<parameter_$parameter>],)*);
-                                }
-                            }
-                        } else if let ($(Some(Some([<parameter_$parameter>])),)*) = (
-                                $([<P$parameter>]::fetch_parameter(&mut [<borrowed_$parameter>]),)*
-                        ) {
-                            (self.function)($([<parameter_$parameter>],)*);
-                        }
+                    let query: Query<($([<P$parameter>],)*)> = Query {
+                        phantom: PhantomData::default(),
+                        world
+                    };
+
+                    let query_iterator = query.try_into_iter().map_err(SystemError::MissingParameter)?;
+                    for ($([<parameter_$parameter>],)*) in query_iterator {
+                        (self.function)($([<parameter_$parameter>],)*);
                     }
+
                     Ok(())
                 }
             }

diff --git a/crates/ecs/src/systems/mod.rs b/crates/ecs/src/systems/mod.rs
@@ -519,6 +519,13 @@ macro_rules! impl_system_parameter_function {
                 type IntoIter = QueryIterator<'a, ($([<P$parameter>],)*)>;
 
                 fn into_iter(self) -> Self::IntoIter {
+                    Self::try_into_iter(self)
+                        .expect("creating `QueryIterator` should work if the archetypes are in a valid state")
+                }
+            }
+
+            impl<'a, $([<P$parameter>]: SystemParameter,)*> Query<'a, ($([<P$parameter>],)*)> {
+                fn try_into_iter(self) -> SystemParameterResult<QueryIterator<'a, ($([<P$parameter>],)*)>> {
                     let base_signature: Vec<TypeId> = [$([<P$parameter>]::base_signature(),)*]
                         .into_iter()
                         .flatten()
@@ -534,16 +541,16 @@ macro_rules! impl_system_parameter_function {
                     .collect();
 
                     let borrowed = (
-                        $([<P$parameter>]::borrow(self.world, &archetypes_indices).unwrap(),)*
+                        $([<P$parameter>]::borrow(self.world, &archetypes_indices)?,)*
                     );
 
                     let iterate_over_entities = $([<P$parameter>]::iterates_over_entities() ||)* false;
 
-                    Self::IntoIter {
+                    Ok(QueryIterator {
                         borrowed,
                         iterate_over_entities,
                         iterated_once: false,
-                    }
+                    })
                 }
             }