fix(ecs): prevent nested queries causing panic

1 files changed, 28 insertions, 31 deletions

diff --git a/ecs/src/component/storage.rs b/ecs/src/component/storage.rs
index f8b5b20..5d9bb06 100644
--- a/ecs/src/component/storage.rs
+++ b/ecs/src/component/storage.rs
@@ -3,6 +3,7 @@ use std::borrow::Borrow;
 use std::cell::RefCell;
 use std::collections::{HashMap, HashSet};
 use std::slice::Iter as SliceIter;
+use std::vec::IntoIter as OwnedVecIter;
 
 use crate::archetype::Id as ArchetypeId;
 use crate::component::{
@@ -13,7 +14,7 @@ use crate::component::{
 };
 use crate::entity::Uid as EntityUid;
 use crate::type_name::TypeName;
-use crate::util::{RefCellRefMap, Sortable};
+use crate::util::Sortable;
 use crate::EntityComponent;
 
 #[derive(Debug, Default)]
@@ -42,18 +43,7 @@ impl Storage
 
         // This looks stupid but the borrow checker complains otherwise
         if self.archetype_lookup.borrow().contains_key(&archetype_id) {
-            let archetype_lookup = self.archetype_lookup.borrow();
-
-            return ArchetypeRefIter {
-                inner: RefCellRefMap::new(archetype_lookup, |archetype_lookup| {
-                    archetype_lookup
-                        .get(&archetype_id)
-                        .unwrap()
-                        .archetype_indices
-                        .iter()
-                }),
-                archetypes: &self.archetypes,
-            };
+            return self.iter_archetypes_by_lookup(&archetype_id);
         }
 
         let comp_ids_set = create_non_opt_component_id_set(components_metadata.borrow());
@@ -79,18 +69,7 @@ impl Storage
             },
         );
 
-        let archetype_lookup = self.archetype_lookup.borrow();
-
-        ArchetypeRefIter {
-            inner: RefCellRefMap::new(archetype_lookup, |archetype_lookup| {
-                archetype_lookup
-                    .get(&archetype_id)
-                    .unwrap()
-                    .archetype_indices
-                    .iter()
-            }),
-            archetypes: &self.archetypes,
-        }
+        self.iter_archetypes_by_lookup(&archetype_id)
     }
 
     #[cfg_attr(feature = "debug", tracing::instrument(skip_all))]
@@ -306,6 +285,28 @@ impl Storage
             })
             .map(|archetype_index| *archetype_index)
     }
+
+    fn iter_archetypes_by_lookup(
+        &self,
+        archetype_id: &ArchetypeId,
+    ) -> ArchetypeRefIter<'_>
+    {
+        let archetype_lookup = self.archetype_lookup.borrow();
+
+        // The archetype indices have to be cloned to prevent a panic when query
+        // iterations are nested. The panic happens because the archetype_lookup RefCell
+        // is borrowed and it tries to mutably borrow it
+        let archetype_indices = archetype_lookup
+            .get(archetype_id)
+            .unwrap()
+            .archetype_indices
+            .clone();
+
+        return ArchetypeRefIter {
+            indices: archetype_indices.into_iter(),
+            archetypes: &self.archetypes,
+        };
+    }
 }
 
 impl TypeName for Storage
@@ -441,11 +442,7 @@ impl ArchetypeEntity
 #[derive(Debug)]
 pub struct ArchetypeRefIter<'component_storage>
 {
-    inner: RefCellRefMap<
-        'component_storage,
-        HashMap<ArchetypeId, ArchetypeLookupEntry>,
-        SliceIter<'component_storage, usize>,
-    >,
+    indices: OwnedVecIter<usize>,
     archetypes: &'component_storage [Archetype],
 }
 
@@ -455,7 +452,7 @@ impl<'component_storage> Iterator for ArchetypeRefIter<'component_storage>
 
     fn next(&mut self) -> Option<Self::Item>
     {
-        let archetype_index = *self.inner.next()?;
+        let archetype_index = self.indices.next()?;
 
         Some(
             self.archetypes