Build your own ECS in Rust, part 2: structural change and queries

main.rs

use std::collections::HashMap;

#[derive(Default, Copy, Clone, Debug, Eq, PartialEq, Hash)]
pub struct Entity {
    pub id: u32,
    pub generation: u32,
}

#[derive(Default)]
pub struct EntityAllocator {
    pub next_id: u32,
    pub free: Vec<(u32, u32)>,
}

impl EntityAllocator {
    pub fn allocate(&mut self) -> Entity {
        if let Some((id, generation)) = self.free.pop() {
            Entity { id, generation }
        } else {
            let id = self.next_id;
            self.next_id += 1;
            Entity { id, generation: 0 }
        }
    }

    pub fn deallocate(&mut self, entity: Entity) {
        self.free.push((entity.id, entity.generation.wrapping_add(1)));
    }
}

#[derive(Default, Copy, Clone)]
pub struct EntityLocation {
    pub generation: u32,
    pub table_index: u32,
    pub array_index: u32,
    pub allocated: bool,
}

#[derive(Default)]
pub struct EntityLocations {
    pub locations: Vec<EntityLocation>,
}

impl EntityLocations {
    pub fn ensure_slot(&mut self, id: u32) {
        let needed = id as usize + 1;
        if needed > self.locations.len() {
            self.locations.resize(needed, EntityLocation::default());
        }
    }

    pub fn get(&self, entity: Entity) -> Option<(usize, usize)> {
        let location = self.locations.get(entity.id as usize)?;
        if !location.allocated || location.generation != entity.generation {
            return None;
        }
        Some((location.table_index as usize, location.array_index as usize))
    }

    pub fn set(&mut self, entity: Entity, table_index: usize, array_index: usize) {
        self.ensure_slot(entity.id);
        self.locations[entity.id as usize] = EntityLocation {
            generation: entity.generation,
            table_index: table_index as u32,
            array_index: array_index as u32,
            allocated: true,
        };
    }

    pub fn mark_deallocated(&mut self, id: u32) {
        if let Some(location) = self.locations.get_mut(id as usize) {
            location.allocated = false;
        }
    }
}

#[derive(Default, Clone, Debug)]
pub struct Position {
    pub x: f32,
    pub y: f32,
}

#[derive(Default, Clone, Debug)]
pub struct Velocity {
    pub x: f32,
    pub y: f32,
}

pub const POSITION: u64 = 1 << 0;
pub const VELOCITY: u64 = 1 << 1;
pub const COMPONENT_COUNT: usize = 2;

pub fn component_index(mask: u64) -> Option<usize> {
    match mask {
        POSITION => Some(0),
        VELOCITY => Some(1),
        _ => None,
    }
}

#[derive(Default)]
pub struct ComponentArrays {
    pub mask: u64,
    pub entities: Vec<Entity>,
    pub positions: Vec<Position>,
    pub velocities: Vec<Velocity>,
}

#[derive(Default, Clone)]
pub struct TableEdges {
    pub add_edges: [Option<usize>; COMPONENT_COUNT],
    pub remove_edges: [Option<usize>; COMPONENT_COUNT],
}

#[derive(Default)]
pub struct World {
    pub allocator: EntityAllocator,
    pub entity_locations: EntityLocations,
    pub tables: Vec<ComponentArrays>,
    pub table_lookup: HashMap<u64, usize>,
    pub table_edges: Vec<TableEdges>,
    pub query_cache: HashMap<u64, Vec<usize>>,
}

impl World {
    pub fn invalidate_query_cache_for_new_table(
        &mut self,
        new_mask: u64,
        new_table_index: usize,
    ) {
        for (query_mask, cached_tables) in self.query_cache.iter_mut() {
            if new_mask & query_mask == *query_mask {
                cached_tables.push(new_table_index);
            }
        }
    }

    pub fn get_or_create_table(&mut self, mask: u64) -> usize {
        if let Some(&index) = self.table_lookup.get(&mask) {
            return index;
        }
        let new_table_index = self.tables.len();
        self.tables.push(ComponentArrays {
            mask,
            ..Default::default()
        });
        self.table_edges.push(TableEdges::default());
        self.table_lookup.insert(mask, new_table_index);

        self.invalidate_query_cache_for_new_table(mask, new_table_index);

        for component_bit in [POSITION, VELOCITY] {
            let Some(component_index_value) = component_index(component_bit) else {
                continue;
            };
            for (existing_index, existing) in self.tables.iter().enumerate() {
                if existing.mask | component_bit == mask {
                    self.table_edges[existing_index].add_edges[component_index_value] =
                        Some(new_table_index);
                }
                if existing.mask & !component_bit == mask {
                    self.table_edges[existing_index].remove_edges[component_index_value] =
                        Some(new_table_index);
                }
            }
        }

        new_table_index
    }

