moirai-for-games 0.1.0

//! Query spec resolution into fingerprints, filters, and traversal drivers.

use alloc::format;
use alloc::string::String;
use alloc::vec::Vec;
use core::any::{type_name, TypeId};
use core::hash::{Hash, Hasher};

use crate::component::ComponentId;
use crate::query::{ExactIdPolicy, QueryError, QuerySpec};
use crate::world::World;

use super::plan::{ResolvedPlan, TraversalSource};
use super::plan_cache::QueryResolveScratch;

struct PreparedQuery1 {
    fingerprint: u64,
    primary_index: usize,
    primary_is_table: bool,
    traversal: TraversalSource,
    added_indices: Vec<usize>,
    changed_indices: Vec<usize>,
    exact_id_policy: Option<ExactIdPolicy>,
}

struct PreparedQuery2 {
    fingerprint: u64,
    primary_index: usize,
    primary_is_table: bool,
    second_index: usize,
    second_is_table: bool,
    traversal: TraversalSource,
    added_indices: Vec<usize>,
    changed_indices: Vec<usize>,
    exact_id_policy: Option<ExactIdPolicy>,
}

pub(crate) fn peek_query1_fingerprint<T: 'static>(
    world: &World,
    spec: &QuerySpec,
    scratch: &mut QueryResolveScratch,
) -> Result<u64, QueryError> {
    Ok(prepare_query1::<T>(world, spec, scratch)?.fingerprint)
}

pub(crate) fn resolve_query1<T: 'static>(
    world: &World,
    spec: &QuerySpec,
    scratch: &mut QueryResolveScratch,
) -> Result<ResolvedPlan, QueryError> {
    let prepared = prepare_query1::<T>(world, spec, scratch)?;
    Ok(ResolvedPlan {
        fingerprint: prepared.fingerprint,
        primary_index: prepared.primary_index,
        primary_is_table: prepared.primary_is_table,
        traversal: prepared.traversal,
        required_indices: scratch.required.clone(),
        without_indices: scratch.without.clone(),
        with_tag_indices: scratch.with_tags.clone(),
        without_tag_indices: scratch.without_tags.clone(),
        added_indices: prepared.added_indices,
        changed_indices: prepared.changed_indices,
        exact_id_policy: prepared.exact_id_policy,
    })
}

fn prepare_query1<T: 'static>(
    world: &World,
    spec: &QuerySpec,
    scratch: &mut QueryResolveScratch,
) -> Result<PreparedQuery1, QueryError> {
    let primary = resolve_component::<T>(world)?;
    let primary_index = primary.index();
    let primary_is_table = world.registry_is_table(&primary);

    let (added_indices, changed_indices) = fill_spec_indices(world, spec, scratch)?;
    if !scratch.required.contains(&primary_index) {
        scratch.required.push(primary_index);
    }
    normalize(&mut scratch.required);
    validate_overlaps(scratch)?;

    let traversal = if let Some(ids) = &spec.exact_ids {
        TraversalSource::Exact { ids: ids.clone() }
    } else if primary_is_table {
        TraversalSource::Table {
            component_index: primary_index,
        }
    } else {
        TraversalSource::Sparse {
            component_index: primary_index,
        }
    };

    let fingerprint = fingerprint_plan(
        &scratch.required,
        &scratch.without,
        &scratch.with_tags,
        &scratch.without_tags,
        &added_indices,
        &changed_indices,
        &traversal,
        Some(primary_index),
        None,
        spec.exact_id_policy,
    );

    Ok(PreparedQuery1 {
        fingerprint,
        primary_index,
        primary_is_table,
        traversal,
        added_indices,
        changed_indices,
        exact_id_policy: spec.exact_id_policy,
    })
}

pub(crate) fn resolve_query2<A: 'static, B: 'static>(
    world: &World,
    spec: &QuerySpec,
    scratch: &mut QueryResolveScratch,
) -> Result<(ResolvedPlan, usize, bool), QueryError> {
    let prepared = prepare_query2::<A, B>(world, spec, scratch)?;
    let plan = ResolvedPlan {
        fingerprint: prepared.fingerprint,
        primary_index: prepared.primary_index,
        primary_is_table: prepared.primary_is_table,
        traversal: prepared.traversal,
        required_indices: scratch.required.clone(),
        without_indices: scratch.without.clone(),
        with_tag_indices: scratch.with_tags.clone(),
        without_tag_indices: scratch.without_tags.clone(),
        added_indices: prepared.added_indices,
        changed_indices: prepared.changed_indices,
        exact_id_policy: prepared.exact_id_policy,
    };

    Ok((plan, prepared.second_index, prepared.second_is_table))
}

