use std::any::type_name; 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::{ Component, IsOptional as ComponentIsOptional, Metadata as ComponentMetadata, }; use crate::type_name::TypeName; use crate::uid::Uid; use crate::util::Sortable; use crate::EntityComponent; #[derive(Debug, Default)] pub struct Storage { archetypes: Vec, archetype_lookup: RefCell>, entity_archetype_lookup: HashMap, } impl Storage { pub fn find_entities( &self, mut components_metadata: CompMetadataList, ) -> ArchetypeRefIter<'_> where CompMetadataList: Sortable, CompMetadataList: AsRef<[ComponentMetadata]>, { components_metadata.sort_by_key_b(|component_metadata| component_metadata.id); let archetype_id = ArchetypeId::from_components_metadata(&components_metadata); // This looks stupid but the borrow checker complains otherwise if self.archetype_lookup.borrow().contains_key(&archetype_id) { return self.iter_archetypes_by_lookup(archetype_id); } let comp_ids_set = create_non_opt_component_id_set(components_metadata.as_ref()); let matching_archetype_indices = self .archetypes .iter() .enumerate() .filter_map(|(index, archetype)| { if archetype.component_ids_is_superset(&comp_ids_set) { return Some(index); } None }) .collect(); self.archetype_lookup.borrow_mut().insert( archetype_id, ArchetypeLookupEntry { component_ids: comp_ids_set, archetype_indices: matching_archetype_indices, }, ); self.iter_archetypes_by_lookup(archetype_id) } pub fn get_entity_archetype(&self, entity_uid: Uid) -> Option<&Archetype> { let archetype_id = self.entity_archetype_lookup.get(&entity_uid)?; let archetype_index = self.find_archetype_index_with_entity(*archetype_id, entity_uid)?; self.archetypes.get(archetype_index) } pub fn remove_entity(&mut self, entity_uid: Uid) { let Some(archetype_id) = self.entity_archetype_lookup.get(&entity_uid) else { return; }; let Some(archetype_index) = self.find_archetype_index_with_entity(*archetype_id, entity_uid) else { return; }; let Some(archetype) = self.archetypes.get_mut(archetype_index) else { return; }; archetype.take_entity(entity_uid); self.entity_archetype_lookup.remove(&entity_uid); } #[cfg_attr(feature = "debug", tracing::instrument(skip_all))] pub fn push_entity( &mut self, entity_uid: Uid, mut components: Vec>, ) -> Result<(ArchetypeId, Uid), Error> { if self.entity_archetype_lookup.contains_key(&entity_uid) { return Err(Error::EntityAlreadyExists(entity_uid)); } components.sort_by_key(|component| component.self_id()); #[cfg(feature = "debug")] tracing::debug!( "Pushing entity with components: ({})", components .iter() .map(|component| component.type_name()) .collect::>() .join(", ") ); let archetype_id = ArchetypeId::from_components_metadata( &components .iter() .map(|component| ComponentMetadata::get(&**component)) .collect::>(), ); let comp_ids_set = create_non_opt_component_id_set( components .iter() .map(|component| ComponentMetadata::get(&**component)), ); let archetype_index = self.get_or_create_archetype(archetype_id, &comp_ids_set, &components); self.populate_matching_archetype_lookup_entries(&comp_ids_set, archetype_index); let archetype = self .archetypes .get_mut(archetype_index) .expect("Archetype is gone"); archetype.push_entity(entity_uid, components); archetype .entity_lookup .insert(entity_uid, archetype.entities.len() - 1); self.entity_archetype_lookup .insert(entity_uid, archetype_id); Ok((archetype_id, entity_uid)) } pub fn add_components_to_entity( &mut self, entity_uid: Uid, components: Vec>, ) -> Option<()> { let archetype_id = self.entity_archetype_lookup.get(&entity_uid)?; let archetype_index = self.find_archetype_index_with_entity(*archetype_id, entity_uid)?; let archetype = self.archetypes.get_mut(archetype_index)?; let contains_component_already = components .iter() .any(|component| archetype.component_ids.contains_key(&component.self_id())); if contains_component_already { let component_cnt = components.len(); // TODO: return error panic!( "Entity with UID {:?} already contains one or more component(s) ({})", entity_uid, components .iter() .flat_map(|component| [component.type_name(), ", "]) .enumerate() .take_while(|(index, _)| { *index < (component_cnt * 2) - 1 }) .map(|(_, component)| component) .collect::() ); } let entity = archetype.take_entity(entity_uid)?; self.entity_archetype_lookup.remove(&entity_uid); self.push_entity( entity_uid, entity .components .into_iter() .map(|component| component.component.into_inner()) .chain(components) .collect(), ) .expect("Not supposed to return Err since the entity is removed"); Some(()) } pub fn remove_components_from_entity( &mut self, entity_uid: Uid, component_ids: impl IntoIterator, ) -> Option<()> { let archetype_id = self.entity_archetype_lookup.get(&entity_uid)?; let archetype_index = self.find_archetype_index_with_entity(*archetype_id, entity_uid)?; let archetype = self.archetypes.get_mut(archetype_index)?; let entity = archetype.take_entity(entity_uid)?; let component_ids_set = component_ids.into_iter().collect::>(); self.entity_archetype_lookup.remove(&entity_uid); self.push_entity( entity_uid, entity .components .into_iter() .map(|component| component.component.into_inner()) .filter(|component| !component_ids_set.contains(&component.self_id())) .collect(), ) .expect("Not supposed to return Err since the entity is removed"); Some(()) } fn populate_matching_archetype_lookup_entries( &mut self, comp_ids_set: &HashSet, archetype_index: usize, ) { let mut archetype_lookup = self.archetype_lookup.borrow_mut(); for (_, lookup_entry) in archetype_lookup.iter_mut() { if &lookup_entry.component_ids == comp_ids_set { continue; } // There shouldn't be duplicate archetype indices in the lookup entry if lookup_entry.archetype_indices.contains(&archetype_index) { continue; } if lookup_entry.component_ids.is_subset(comp_ids_set) { lookup_entry.archetype_indices.push(archetype_index); } } } fn get_or_create_archetype( &mut self, archetype_id: ArchetypeId, comp_ids_set: &HashSet, components: &[Box], ) -> usize { let mut archetype_lookup = self.archetype_lookup.borrow_mut(); let lookup_entry = archetype_lookup.entry(archetype_id).or_insert_with(|| { ArchetypeLookupEntry { component_ids: comp_ids_set.clone(), archetype_indices: Vec::with_capacity(1), } }); if lookup_entry.archetype_indices.is_empty() { self.archetypes.push(Archetype::new( components.iter().map(|component| component.self_id()), )); lookup_entry .archetype_indices .push(self.archetypes.len() - 1); } // SAFETY: Above, we push a archetype index if archetype_indices is empty so this // cannot fail unsafe { *lookup_entry.archetype_indices.first().unwrap_unchecked() } } fn find_archetype_index_with_entity( &self, archetype_id: ArchetypeId, entity_uid: Uid, ) -> Option { let archetype_lookup = self.archetype_lookup.borrow_mut(); let archetype_lookup_entry = archetype_lookup.get(&archetype_id)?; // TODO: Also have a hashmap for precise archetype ID -> archetype index lookup. // This way is slow archetype_lookup_entry .archetype_indices .iter() .find(|archetype_index| { let Some(archetype) = self.archetypes.get(**archetype_index) else { return false; }; archetype.has_entity(entity_uid) }) .copied() } 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(); ArchetypeRefIter { indices: archetype_indices.into_iter(), archetypes: &self.archetypes, } } } /// Component storage error #[derive(Debug, Clone, thiserror::Error)] pub enum Error { #[error("Entity already exists")] EntityAlreadyExists(Uid), } impl TypeName for Storage { fn type_name(&self) -> &'static str { type_name::() } } #[derive(Debug)] struct ArchetypeLookupEntry { component_ids: HashSet, archetype_indices: Vec, } #[derive(Debug)] pub struct Archetype { component_ids: HashMap, entity_lookup: HashMap, entities: Vec, } impl Archetype { fn new(component_ids: impl IntoIterator) -> Self { Self { component_ids: component_ids .into_iter() .enumerate() .map(|(index, component_id)| (component_id, index)) .collect(), entity_lookup: HashMap::new(), entities: Vec::new(), } } pub fn component_ids_is_superset(&self, other_component_ids: &HashSet) -> bool { if other_component_ids.len() <= self.component_ids.len() { other_component_ids .iter() .all(|v| self.component_ids.contains_key(v)) } else { false } } pub fn get_entity(&self, entity_uid: Uid) -> Option<&ArchetypeEntity> { let entity_index = *self.entity_lookup.get(&entity_uid)?; self.entities.get(entity_index) } pub fn entities(&self) -> EntityIter<'_> { EntityIter { iter: self.entities.iter() } } pub fn get_index_for_component(&self, component_id: Uid) -> Option { self.component_ids.get(&component_id).copied() } fn push_entity( &mut self, entity_uid: Uid, components: impl IntoIterator>, ) { self.entities.push(ArchetypeEntity { uid: entity_uid, components: components.into_iter().map(Into::into).collect(), }); } pub fn take_entity(&mut self, entity_uid: Uid) -> Option { let entity_index = self.entity_lookup.remove(&entity_uid)?; let entity = self.entities.remove(entity_index); for index in self.entity_lookup.values_mut() { if *index > entity_index { *index -= 1; } } Some(entity) } fn has_entity(&self, entity_uid: Uid) -> bool { self.entity_lookup.contains_key(&entity_uid) } } #[derive(Debug)] pub struct ArchetypeEntity { uid: Uid, components: Vec, } impl ArchetypeEntity { pub fn uid(&self) -> Uid { self.uid } pub fn components(&self) -> &[EntityComponent] { &self.components } pub fn get_component(&self, index: usize) -> Option<&EntityComponent> { self.components.get(index) } } #[derive(Debug)] pub struct ArchetypeRefIter<'component_storage> { indices: OwnedVecIter, archetypes: &'component_storage [Archetype], } impl<'component_storage> Iterator for ArchetypeRefIter<'component_storage> { type Item = &'component_storage Archetype; fn next(&mut self) -> Option { let archetype_index = self.indices.next()?; Some( self.archetypes .get(archetype_index) .expect("Archetype index in archetype lookup entry was not found"), ) } } #[derive(Debug)] pub struct EntityIter<'archetype> { iter: SliceIter<'archetype, ArchetypeEntity>, } impl<'archetype> Iterator for EntityIter<'archetype> { type Item = &'archetype ArchetypeEntity; fn next(&mut self) -> Option { self.iter.next() } } fn create_non_opt_component_id_set( component_metadata_iter: impl IntoIterator, ) -> HashSet where Item: Borrow, { component_metadata_iter .into_iter() .filter_map(|item| { let component_metadata = item.borrow(); if component_metadata.is_optional == ComponentIsOptional::Yes { return None; } Some(component_metadata.id) }) .collect::>() } #[cfg(test)] mod tests { use ecs_macros::Component; use super::Storage; use crate::archetype::Id as ArchetypeId; use crate::component::{ Component, IsOptional as ComponentIsOptional, Metadata as ComponentMetadata, }; use crate::uid::{Kind as UidKind, Uid}; #[derive(Debug, Component)] struct HealthPotion { _hp_restoration: u32, } #[derive(Debug, Component)] struct Hookshot { _range: u32, } #[derive(Debug, Component)] struct DekuNut { _throwing_damage: u32, } #[derive(Debug, Component)] struct Bow { _damage: u32, } #[derive(Debug, Component)] struct IronBoots; #[test] fn push_entity_works() { let mut component_storage = Storage::default(); component_storage .push_entity( Uid::new_unique(UidKind::Entity), vec![ Box::new(HealthPotion { _hp_restoration: 12 }), Box::new(Hookshot { _range: 50 }), ], ) .expect("Expected Ok"); assert_eq!(component_storage.archetypes.len(), 1); let archetype = component_storage .archetypes .first() .expect("Expected a archetype in archetypes Vec"); assert_eq!(archetype.component_ids.len(), 2); // One entity assert_eq!(archetype.entities.len(), 1); let entity_components = archetype .entities .first() .expect("Expected a entity in archetype"); assert_eq!(entity_components.components.len(), 2); assert_eq!(component_storage.archetype_lookup.borrow().len(), 1); let mut components_metadata = [ ComponentMetadata { id: HealthPotion::id(), is_optional: ComponentIsOptional::No, }, ComponentMetadata { id: Hookshot::id(), is_optional: ComponentIsOptional::No, }, ]; components_metadata.sort_by_key(|comp_metadata| comp_metadata.id); let archetype_lookup = component_storage.archetype_lookup.borrow(); let lookup_entry = archetype_lookup .get(&ArchetypeId::from_components_metadata(&components_metadata)) .expect("Expected entry in archetype lookup map"); let first_archetype_index = lookup_entry .archetype_indices .first() .expect("Expected archetype lookup to contain a archetype reference"); assert_eq!(*first_archetype_index, 0); } }