    pub fn cached_tables(&mut self, mask: u64) -> &[usize] {
        if !self.query_cache.contains_key(&mask) {
            let matching: Vec<usize> = self
                .tables
                .iter()
                .enumerate()
                .filter(|(_, table)| table.mask & mask == mask)
                .map(|(index, _)| index)
                .collect();
            self.query_cache.insert(mask, matching);
        }
        &self.query_cache[&mask]
    }

    pub fn spawn(&mut self, mask: u64) -> Entity {
        let entity = self.allocator.allocate();
        let table_index = self.get_or_create_table(mask);
        let table = &mut self.tables[table_index];
        let array_index = table.entities.len();

        table.entities.push(entity);
        if mask & POSITION != 0 {
            table.positions.push(Position::default());
        }
        if mask & VELOCITY != 0 {
            table.velocities.push(Velocity::default());
        }

        self.entity_locations.set(entity, table_index, array_index);
        entity
    }

    pub fn get_position(&self, entity: Entity) -> Option<&Position> {
        let (table_index, array_index) = self.entity_locations.get(entity)?;
        let table = &self.tables[table_index];
        if table.mask & POSITION == 0 {
            return None;
        }
        Some(&table.positions[array_index])
    }

    pub fn get_position_mut(&mut self, entity: Entity) -> Option<&mut Position> {
        let (table_index, array_index) = self.entity_locations.get(entity)?;
        let table = &mut self.tables[table_index];
        if table.mask & POSITION == 0 {
            return None;
        }
        Some(&mut table.positions[array_index])
    }

    pub fn get_velocity(&self, entity: Entity) -> Option<&Velocity> {
        let (table_index, array_index) = self.entity_locations.get(entity)?;
        let table = &self.tables[table_index];
        if table.mask & VELOCITY == 0 {
            return None;
        }
        Some(&table.velocities[array_index])
    }

    pub fn get_velocity_mut(&mut self, entity: Entity) -> Option<&mut Velocity> {
        let (table_index, array_index) = self.entity_locations.get(entity)?;
        let table = &mut self.tables[table_index];
        if table.mask & VELOCITY == 0 {
            return None;
        }
        Some(&mut table.velocities[array_index])
    }

    pub fn set_position(&mut self, entity: Entity, value: Position) {
        let Some((table_index, array_index)) = self.entity_locations.get(entity) else {
            return;
        };
        if self.tables[table_index].mask & POSITION != 0 {
            self.tables[table_index].positions[array_index] = value;
            return;
        }
        self.add_components(entity, POSITION);
        if let Some((table_index, array_index)) = self.entity_locations.get(entity) {
            self.tables[table_index].positions[array_index] = value;
        }
    }

    pub fn set_velocity(&mut self, entity: Entity, value: Velocity) {
        let Some((table_index, array_index)) = self.entity_locations.get(entity) else {
            return;
        };
        if self.tables[table_index].mask & VELOCITY != 0 {
            self.tables[table_index].velocities[array_index] = value;
            return;
        }
        self.add_components(entity, VELOCITY);
        if let Some((table_index, array_index)) = self.entity_locations.get(entity) {
            self.tables[table_index].velocities[array_index] = value;
        }
    }

    pub fn move_entity(
        &mut self,
        entity: Entity,
        from_table: usize,
        from_index: usize,
        to_table: usize,
    ) {
        let from_mask = self.tables[from_table].mask;

        let position = if from_mask & POSITION != 0 {
            Some(std::mem::take(
                &mut self.tables[from_table].positions[from_index],
            ))
        } else {
            None
        };
        let velocity = if from_mask & VELOCITY != 0 {
            Some(std::mem::take(
                &mut self.tables[from_table].velocities[from_index],
            ))
        } else {
            None
        };

        let to_array_index = {
            let to = &mut self.tables[to_table];
            let array_index = to.entities.len();
            to.entities.push(entity);
            if to.mask & POSITION != 0 {
                to.positions.push(position.unwrap_or_default());
            }
            if to.mask & VELOCITY != 0 {
                to.velocities.push(velocity.unwrap_or_default());
            }
            array_index
        };

        self.entity_locations.set(entity, to_table, to_array_index);

        let from = &mut self.tables[from_table];
        let last_index = from.entities.len() - 1;
        let swapped = if from_index < last_index {
            Some(from.entities[last_index])
        } else {
            None
        };

        from.entities.swap_remove(from_index);
        if from.mask & POSITION != 0 {
            from.positions.swap_remove(from_index);
        }
        if from.mask & VELOCITY != 0 {
            from.velocities.swap_remove(from_index);
        }

        if let Some(moved) = swapped {
            self.entity_locations.set(moved, from_table, from_index);
        }
    }