pub(crate) fn peek_query2_fingerprint<A: 'static, B: 'static>(
    world: &World,
    spec: &QuerySpec,
    scratch: &mut QueryResolveScratch,
) -> Result<(u64, usize, bool), QueryError> {
    let prepared = prepare_query2::<A, B>(world, spec, scratch)?;
    Ok((
        prepared.fingerprint,
        prepared.second_index,
        prepared.second_is_table,
    ))
}

fn prepare_query2<A: 'static, B: 'static>(
    world: &World,
    spec: &QuerySpec,
    scratch: &mut QueryResolveScratch,
) -> Result<PreparedQuery2, QueryError> {
    let primary_a = resolve_component::<A>(world)?;
    let primary_b = resolve_component::<B>(world)?;
    let second_index = primary_b.index();
    let second_is_table = world.registry_is_table(&primary_b);

    let (added_indices, changed_indices) = fill_spec_indices(world, spec, scratch)?;
    for index in [primary_a.index(), primary_b.index()] {
        if !scratch.required.contains(&index) {
            scratch.required.push(index);
        }
    }
    normalize(&mut scratch.required);
    validate_overlaps(scratch)?;

    let primary_index = primary_a.index();
    let primary_is_table = world.registry_is_table(&primary_a);

    let traversal = if let Some(ids) = &spec.exact_ids {
        TraversalSource::Exact { ids: ids.clone() }
    } else {
        let primary_len = world.query_component_population(primary_index, primary_is_table);
        let second_len = world.query_component_population(second_index, second_is_table);
        let (driver_index, driver_is_table) = if second_len < primary_len {
            (second_index, second_is_table)
        } else {
            (primary_index, primary_is_table)
        };
        if driver_is_table {
            TraversalSource::Table {
                component_index: driver_index,
            }
        } else {
            TraversalSource::Sparse {
                component_index: driver_index,
            }
        }
    };

    let fingerprint = fingerprint_plan(
        &scratch.required,
        &scratch.without,
        &scratch.with_tags,
        &scratch.without_tags,
        &added_indices,
        &changed_indices,
        &traversal,
        Some(primary_index),
        Some(second_index),
        spec.exact_id_policy,
    );

    Ok(PreparedQuery2 {
        fingerprint,
        primary_index,
        primary_is_table,
        traversal,
        added_indices,
        changed_indices,
        exact_id_policy: spec.exact_id_policy,
        second_index,
        second_is_table,
    })
}

pub(crate) fn peek_entities_fingerprint(
    world: &World,
    spec: &QuerySpec,
    scratch: &mut QueryResolveScratch,
) -> Result<u64, QueryError> {
    Ok(resolve_entities(world, spec, scratch)?.fingerprint)
}

pub(crate) fn resolve_entities(
    world: &World,
    spec: &QuerySpec,
    scratch: &mut QueryResolveScratch,
) -> Result<ResolvedPlan, QueryError> {
    let (added_indices, changed_indices) = fill_spec_indices(world, spec, scratch)?;
    validate_overlaps(scratch)?;

    let traversal = if let Some(ids) = &spec.exact_ids {
        TraversalSource::Exact { ids: ids.clone() }
    } else if let Some(&index) = scratch.required.first() {
        traversal_for_index(world, index)
    } else if let Some(&index) = scratch.with_tags.first() {
        TraversalSource::Sparse {
            component_index: index,
        }
    } else {
        TraversalSource::All
    };
    let fingerprint = fingerprint_plan(
        &scratch.required,
        &scratch.without,
        &scratch.with_tags,
        &scratch.without_tags,
        &added_indices,
        &changed_indices,
        &traversal,
        None,
        None,
        spec.exact_id_policy,
    );
    Ok(ResolvedPlan {
        fingerprint,
        primary_index: usize::MAX,
        primary_is_table: false,
        traversal,
        required_indices: scratch.required.clone(),
        without_indices: scratch.without.clone(),
        with_tag_indices: scratch.with_tags.clone(),
        without_tag_indices: scratch.without_tags.clone(),
        added_indices,
        changed_indices,
        exact_id_policy: spec.exact_id_policy,
    })
}

