use crate::ecs::{
component::{Component, ComponentId},
entity::Entity,
query::{QueryItem, QueryItem2, QueryItemMut},
storage::table::TableComponentKey,
};
use super::{Archetype, TableComponentValue, TableRowLocation};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct TableEntityLocation {
pub(crate) archetype: usize,
pub(crate) row: TableRowLocation,
}
pub(crate) struct TableComponentRemoval {
pub(crate) removed_entity: Entity,
pub(crate) new_location: Option<TableEntityLocation>,
pub(crate) moved_entity: Option<Entity>,
}
pub(crate) struct TableComponentInsertion {
pub(crate) entity: Entity,
pub(crate) new_location: TableEntityLocation,
pub(crate) moved_entity: Option<Entity>,
}
#[derive(Default)]
pub(crate) struct TableStorage {
archetypes: Vec<Archetype>,
chunk_capacity: usize,
}
impl TableStorage {
#[must_use]
pub(crate) fn new(chunk_capacity: usize) -> Self {
assert!(
chunk_capacity > 0,
"table chunk capacity must be greater than zero"
);
Self {
archetypes: Vec::new(),
chunk_capacity,
}
}
pub(crate) fn query<T: Component>(&self, component_id: ComponentId) -> Vec<QueryItem<'_, T>> {
let mut results = Vec::new();
for archetype in &self.archetypes {
if !archetype.has_component(component_id) {
continue;
}
for chunk in archetype.chunks() {
for row in 0..chunk.row_count() {
let Some(entity) = chunk.entity(row) else {
continue;
};
let Some(component) = chunk.get::<T>(component_id, row) else {
continue;
};
results.push(QueryItem { entity, component });
}
}
}
results
}
pub(crate) fn query2<A: Component, B: Component>(
&self,
first_id: ComponentId,
second_id: ComponentId,
) -> Vec<QueryItem2<'_, A, B>> {
let mut results = Vec::new();
for archetype in &self.archetypes {
if !archetype.has_component(first_id) || !archetype.has_component(second_id) {
continue;
}
for chunk in archetype.chunks() {
for row in 0..chunk.row_count() {
let Some(entity) = chunk.entity(row) else {
continue;
};
let Some(first) = chunk.get::<A>(first_id, row) else {
continue;
};
let Some(second) = chunk.get::<B>(second_id, row) else {
continue;
};
results.push(QueryItem2 {
entity,
first,
second,
});
}
}
}
results
}
pub(crate) fn query_mut<T: Component>(
&mut self,
component_id: ComponentId,
) -> Vec<QueryItemMut<'_, T>> {
let mut results = Vec::new();
for archetype in &mut self.archetypes {
if !archetype.has_component(component_id) {
continue;
}
results.extend(archetype.query_mut(component_id));
}
results
}
pub(crate) fn get<T: Component>(
&self,
location: TableEntityLocation,
component_id: ComponentId,
) -> Option<&T> {
self.archetypes
.get(location.archetype)?
.get(component_id, location.row)
}
pub(crate) fn get_mut<T: Component>(
&mut self,
location: TableEntityLocation,
component_id: ComponentId,
) -> Option<&mut T> {
self.archetypes
.get_mut(location.archetype)?
.get_mut(component_id, location.row)
}
pub(crate) fn insert(
&mut self,
entity: Entity,
components: Vec<TableComponentValue>,
keys: Vec<TableComponentKey>,
) -> TableEntityLocation {
let mut keys = keys;
keys.sort_by_key(|key| key.order());
let archetype = self.find_or_create_archetype(keys);
let row = self.archetypes[archetype].push_row(
entity,
components
.into_iter()
.map(TableComponentValue::into_parts)
.collect(),
);
TableEntityLocation { archetype, row }
}
pub(crate) fn insert_component(
&mut self,
location: TableEntityLocation,
component: TableComponentValue,
key: TableComponentKey,
) -> Option<TableComponentInsertion> {
let component_id = component.id();
if self
.archetypes
.get(location.archetype)?
.contains_component(component_id)
{
self.archetypes[location.archetype].replace(
component_id,
location.row,
component.into_parts().1,
)?;
return Some(TableComponentInsertion {
entity: self.archetypes[location.archetype].entity(location.row)?,
new_location: location,
moved_entity: None,
});
}
let destination_keys = self.archetypes[location.archetype].component_keys_with(key);
let removed_row = self.archetypes[location.archetype].take_row(location.row);
let mut components = removed_row
.components
.into_iter()
.map(|(id, value)| TableComponentValue::from_erased(id, value))
.collect::<Vec<_>>();
components.push(component);
let new_location = self.insert(removed_row.entity, components, destination_keys);
Some(TableComponentInsertion {
entity: removed_row.entity,
new_location,
moved_entity: removed_row.moved_entity,
})
}
pub(crate) fn remove(&mut self, location: TableEntityLocation) -> Option<Entity> {
let archetype = self
.archetypes
.get_mut(location.archetype)
.expect("table entity location should reference an existing archetype");
archetype.remove_row(location.row)
}
pub(crate) fn remove_component(
&mut self,
location: TableEntityLocation,
component_id: ComponentId,
) -> Option<TableComponentRemoval> {
let source = self.archetypes.get(location.archetype)?;
if !source.contains_component(component_id) {
return None;
}
let destination_keys = source.component_keys_without(component_id);
let removed_row = self.archetypes[location.archetype].take_row(location.row);
let remaining_components = removed_row
.components
.into_iter()
.filter(|(id, _)| *id != component_id)
.map(|(id, value)| TableComponentValue::from_erased(id, value))
.collect::<Vec<_>>();
let new_location = if destination_keys.is_empty() {
None
} else {
Some(self.insert(removed_row.entity, remaining_components, destination_keys))
};
Some(TableComponentRemoval {
removed_entity: removed_row.entity,
new_location,
moved_entity: removed_row.moved_entity,
})
}
pub(crate) fn archetype_count(&self) -> usize {
self.archetypes.len()
}
pub(crate) fn len(&self) -> usize {
self.archetypes.iter().map(Archetype::len).sum()
}
fn find_or_create_archetype(&mut self, keys: Vec<TableComponentKey>) -> usize {
if let Some(index) = self
.archetypes
.iter()
.position(|archetype| archetype.components() == keys)
{
return index;
}
let index = self.archetypes.len();
self.archetypes
.push(Archetype::new(keys, self.chunk_capacity));
index
}
}