1 files changed, 145 insertions, 0 deletions
diff --git a/src/data/entity_arena.rs b/src/data/entity_arena.rs
new file mode 100644
index 0000000..ad6e451
--- /dev/null
+++ b/src/data/entity_arena.rs
@@ -0,0 +1,145 @@
+use crate::data::arena::Index;
+use crate::data::Arena;
+use crate::data::Coarena;
+use bevy::prelude::Entity;
+use std::ops::IndexMut;
+
+impl From<Entity> for Index {
+    #[inline]
+    fn from(value: Entity) -> Self {
+        Index::from_raw_parts(value.index(), value.generation())
+    }
+}
+
+impl From<Index> for Entity {
+    #[inline]
+    fn from(value: Index) -> Self {
+        let (id, gen) = value.into_raw_parts();
+        Entity::from_bits(id as u64 | ((gen as u64) << u32::BITS))
+    }
+}
+
+#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
+#[derive(Clone, Debug, Default)]
+struct EntityMap {
+    entity_to_index: Coarena<Index>,
+    index_to_entity: Coarena<Entity>, // Is this really needed?
+}
+
+impl EntityMap {
+    pub fn get_index(&self, entity: Entity) -> Option<Index> {
+        self.entity_to_index.get(entity.into()).copied()
+    }
+
+    pub fn get_entity(&self, index: Index) -> Option<Entity> {
+        self.index_to_entity.get(index).copied()
+    }
+
+    pub fn insert(&mut self, entity: Entity, index: Index) {
+        self.entity_to_index.insert(entity.into(), index);
+        self.index_to_entity
+            .insert_with_default(index, entity, Entity::PLACEHOLDER);
+    }
+
+    pub fn remove(&mut self, entity: Entity) {
+        if let Some(index) = self.entity_to_index.remove(
+            entity.into(),
+            Index::from_raw_parts(crate::INVALID_U32, crate::INVALID_U32),
+        ) {
+            self.index_to_entity.remove(index, Entity::PLACEHOLDER);
+        }
+    }
+}
+
+impl std::ops::Index<Index> for EntityMap {
+    type Output = Entity;
+
+    fn index(&self, index: Index) -> &Self::Output {
+        self.index_to_entity.get(index).expect("Index not found.")
+    }
+}
+
+impl std::ops::Index<Entity> for EntityMap {
+    type Output = Index;
+
+    fn index(&self, entity: Entity) -> &Self::Output {
+        self.entity_to_index
+            .get(entity.into())
+            .expect("Index not found.")
+    }
+}
+
+#[derive(Clone, Debug)]
+#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
+pub struct EntityArena<T> {
+    entity_map: EntityMap,
+    elements: Arena<T>,
+}
+
+impl<T> Default for EntityArena<T> {
+    fn default() -> Self {
+        Self {
+            entity_map: EntityMap::default(),
+            elements: Arena::default(),
+        }
+    }
+}
+
+impl<T> EntityArena<T> {
+    pub fn insert(&mut self, entity: Entity, value: T) {
+        let index = self.elements.insert(value);
+        self.entity_map.insert(entity, index);
+    }
+
+    pub fn remove(&mut self, entity: Entity) -> Option<T> {
+        let elt = self.elements.remove(self.entity_map.get_index(entity)?)?;
+        self.entity_map.remove(entity);
+        Some(elt)
+    }
+
+    pub fn contains(&self, entity: Entity) -> bool {
+        self.entity_map.get_index(entity).is_some()
+    }
+
+    pub fn len(&self) -> usize {
+        self.elements.len()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.elements.is_empty()
+    }
+
+    pub fn get(&self, entity: Entity) -> Option<&T> {
+        self.elements.get(self.entity_map.get_index(entity)?)
+    }
+
+    pub fn get_mut(&mut self, entity: Entity) -> Option<&mut T> {
+        self.elements.get_mut(self.entity_map.get_index(entity)?)
+    }
+
+    pub fn iter(&self) -> impl ExactSizeIterator<Item = (Entity, &T)> {
+        self.elements
+            .iter()
+            .map(|(h, data)| (self.entity_map[h], data))
+    }
+
+    pub fn iter_mut(&mut self) -> impl ExactSizeIterator<Item = (Entity, &mut T)> {
+        self.elements
+            .iter_mut()
+            .map(|(h, data)| (self.entity_map[h], data))
+    }
+}
+
+impl<T> std::ops::Index<Entity> for EntityArena<T> {
+    type Output = T;
+
+    fn index(&self, index: Entity) -> &Self::Output {
+        &self.elements[self.entity_map[index]]
+    }
+}
+
+impl<T> IndexMut<Entity> for EntityArena<T> {
+    fn index_mut(&mut self, index: Entity) -> &mut Self::Output {
+        &mut self.elements[self.entity_map[index]]
+    }
+}