bevy_entity_uuid 0.18.0

//! Support uuids to identify entities in bevy.
//!
//! Bevy does not easily expose control over manually creating an entity with
//! an arbritrary value due to its internal requirements. This will allow to
//! keep track of entities with some manually specified id instead.
//!
//! This plugin adds a new component that can mark an entity with a uuid.
//! The plugin will be able to do fast lookups uuid<>entity and keep track of
//! entity creation or destroy.
//!
//! This can be a useful plugin when things need to be synchronized across the
//! network clients where entity ids are different but they need somehow to be
//! unified under one id representation.
//!
//! The normal method to create an entity in bevy is to do so from a Command.
//! This module will add an extra method to create entities from a Uuid. It
//! will still call the normal entity creation underneath, but also take care
//! to add relevant component and tracking resources to bind it to the wanted
//! uuid on creation.
//!
//! To attach a uuid to an entity instead, just add the [EntityUuid] component
//! to the wanted entity.
//!
//! Two methods are added to Commands with the trait [CommandsEntityUuid]:
//! - spawn_uuid(uuid, ...):
//!   equivalent of cmd.spawn(...) but with added uuid
//! - spawn_empty_uuid(uuid):
//!   equivalent of cmd.spawn_empty() but with added uuid
//!
//! To setup this plugin, add the plugin [EntityUuidPlugin] to bevy:
//!
//! ```
//! use bevy_entity_uuid::EntityUuidPlugin;
//! use bevy_ecs::prelude::*;
//! use bevy_app::prelude::*;
//!
//! let mut app = App::new();
//! app.add_plugins(EntityUuidPlugin);
//! ```
//!
//! Then when using Commands, spawn with uuid variants:
//!
//! ```
//! use bevy_entity_uuid::CommandsEntityUuid;
//! use bevy::prelude::*;
//! use uuid::Uuid;
//!
//! fn system_with_empty(mut cmd: Commands) {
//!     let uuid = Uuid::new_v4();
//!     cmd.spawn_empty_uuid(uuid);
//! }
//!
//! fn system_with_bundle(mut cmd: Commands) {
//!     let uuid = Uuid::new_v4();
//!     cmd.spawn_uuid(uuid, Transform::default());
//! }
//! ```

use bevy_app::prelude::*;
use bevy_ecs::component::Immutable;
use bevy_ecs::component::StorageType;
use bevy_ecs::lifecycle::ComponentHook;
use bevy_ecs::prelude::*;

use std::collections::HashMap;
use uuid::Uuid;

/// Added methods to spawn entities with a uuid.
pub trait CommandsEntityUuid {
    /// Same as cmd.spawn(bundle) but with added uuid
    fn spawn_uuid<T: Bundle>(&mut self, uuid: Uuid, bundle: T) -> EntityCommands<'_>;
    /// Same as cmd.spawn_empty() but with added uuid
    fn spawn_empty_uuid(&mut self, uuid: Uuid) -> EntityCommands<'_>;
}

impl CommandsEntityUuid for Commands<'_, '_> {
    fn spawn_uuid<T: Bundle>(&mut self, uuid: Uuid, bundle: T) -> EntityCommands<'_> {
        let mut cmd = self.spawn(bundle);
        cmd.insert(EntityUuid { uuid });
        cmd
    }

    fn spawn_empty_uuid(&mut self, uuid: Uuid) -> EntityCommands<'_> {
        let mut cmd = self.spawn_empty();
        cmd.insert(EntityUuid { uuid });
        cmd
    }
}

/// This trait is for solving the eventual consistency of Commands: using world
/// the same uuid won't be spawned twice, but return the original first entity.
pub trait WorldEntityUuid {
    fn spawn_uuid<B: Bundle>(&mut self, uuid: Uuid, bundle: B) -> EntityWorldMut<'_>;
    fn spawn_empty_uuid(&mut self, uuid: Uuid) -> EntityWorldMut<'_>;
}

impl WorldEntityUuid for World {
    fn spawn_uuid<B: Bundle>(&mut self, uuid: Uuid, bundle: B) -> EntityWorldMut<'_> {
        let map = self.resource_mut::<EntityLookup>();
        if let Some(id) = map.uuid_to_entity(&uuid) {
            return self.entity_mut(id);
        }
        let mut cmd = self.spawn(bundle);
        cmd.insert(EntityUuid { uuid });
        cmd
    }

    fn spawn_empty_uuid(&mut self, uuid: Uuid) -> EntityWorldMut<'_> {
        let map = self.resource_mut::<EntityLookup>();
        if let Some(id) = map.uuid_to_entity(&uuid) {
            return self.entity_mut(id);
        }
        let mut cmd = self.spawn_empty();
        cmd.insert(EntityUuid { uuid });
        cmd
    }
}

/// Component used to keep track of which uuid is in this entity.
pub struct EntityUuid {
    pub uuid: Uuid,
}

impl Component for EntityUuid {
    const STORAGE_TYPE: StorageType = StorageType::Table;
    type Mutability = Immutable;

    fn on_add() -> Option<ComponentHook> {
        Some(|mut world, ctx| {
            let uuid = world.entity(ctx.entity).get::<EntityUuid>().unwrap().uuid;
            let mut map = world.resource_mut::<EntityLookup>();
            map.add_new(ctx.entity, uuid);
        })
    }

    fn on_remove() -> Option<ComponentHook> {
        Some(|mut world, ctx| {
            let uuid = world.entity(ctx.entity).get::<EntityUuid>().unwrap().uuid;
            let mut map = world.resource_mut::<EntityLookup>();
            map.remove(ctx.entity, uuid);
            world.write_message(EntityUuidDeleted(uuid));
        })
    }
}

/// Event triggered when a EventUuid is deleted.
/// This can be useful on entity despawn situations because otherwise it's not
/// possible to obtain the original deleted uuid: the map resourece is already
/// updated and the uuid won't be found in it.
#[derive(Message)]
pub struct EntityUuidDeleted(pub Uuid);

/// This resource allows for fast lookups uuid<>entity.
#[derive(Resource, Default, Clone, Debug)]
pub struct EntityLookup {
    uuid_to_entity: HashMap<Uuid, Entity>,
    entity_to_uuid: HashMap<Entity, Uuid>,
}

impl EntityLookup {
    /// Find an entity id given a uuid.
    pub fn uuid_to_entity(&self, uuid: &Uuid) -> Option<Entity> {
        self.uuid_to_entity.get(uuid).copied()
    }

    /// Find a uuid from an entity id.
    pub fn entity_to_uuid(&self, id: &Entity) -> Option<Uuid> {
        self.entity_to_uuid.get(id).copied()
    }

    pub(crate) fn add_new(&mut self, id: Entity, uuid: Uuid) {
        self.uuid_to_entity.insert(uuid, id);
        self.entity_to_uuid.insert(id, uuid);
    }

    pub(crate) fn remove(&mut self, id: Entity, uuid: Uuid) {
        if let Some(id) = self.uuid_to_entity.remove(&uuid) {
            self.entity_to_uuid.remove(&id);
        }
        if let Some(uuid) = self.entity_to_uuid.remove(&id) {
            self.uuid_to_entity.remove(&uuid);
        }
    }
}

/// Plugin to setup EntityUuid functonality.
pub struct EntityUuidPlugin;

impl Plugin for EntityUuidPlugin {
    fn build(&self, app: &mut App) {
        app.add_message::<EntityUuidDeleted>();
        app.init_resource::<EntityLookup>();
    }
}