fn traversal_for_index(world: &World, index: usize) -> TraversalSource {
    let id = ComponentId::new(world.owner_token(), index as u32);
    if world.registry_is_table(&id) {
        TraversalSource::Table {
            component_index: index,
        }
    } else {
        TraversalSource::Sparse {
            component_index: index,
        }
    }
}

fn fill_spec_indices(
    world: &World,
    spec: &QuerySpec,
    scratch: &mut QueryResolveScratch,
) -> Result<(Vec<usize>, Vec<usize>), QueryError> {
    fill_type_and_component_ids(
        world,
        &spec.required,
        &spec.required_ids,
        false,
        &mut scratch.required,
    )?;
    fill_type_and_component_ids(
        world,
        &spec.without,
        &spec.without_ids,
        false,
        &mut scratch.without,
    )?;
    fill_type_and_component_ids(
        world,
        &spec.with_tags,
        &spec.with_tag_ids,
        true,
        &mut scratch.with_tags,
    )?;
    fill_type_and_component_ids(
        world,
        &spec.without_tags,
        &spec.without_tag_ids,
        true,
        &mut scratch.without_tags,
    )?;

    if spec.exact_ids.is_some() && spec.exact_id_policy.is_none() {
        return Err(QueryError::WrongQuery {
            detail: String::from("exact-id queries require an explicit ExactIdPolicy"),
        });
    }

    let mut added = resolve_selector_group(world, &spec.added, &spec.added_ids)?;
    let mut changed = resolve_selector_group(world, &spec.changed, &spec.changed_ids)?;
    normalize(&mut added);
    normalize(&mut changed);
    if !added.is_empty() && !changed.is_empty() {
        return Err(QueryError::ConflictingFilters {
            detail: String::from("added and changed filters are mutually exclusive"),
        });
    }
    Ok((added, changed))
}

fn fill_type_and_component_ids(
    world: &World,
    type_ids: &[TypeId],
    component_ids: &[ComponentId],
    require_tag: bool,
    out: &mut Vec<usize>,
) -> Result<(), QueryError> {
    out.clear();
    for &type_id in type_ids {
        let id = resolve_type_id(world, type_id)?;
        validate_tag_kind(world, &id, require_tag)?;
        out.push(id.index());
    }
    for id in component_ids {
        validate_component_id(world, id)?;
        validate_tag_kind(world, id, require_tag)?;
        out.push(id.index());
    }
    normalize(out);
    Ok(())
}

fn resolve_selector_group(
    world: &World,
    type_ids: &[TypeId],
    component_ids: &[ComponentId],
) -> Result<Vec<usize>, QueryError> {
    let mut out = Vec::with_capacity(type_ids.len() + component_ids.len());
    for &type_id in type_ids {
        out.push(resolve_type_id(world, type_id)?.index());
    }
    for id in component_ids {
        validate_component_id(world, id)?;
        out.push(id.index());
    }
    Ok(out)
}

fn validate_component_id(world: &World, id: &ComponentId) -> Result<(), QueryError> {
    if id.validate_owner(world.owner()).is_err() {
        return Err(QueryError::WrongOwner);
    }
    if !world.registry_contains(id) {
        return Err(QueryError::UnregisteredComponent {
            name: String::from("<stale component id>"),
        });
    }
    Ok(())
}

fn validate_tag_kind(world: &World, id: &ComponentId, require_tag: bool) -> Result<(), QueryError> {
    if require_tag && !world.is_tag_component(id) {
        return Err(QueryError::WrongStorageKind {
            name: world.registry_component_name(id),
        });
    }
    Ok(())
}

fn normalize(indices: &mut Vec<usize>) {
    indices.sort_unstable();
    indices.dedup();
}

fn validate_overlaps(scratch: &QueryResolveScratch) -> Result<(), QueryError> {
    validate_no_overlap(&scratch.required, &scratch.without, "required", "without")?;
    validate_no_overlap(
        &scratch.required,
        &scratch.without_tags,
        "required",
        "without_tag",
    )?;
    validate_no_overlap(&scratch.with_tags, &scratch.without, "with_tag", "without")?;
    validate_no_overlap(
        &scratch.with_tags,
        &scratch.without_tags,
        "with_tag",
        "without_tag",
    )
}

