fix(ecs): prevent nested queries causing panic

2 files changed, 28 insertions, 81 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
diff --git a/ecs/src/util.rs b/ecs/src/util.rs
index 5ec21c7..f4f8632 100644
--- a/ecs/src/util.rs
+++ b/ecs/src/util.rs
@@ -1,6 +1,3 @@
-use std::cell::Ref;
-use std::ops::{Deref, DerefMut};
-
 pub trait Sortable
 {
     type Item;
@@ -49,50 +46,3 @@ impl<Item> Sortable for Vec<Item>
         self.sort_by_key(func)
     }
 }
-
-#[derive(Debug)]
-pub struct RefCellRefMap<'a, Value: 'a, MappedValue: 'a>
-{
-    mapped_value: MappedValue,
-    _refcell_ref: Ref<'a, Value>,
-}
-
-impl<'a, Value: 'a, MappedValue: 'a> RefCellRefMap<'a, Value, MappedValue>
-{
-    pub fn new(
-        refcell_ref: Ref<'a, Value>,
-        map_fn: impl Fn(&'a Value) -> MappedValue,
-    ) -> Self
-    {
-        // Convert the lifetime to 'static. This is necessary to make the compiler not
-        // complain about refcell_ref being moved
-        //
-        // SAFETY: This is fine since we pass it to map_fn with the original lifetime 'a
-        let val_ref = unsafe { &*(&*refcell_ref as *const Value) };
-
-        let mapped_value = map_fn(val_ref);
-
-        Self {
-            mapped_value,
-            _refcell_ref: refcell_ref,
-        }
-    }
-}
-
-impl<'a, Value: 'a, MappedValue: 'a> Deref for RefCellRefMap<'a, Value, MappedValue>
-{
-    type Target = MappedValue;
-
-    fn deref(&self) -> &Self::Target
-    {
-        &self.mapped_value
-    }
-}
-
-impl<'a, Value: 'a, MappedValue: 'a> DerefMut for RefCellRefMap<'a, Value, MappedValue>
-{
-    fn deref_mut(&mut self) -> &mut Self::Target
-    {
-        &mut self.mapped_value
-    }
-}