    pub fn add_components(&mut self, entity: Entity, mask: u64) -> bool {
        let Some((table_index, array_index)) = self.entity_locations.get(entity) else {
            return false;
        };
        let current_mask = self.tables[table_index].mask;
        if current_mask & mask == mask {
            return true;
        }

        let cached = if mask.count_ones() == 1 {
            component_index(mask)
                .and_then(|index| self.table_edges[table_index].add_edges[index])
        } else {
            None
        };

        let new_table_index = match cached {
            Some(index) => index,
            None => self.get_or_create_table(current_mask | mask),
        };

        self.move_entity(entity, table_index, array_index, new_table_index);
        true
    }

    pub fn remove_components(&mut self, entity: Entity, mask: u64) -> bool {
        let Some((table_index, array_index)) = self.entity_locations.get(entity) else {
            return false;
        };
        let current_mask = self.tables[table_index].mask;
        if current_mask & mask == 0 {
            return true;
        }

        let cached = if mask.count_ones() == 1 {
            component_index(mask)
                .and_then(|index| self.table_edges[table_index].remove_edges[index])
        } else {
            None
        };

        let new_table_index = match cached {
            Some(index) => index,
            None => self.get_or_create_table(current_mask & !mask),
        };

        self.move_entity(entity, table_index, array_index, new_table_index);
        true
    }

    pub fn despawn(&mut self, entity: Entity) -> bool {
        let Some((table_index, array_index)) = self.entity_locations.get(entity) else {
            return false;
        };

        self.entity_locations.mark_deallocated(entity.id);
        self.allocator.deallocate(entity);

        let table = &mut self.tables[table_index];
        let last_index = table.entities.len() - 1;
        let swapped = if array_index < last_index {
            Some(table.entities[last_index])
        } else {
            None
        };

        table.entities.swap_remove(array_index);
        if table.mask & POSITION != 0 {
            table.positions.swap_remove(array_index);
        }
        if table.mask & VELOCITY != 0 {
            table.velocities.swap_remove(array_index);
        }

        if let Some(moved) = swapped {
            self.entity_locations.set(moved, table_index, array_index);
        }

        true
    }

    pub fn query_entities(&self, mask: u64) -> impl Iterator<Item = Entity> + '_ {
        self.tables
            .iter()
            .filter(move |table| table.mask & mask == mask)
            .flat_map(|table| table.entities.iter().copied())
    }

    pub fn for_each<F>(&self, include: u64, exclude: u64, mut f: F)
    where
        F: FnMut(Entity, &ComponentArrays, usize),
    {
        for table in &self.tables {
            if table.mask & include != include || table.mask & exclude != 0 {
                continue;
            }
            for array_index in 0..table.entities.len() {
                let entity = table.entities[array_index];
                f(entity, table, array_index);
            }
        }
    }

    pub fn for_each_mut<F>(&mut self, include: u64, exclude: u64, mut f: F)
    where
        F: FnMut(Entity, &mut ComponentArrays, usize),
    {
        let table_indices: Vec<usize> = self.cached_tables(include).to_vec();
        for table_index in table_indices {
            let table = &mut self.tables[table_index];
            if table.mask & exclude != 0 {
                continue;
            }
            for array_index in 0..table.entities.len() {
                let entity = table.entities[array_index];
                f(entity, table, array_index);
            }
        }
    }
}

fn main() {
    let mut world = World::default();

    let mover = world.spawn(POSITION);
    world.set_position(mover, Position { x: 1.0, y: 2.0 });

    world.add_components(mover, VELOCITY);
    world.set_velocity(mover, Velocity { x: 0.5, y: 0.25 });

    let landmark = world.spawn(POSITION);
    world.set_position(landmark, Position { x: 100.0, y: 100.0 });

    for _ in 0..4 {
        world.for_each_mut(POSITION | VELOCITY, 0, |_entity, table, index| {
            table.positions[index].x += table.velocities[index].x;
            table.positions[index].y += table.velocities[index].y;
        });
    }

    let mover_position = world.get_position(mover).unwrap();
    println!("mover after 4 steps: ({}, {})", mover_position.x, mover_position.y);

    let landmark_position = world.get_position(landmark).unwrap();
    println!("landmark unchanged: ({}, {})", landmark_position.x, landmark_position.y);

    world.remove_components(mover, VELOCITY);
    assert!(world.get_velocity(mover).is_none());
    assert!(world.get_position(mover).is_some());

    for entity in world.query_entities(POSITION) {
        let position = world.get_position(entity).unwrap();
        println!("{entity:?} at ({}, {})", position.x, position.y);
    }

    println!("ok");
}