fn resolve_component<T: 'static>(world: &World) -> Result<ComponentId, QueryError> {
    world
        .registry_id_of::<T>()
        .ok_or_else(|| QueryError::UnregisteredComponent {
            name: String::from(type_name::<T>()),
        })
}

fn resolve_type_id(world: &World, type_id: TypeId) -> Result<ComponentId, QueryError> {
    world
        .registry_id_of_type(type_id)
        .ok_or_else(|| QueryError::UnregisteredComponent {
            name: String::from("<unregistered component>"),
        })
}

fn validate_no_overlap(
    left: &[usize],
    right: &[usize],
    left_name: &str,
    right_name: &str,
) -> Result<(), QueryError> {
    for index in left {
        if right.contains(index) {
            return Err(QueryError::ConflictingFilters {
                detail: format!("{left_name} and {right_name} both reference index {index}"),
            });
        }
    }
    Ok(())
}

#[allow(clippy::too_many_arguments)]
fn fingerprint_plan(
    required: &[usize],
    without: &[usize],
    with_tags: &[usize],
    without_tags: &[usize],
    added: &[usize],
    changed: &[usize],
    traversal: &TraversalSource,
    primary: Option<usize>,
    secondary: Option<usize>,
    exact_id_policy: Option<ExactIdPolicy>,
) -> u64 {
    let mut hasher = FnvHasher::new();
    match (primary, secondary) {
        (None, None) => 0u8.hash(&mut hasher),
        (Some(_), None) => 1u8.hash(&mut hasher),
        (Some(_), Some(_)) => 2u8.hash(&mut hasher),
        (None, Some(_)) => unreachable!("secondary query component requires a primary"),
    }
    primary.hash(&mut hasher);
    secondary.hash(&mut hasher);
    for index in required {
        index.hash(&mut hasher);
    }
    for index in without {
        index.hash(&mut hasher);
        1u8.hash(&mut hasher);
    }
    for index in with_tags {
        index.hash(&mut hasher);
        2u8.hash(&mut hasher);
    }
    for index in without_tags {
        index.hash(&mut hasher);
        3u8.hash(&mut hasher);
    }
    added.hash(&mut hasher);
    changed.hash(&mut hasher);
    exact_id_policy.hash(&mut hasher);
    match traversal {
        TraversalSource::All => {
            0u8.hash(&mut hasher);
        }
        TraversalSource::Sparse { .. } | TraversalSource::Table { .. } => {
            // Sparse/table traversal and the selected component are physical
            // execution choices. The semantic identity is already captured by
            // the primary and normalized filter component sets above.
            1u8.hash(&mut hasher);
        }
        TraversalSource::Exact { ids } => {
            3u8.hash(&mut hasher);
            ids.len().hash(&mut hasher);
            for id in ids {
                id.hash(&mut hasher);
            }
        }
    }
    hasher.finish()
}

struct FnvHasher(u64);

impl FnvHasher {
    fn new() -> Self {
        Self(0xcbf29ce484222325)
    }
}

impl Hasher for FnvHasher {
    fn finish(&self) -> u64 {
        self.0
    }

    fn write(&mut self, bytes: &[u8]) {
        for byte in bytes {
            self.0 ^= *byte as u64;
            self.0 = self.0.wrapping_mul(0x100000001b3);
        }
    }

    fn write_u64(&mut self, i: u64) {
        self.0 ^= i;
        self.0 = self.0.wrapping_mul(0x100000001b3);
    }

