azalea/

nearest_entity.rs

use azalea_core::entity_id::MinecraftEntityId;
use azalea_entity::Position;
use azalea_world::WorldName;
use bevy_ecs::{
    prelude::Entity,
    query::{QueryFilter, With},
    system::{Query, SystemParam},
};

/// This system parameter can be used as a shorthand for quickly finding an
/// entity, (or several) close to a given position.
///
/// This system parameter allows for additional filtering of entities based off
/// of ECS marker components, such as `With<T>`, `Without<T>`, or `Added<T>`,
/// etc. All functions used by this system parameter instance will respect the
/// applied filter.
///
/// ```
/// use azalea::{chat::SendChatEvent, nearest_entity::EntityFinder};
/// use azalea_entity::{
///     LocalEntity,
///     metadata::{AbstractMonster, Player},
/// };
/// use bevy_ecs::{
///     prelude::{Entity, MessageWriter},
///     query::With,
///     system::Query,
/// };
///
/// /// All bots near aggressive mobs will scream in chat.
/// pub fn bots_near_aggressive_mobs(
///     bots: Query<Entity, (With<LocalEntity>, With<Player>)>,
///     entity_finder: EntityFinder<With<AbstractMonster>>,
///     mut chat_events: MessageWriter<SendChatEvent>,
/// ) {
///     for bot_id in bots.iter() {
///         let Some(nearest) = entity_finder.nearest_to_entity(bot_id, 16.0) else {
///             continue;
///         };
///
///         chat_events.write(SendChatEvent {
///             entity: bot_id,
///             content: String::from("Ahhh!"),
///         });
///     }
/// }
/// ```
#[derive(SystemParam)]
pub struct EntityFinder<'w, 's, F = ()>
where
    F: QueryFilter + 'static,
{
    all_entities: Query<'w, 's, (&'static Position, &'static WorldName), With<MinecraftEntityId>>,

    filtered_entities: Query<
        'w,
        's,
        (Entity, &'static WorldName, &'static Position),
        (With<MinecraftEntityId>, F),
    >,
}

impl<'a, F> EntityFinder<'_, '_, F>
where
    F: QueryFilter + 'static,
{
    /// Gets the nearest entity to the given position and with the given world
    /// name.
    ///
    /// This method will return `None` if there are no entities within range. If
    /// multiple entities are within range, only the closest one is returned.
    pub fn nearest_to_position(
        &'a self,
        position: Position,
        world_name: &WorldName,
        max_distance: f64,
    ) -> Option<Entity> {
        let mut nearest_entity = None;
        let mut min_distance = max_distance;

        for (target_entity, e_world, e_pos) in self.filtered_entities.iter() {
            if e_world != world_name {
                continue;
            }

            let target_distance = position.distance_to(**e_pos);
            if target_distance < min_distance {
                nearest_entity = Some(target_entity);
                min_distance = target_distance;
            }
        }

        nearest_entity
    }

    /// Gets the nearest entity to the given entity.
    ///
    /// This method will return `None` if there are no entities within range. If
    /// multiple entities are within range, only the closest one is
    /// returned.
    pub fn nearest_to_entity(&'a self, entity: Entity, max_distance: f64) -> Option<Entity> {
        let Ok((position, world_name)) = self.all_entities.get(entity) else {
            return None;
        };

        let mut nearest_entity = None;
        let mut min_distance = max_distance;

        for (target_entity, e_world, e_pos) in self.filtered_entities.iter() {
            if entity == target_entity {
                continue;
            };

            if e_world != world_name {
                continue;
            }

            let target_distance = position.distance_to(**e_pos);
            if target_distance < min_distance {
                nearest_entity = Some(target_entity);
                min_distance = target_distance;
            }
        }

        nearest_entity
    }

    /// This function get an iterator over all nearby entities to the given
    /// position within the given maximum distance.
    ///
    /// The entities in this iterator are not returned in any specific order.
    ///
    /// This function returns the Entity ID of nearby entities and their
    /// distance away.
    pub fn nearby_entities_to_position(
        &'a self,
        position: &'a Position,
        world_name: &'a WorldName,
        max_distance: f64,
    ) -> impl Iterator<Item = (Entity, f64)> + 'a {
        self.filtered_entities
            .iter()
            .filter_map(move |(target_entity, e_world, e_pos)| {
                if e_world != world_name {
                    return None;
                }

                let distance = position.distance_to(**e_pos);
                if distance < max_distance {
                    Some((target_entity, distance))
                } else {
                    None
                }
            })
    }

    /// This function get an iterator over all nearby entities to the given
    /// entity within the given maximum distance.
    ///
    /// The entities in this iterator are not returned in any specific order.
    ///
    /// This function returns the Entity ID of nearby entities and their
    /// distance away.
    pub fn nearby_entities_to_entity(
        &'a self,
        entity: Entity,
        max_distance: f64,
    ) -> impl Iterator<Item = (Entity, f64)> + 'a {
        let position;
        let world_name;
        if let Ok((p, w)) = self.all_entities.get(entity) {
            position = *p;
            world_name = Some(w);
        } else {
            position = Position::default();
            world_name = None;
        };

        self.filtered_entities
            .iter()
            .filter_map(move |(target_entity, e_world, e_pos)| {
                if entity == target_entity {
                    return None;
                }

                if Some(e_world) != world_name {
                    return None;
                }

                let distance = position.distance_to(**e_pos);
                if distance < max_distance {
                    Some((target_entity, distance))
                } else {
                    None
                }
            })
    }
}