    fn write_usize(&mut self, i: usize) {
        self.write_u64(i as u64);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::component::ComponentOptions;
    use crate::query::{ExactIdPolicy, QueryError, QuerySpec};
    use crate::world::WorldBuilder;
    use alloc::vec;

    #[derive(Clone, Copy)]
    struct Position(#[allow(dead_code)] i32);

    #[derive(Clone, Copy)]
    struct Velocity(#[allow(dead_code)] i32);

    #[derive(Clone, Copy)]
    struct TablePosition(#[allow(dead_code)] i32);

    #[derive(Clone, Copy)]
    struct TableVelocity(#[allow(dead_code)] i32);

    #[derive(Clone, Copy)]
    struct Player;

    #[derive(Clone, Copy)]
    struct Ghost;

    fn world() -> World {
        let mut builder = WorldBuilder::new();
        builder
            .register_component::<Position>(ComponentOptions::sparse())
            .expect("pos");
        builder
            .register_component::<Velocity>(ComponentOptions::sparse())
            .expect("vel");
        builder
            .register_component::<TablePosition>(ComponentOptions::table())
            .expect("table pos");
        builder
            .register_component::<TableVelocity>(ComponentOptions::table())
            .expect("table vel");
        builder
            .register_component::<Player>(ComponentOptions::tag())
            .expect("tag");
        builder.build().expect("build")
    }

    #[test]
    fn exact_ids_without_policy_is_rejected() {
        let mut world = world();
        let mut spec = QuerySpec::new();
        spec.exact_ids = Some(vec![]);
        assert!(matches!(
            world.resolve_query1_plan::<Position>(&spec),
            Err(QueryError::WrongQuery { .. })
        ));
        assert!(matches!(
            world.resolve_query2_plan::<Position, Velocity>(&spec),
            Err(QueryError::WrongQuery { .. })
        ));
    }

    #[test]
    fn added_and_changed_filters_conflict() {
        let mut world = world();
        let spec = QuerySpec::new().added::<Position>().changed::<Velocity>();
        assert!(matches!(
            world.resolve_query1_plan::<Position>(&spec),
            Err(QueryError::ConflictingFilters { .. })
        ));
        assert!(matches!(
            world.resolve_query2_plan::<Position, Velocity>(&spec),
            Err(QueryError::ConflictingFilters { .. })
        ));
    }

    #[test]
    fn overlapping_required_and_without_conflict() {
        let mut world = world();
        let spec = QuerySpec::new().with::<Position>().without::<Position>();
        assert!(matches!(
            world.resolve_query1_plan::<Position>(&spec),
            Err(QueryError::ConflictingFilters { .. })
        ));
    }

    #[test]
    fn query2_overlapping_tag_filters_conflict() {
        let mut world = world();
        let spec = QuerySpec::new()
            .with_tag::<Player>()
            .without_tag::<Player>();
        assert!(matches!(
            world.resolve_query2_plan::<Position, Velocity>(&spec),
            Err(QueryError::ConflictingFilters { .. })
        ));
    }

    #[test]
    fn unregistered_filter_type_is_rejected() {
        let mut world = world();
        let spec = QuerySpec::new().without::<Ghost>();
        assert!(matches!(
            world.resolve_query1_plan::<Position>(&spec),
            Err(QueryError::UnregisteredComponent { .. })
        ));
    }

    #[test]
    fn non_tag_with_tag_filter_is_wrong_storage_kind() {
        let mut world = world();
        let spec = QuerySpec::new().with_tag::<Position>();
        assert!(matches!(
            world.resolve_query1_plan::<Position>(&spec),
            Err(QueryError::WrongStorageKind { .. })
        ));
    }

    #[test]
    fn unregistered_without_tag_filter_is_rejected() {
        let mut world = world();
        let spec = QuerySpec::new().without_tag::<Ghost>();
        assert!(matches!(
            world.resolve_query1_plan::<Position>(&spec),
            Err(QueryError::UnregisteredComponent { .. })
        ));
    }

    #[test]
    fn query2_exact_ids_use_exact_traversal() {
        let mut world = world();
        let entity = world.spawn().expect("spawn");
        world.insert(entity, Position(1)).expect("insert");
        let spec = QuerySpec::new().exact_ids(vec![entity], ExactIdPolicy::SkipUnavailable);
        let (plan, _, _) = world
            .resolve_query2_plan::<Position, Velocity>(&spec)
            .expect("plan");
        assert!(matches!(plan.traversal, TraversalSource::Exact { .. }));
    }

    #[test]
    fn query2_uses_smaller_driver_across_storage_pairs() {
        let mut world = world();
        let first = world.spawn().expect("first");
        for value in 0..3 {
            let entity = if value == 0 {
                first
            } else {
                world.spawn().expect("entity")
            };
            world.insert(entity, Position(value)).expect("sparse A");
            world.insert(entity, TablePosition(value)).expect("table A");
        }
        world.insert(first, Velocity(1)).expect("sparse B");
        world.insert(first, TableVelocity(1)).expect("table B");

        let sparse_b = world.component_index::<Velocity>().expect("sparse B index");
        let table_b = world
            .component_index::<TableVelocity>()
            .expect("table B index");

        let (ss, _, _) = world
            .resolve_query2_plan::<Position, Velocity>(&QuerySpec::new())
            .expect("sparse/sparse");
        assert!(matches!(
            ss.traversal,
            TraversalSource::Sparse { component_index } if component_index == sparse_b
        ));

        let (st, _, _) = world
            .resolve_query2_plan::<Position, TableVelocity>(&QuerySpec::new())
            .expect("sparse/table");
        assert!(matches!(
            st.traversal,
            TraversalSource::Table { component_index } if component_index == table_b
        ));

        let (ts, _, _) = world
            .resolve_query2_plan::<TablePosition, Velocity>(&QuerySpec::new())
            .expect("table/sparse");
        assert!(matches!(
            ts.traversal,
            TraversalSource::Sparse { component_index } if component_index == sparse_b
        ));

        let (tt, _, _) = world
            .resolve_query2_plan::<TablePosition, TableVelocity>(&QuerySpec::new())
            .expect("table/table");
        assert!(matches!(
            tt.traversal,
            TraversalSource::Table { component_index } if component_index == table_b
        ));
    }

    #[test]
    fn entity_specs_resolve_dynamic_ids_and_storage_traversal() {
        let mut builder = WorldBuilder::new();
        let sparse = builder
            .register_component::<Position>(ComponentOptions::sparse())
            .expect("sparse");
        let table = builder
            .register_component::<TablePosition>(ComponentOptions::table())
            .expect("table");
        let tag = builder
            .register_component::<Player>(ComponentOptions::tag())
            .expect("tag");
        let mut world = builder.build().expect("build");

        let sparse_plan = world
            .resolve_entity_plan(
                &QuerySpec::new()
                    .with_id(sparse.clone())
                    .without_tag_id(tag.clone())
                    .added_id(sparse),
            )
            .expect("sparse dynamic plan");
        assert!(matches!(
            sparse_plan.traversal,
            TraversalSource::Sparse { .. }
        ));
        assert_eq!(sparse_plan.added_indices.len(), 1);

        let table_plan = world
            .resolve_entity_plan(&QuerySpec::new().with_id(table.clone()).changed_id(table))
            .expect("table dynamic plan");
        assert!(matches!(
            table_plan.traversal,
            TraversalSource::Table { .. }
        ));

        let tag_plan = world
            .resolve_entity_plan(&QuerySpec::new().with_tag_id(tag))
            .expect("tag dynamic plan");
        assert!(matches!(tag_plan.traversal, TraversalSource::Sparse { .. }));
    }

    #[test]
    fn dynamic_component_ids_reject_foreign_stale_and_overlapping_selectors() {
        let mut builder = WorldBuilder::new();
        let player = builder
            .register_component::<Player>(ComponentOptions::tag())
            .expect("player");
        let mut world = builder.build().expect("build");

        let mut foreign_builder = WorldBuilder::new();
        let foreign = foreign_builder
            .register_component::<Position>(ComponentOptions::sparse())
            .expect("foreign");
        assert!(matches!(
            world.resolve_entity_plan(&QuerySpec::new().with_id(foreign)),
            Err(QueryError::WrongOwner)
        ));

        let stale = ComponentId::new(world.owner_token(), 999);
        assert!(matches!(
            world.resolve_entity_plan(&QuerySpec::new().with_id(stale)),
            Err(QueryError::UnregisteredComponent { .. })
        ));

        assert!(matches!(
            world.resolve_entity_plan(
                &QuerySpec::new()
                    .with_id(player.clone())
                    .without_tag_id(player)
            ),
            Err(QueryError::ConflictingFilters { .. })
        ));
    }

    #[test]
    #[should_panic(expected = "secondary query component requires a primary")]
    fn fingerprint_rejects_secondary_without_primary() {
        let _ = fingerprint_plan(
            &[],
            &[],
            &[],
            &[],
            &[],
            &[],
            &TraversalSource::All,
            None,
            Some(0),
            None,
        );
    }
}