Unique Components Enforcement in Bevy 0.19

unique_enum.rs

// License
// 
// The following is free and open source. All code in this file is dual-licensed under either:
// 
// - MIT License (<http://opensource.org/licenses/MIT>)
// - Apache License, Version 2.0 (<http://www.apache.org/licenses/LICENSE-2.0>)
// 
// at your option.

use bevy::prelude::*;

use core::{any::type_name, marker::PhantomData};
use std::mem::take;
use bevy::ecs::{world::DeferredWorld, lifecycle::HookContext};
use smallvec::SmallVec;

/// Wrap a component during insertion/spawn with a variant of this enum to maintain uniqueness of that component.
/// You should also add [`EnforceUniqueComponent`] to your component's required components.
/// An error is logged if the component does not have [`EnforceUniqueComponent`]
///
/// This wrapper's variants exist for specific use cases. See the following:
///
/// [`Unique::Always`]
/// [`Unique::Error`]
/// [`Unique::Toggle`]
/// [`Unique::IfNew`]
/// [`Unique::Unenforced`]
/// [`Unique::IgnoreAdditions`]
/// [`Unique::DespawnAdditions`]
/// [`Unique::DespawnOld`]
///
/// ```rust
/// #[derive(Component)]
/// // Enforce uniqueness and error if a Unique wrapper is not used
/// #[require(EnforceUniqueComponent::<Self>)]
/// struct MyUniqueComponent;
///
/// #[derive(Component)]
/// // Wrapper use in required components
/// #[require(
///     Unique::<_>::Always(MyUniqueComponent),
/// )]
/// struct MyOtherComponent;
///
/// fn insert_unique_system(mut commands: Commands) {
///     // Wrapper use in spawn
///     commands.spawn(Unique::Always(MyUniqueComponent))
///     // Wrapper use in insert
///     commands.spawn_empty().insert(Unique::Always(MyUniqueComponent))
///     // Wrapper use in bsn
///     commands.spawn_scene(bsn! {
///         template_value(Unique::Always(MyUniqueComponent))
///     });
/// }
/// ```
#[derive(Component, Default, Clone)]
#[component(storage = "SparseSet", on_insert = Self::unique_insert_hook)]
#[doc(hidden)]
pub enum Unique<T: Component> {
    /// A [`Unique`] wrapper variant.
    ///
    /// Wrap a component with this to maintain uniqueness of that component.
    ///
    /// The component will be removed from any entities that currently have it.
    ///
    /// However, the new entity will still recieve the unique component.
    Always(T),

    /// A [`Unique`] wrapper variant.
    ///
    /// Wrap a component with this to maintain uniqueness of that component.
    ///
    /// Functions the same as the [`Unique::Always`] wrapper
    ///
    /// This variant will log an error if the component already exists on an entity.
    ///
    /// However, the new entity will still recieve the unique component.
    Error(T),

    /// A [`Unique`] wrapper variant.
    ///
    /// Wrap a component with this to maintain uniqueness of that component.
    ///
    /// Functions the same as the [`Unique::Always`] wrapper
    ///
    /// This variant will remove the component if it already exists on that entity
    Toggle(T),

    /// A [`Unique`] wrapper variant.
    ///
    /// Wrap a component with this to maintain uniqueness of that component.
    ///
    /// Functions the same as the [`Unique::Always`] wrapper
    ///
    /// This variant will not attempt to reinsert the component if it already exists on that entity
    IfNew(T),

    /// A [`Unique`] wrapper variant.
    ///
    /// Wrap a component with this to maintain uniqueness of that component.
    ///
    /// Functions the same as the [`Unique::Always`] wrapper
    ///
    /// This variant will not enforce uniqueness via a required component check.
    /// This is useful for component types you don't have control of e.g. from a library.
    Unenforced(T),

    /// A [`Unique`] wrapper variant.
    ///
    /// Wrap a component with this to maintain uniqueness of that component.
    ///
    /// Functions the same as the [`Unique::Always`] wrapper
    ///
    /// This variant will not allow new insertions after the first.
    /// This applies to all Unique wrapper variants.
    /// If a wrapper is not used, there will be a momentary duplicate before
    /// it is removed from the new location.
    IgnoreAdditions(T),

    /// A [`Unique`] wrapper variant.
    ///
    /// Wrap a component with this to maintain uniqueness of that component.
    ///
    /// Functions the same as the [`Unique::Always`] wrapper
    ///
    /// This variant will despawn entities with new insertions after the first.
    /// This applies to all Unique wrapper variants.
    /// If a wrapper is not used, there will be a momentary duplicate before
    /// the new entity is despawned
    DespawnAdditions(T),

    /// A [`Unique`] wrapper variant.
    ///
    /// Wrap a component with this to maintain uniqueness of that component.
    ///
    /// Functions the same as the [`Unique::Always`] wrapper
    ///
    /// This variant will despawn old entities with the component.
    DespawnOld(T),

    #[default]
    #[doc(hidden)]
    /// This is only for internal use inside [Unique::unique_insert_hook].
    /// This will take the place of the wrapped component during insertion
    _Taken,
}

impl <T: Component> Unique<T> {
    fn should_despawn(&self) -> bool {
        match self {
            Unique::DespawnAdditions(_) => true,
            Unique::DespawnOld(_) => true,

            Unique::Always(_) => false,
            Unique::Error(_) => false,
            Unique::Toggle(_) => false,
            Unique::IfNew(_) => false,
            Unique::Unenforced(_) => false,
            Unique::IgnoreAdditions(_) => false,
            Unique::_Taken => false,
        }
    }

    fn should_error(&self) -> bool {
        match self {
            Unique::Error(_) => true,
            Unique::IgnoreAdditions(_) => true,
            Unique::DespawnAdditions(_) => true,
            Unique::DespawnOld(_) => true,

            Unique::Always(_) => false,
            Unique::IfNew(_) => false,
            Unique::Toggle(_) => false,
            Unique::Unenforced(_) => false,
            Unique::_Taken => false,
        }
    }

    fn should_reinsert(&self) -> bool {
        match self {
            Unique::Always(_) => true,
            Unique::Error(_) => true,
            Unique::Unenforced(_) => true,
            Unique::DespawnOld(_) => true,

            Unique::IfNew(_) => false,
            Unique::Toggle(_) => false,
            Unique::IgnoreAdditions(_) => false,
            Unique::DespawnAdditions(_) => false,
            Unique::_Taken => false,
        }
    }

    fn unique_insert_hook(
        mut world: DeferredWorld,
        context: HookContext,
    ) {
        let mut inserted_query = world.get_mut::<Self>(context.entity);
        let Some(inserted) = inserted_query.as_deref_mut() else {
            error!("Unable to find unique component wrapper for:\n{}", type_name::<T>());
            return;
        };

        // remove the value from the entity and replace with default variant (Unique::_Taken)
        let wrapper = take(inserted);
        // remove the old wrapper from the entity
        world.commands().entity(context.entity).remove::<Self>();

        // Find all existing entities of type T
        // SmallVec (suggested by Eugineerd on Discord) with size 2 (1 for new, 1 for existing)
        // Size over 2 will still be handled, but will exist on the heap instead
        let entities: SmallVec<[Entity; 2]> = match world.try_query_filtered::<Entity, With<T>>() {
            Some(mut existing_query) => existing_query.iter(&world).collect(),
            // The query could fail if we haven't yet registered the component
            // e.g. when spawning first time via BSN, so default to an empty vec
            None => default(),
        };

        let mut has_existing = false;

        for entity in &entities {
            // clear existing entities except the one we're inserting the unique component on
            // (unless toggling)
            if context.entity != *entity || matches!(wrapper, Self::Toggle(_)) {
                has_existing = true;

                if wrapper.should_despawn() {
                    world.commands().entity(*entity).try_despawn();
                } else {
                    world.commands().entity(*entity).try_remove::<T>();
                }
            }
        }

        if wrapper.should_error() && has_existing {
            error!("Unique component was unexpectedly inserted on a new entity:\n{}", type_name::<T>());
        }

        let ignore = world.try_query_filtered::<(), With<IgnoreUniqueComponent<T>>>()
            .is_some_and(|mut query| query.iter(&world).next().is_some());

        let insert = !ignore && (wrapper.should_reinsert() || !entities.contains(&context.entity));

        if ignore {
            let despawn = world.try_query_filtered::<Entity, With<DespawnUniqueComponent<T>>>()
                .is_some_and(|mut query| query.iter(&world).next()
                    .is_some_and(|entity| entity != context.entity));

            if despawn {
                error!("Component was already inserted with Unique::DespawnAdditions, the new entity will be despawned:\n{}", type_name::<T>());
                world.commands().entity(context.entity).try_despawn();
            } else {
                error!("Component was already inserted with Unique::IgnoreAdditions or Unique::DespawnAdditions, the new one won't be inserted:\n{}", type_name::<T>());
            }
        }

        // must set this to a variable here to not borrow after taking the value from the wrapper
        let is_enforced = !matches!(wrapper, Self::Unenforced(_));

        if insert {
            // must do this first to not borrow after taking the value from the wrapper
            if matches!(wrapper, Self::IgnoreAdditions(_)) {
                world.commands().entity(context.entity).try_insert(IgnoreUniqueComponent::<T>::default());
            }

            // must do this first to not borrow after taking the value from the wrapper
            if matches!(wrapper, Self::DespawnAdditions(_)) {
                world.commands().entity(context.entity).try_insert(DespawnUniqueComponent::<T>::default());
            }

            let new_component = match wrapper {
                Self::Always(c)
                | Self::Toggle(c)
                | Self::IfNew(c)
                | Self::Error(c)
                | Self::Unenforced(c)
                | Self::IgnoreAdditions(c)
                | Self::DespawnOld(c)
                | Self::DespawnAdditions(c) => c,

                Self::_Taken => {
                    error!("Unable to extract unique component - value already taken:\n{}", type_name::<T>());
                    return;
                },
            };

            world.commands().entity(context.entity)
                // add a new component to verify that we inserted with the unique wrapper
                .try_insert(VerifyUniqueComponent::<T>::default())
                // insert the new unique component with the contents from the wrapper
                .try_insert(new_component);
        };

        if is_enforced {
            // This is done last and queued because the `EnforceUniqueComponent<T>` may not yet be registered
            // until the component is inserted, otherwise resulting in false negatives.
            world.commands().entity(context.entity).queue(move |mut entity_world: EntityWorldMut| {
                // check if the component has EnforceUnique<T> as a required component
                entity_world.world_scope(|scoped_world| {
                    let has_enforced = match scoped_world.component_id::<EnforceUniqueComponent<T>>() {
                        // find the EnforceUnique required component
                        Some(enforce_id) => match scoped_world.get_required_components::<T>() {
                            Some(required) => required.iter_ids().any(|id| id == enforce_id),
                            None => false,
                        },
                        None => false,
                    };

                    if !has_enforced {
                        error!("Component inserted with UniqueComponent is not EnforceUnique:\n{}", type_name::<T>());
                    }
                });
            });
        }

    }
}

/// This hook is called when [`EnforceUniqueComponent`] is added to a component
/// In the case a [`Unique`] wrapper was used to insert the component,
/// Then the [`VerifyUniqueComponent`] will be present and that means we were able
/// to insert the component without any risk of momentary duplicates.
/// In the case where a [`Unique`] wrapper was not used properly,
/// the fallback plan is to try to remove duplicates anyway and error since that
/// could result in a momentary duplicate which would violate Observers with a `Single`
/// query of the unique component.
///
/// In the case [`Unique::DespawnAdditions`] or [`Unique::IgnoreAdditions`] were used
/// before the current invocation, the new entity will be despawned
/// or the new component will be removed respectively.
/// These conditions are caught with the existence of [`IgnoreUniqueComponent`] and
/// [`DespawnUniqueComponent`] added by the use of the original wrapper.
fn enforce_unique_hook<T: Component>(
    mut world: DeferredWorld,
    context: HookContext,
) {
    let is_verified = match world.try_query_filtered::<Entity, With<VerifyUniqueComponent<T>>>() {
        Some(mut verify_query) => verify_query.get(&world, context.entity).is_ok(),
        None => {
            error!("Unable to get verified unique components:\n{}", type_name::<T>());
            false
        },
    };

    world.commands().entity(context.entity)
        // We've now verified this hook, remove the verification marker
        .remove::<VerifyUniqueComponent<T>>();

    if is_verified {
        return;
    }

    error!("Component marked with EnforceUnique was inserted without UniqueComponent wrapper:\n{}", type_name::<T>());

    let ignore = world.try_query_filtered::<Entity, With<IgnoreUniqueComponent<T>>>()
        .is_some_and(|mut query| query.iter(&world)
            .next()
            .inspect(|entity| if *entity != context.entity {
                // remove the commponent from attempted impostors
                world.commands().entity(context.entity).try_remove::<T>();
            })
            .is_some()
        );

    if ignore {
        let despawn = world.try_query_filtered::<(), With<DespawnUniqueComponent<T>>>()
            .is_some_and(|mut query| query.iter(&world).next().is_some());

        if despawn {
            error!("Attempted to insert a unique component that was already inserted with Unique::DespawnAdditions, the new entity will be despawned:\n{}", type_name::<T>());
            world.commands().entity(context.entity).try_despawn();
        } else {
            error!("Attempted to insert a unique component that was already inserted with Unique::IgnoreAdditions, the component won't be inserted:\n{}", type_name::<T>());
        }

        return;
    }

    let Some(mut existing_query) = world.try_query_filtered::<Entity, With<T>>() else {
        error!("Unable to clear existing unique component:\n{}", type_name::<T>()); return;
    };

    let entities: Vec<Entity> = existing_query.iter(&world).collect();

    // remove all existing entities just in case the UniqueComponent component was forgotten
    for entity in &entities {
        if context.entity != *entity {
            world.commands().entity(*entity).try_remove::<T>();
        }
    }
}

/// Add this component as a required component maintain uniqueness of that component.
/// Use a [`Unique`] wrapper variant
/// when inserting/spawning that component to gurantee uniqueness
///
/// An error is logged if the component is not spawned with a [`Unique`]
/// wrapper variant.
///
/// A fallback hook ([`enforce_unique_hook`]) will attempt to remove duplicates regardless,
/// but there may be a momentary duplicate during hooks and observers.
///
/// ```rust
/// #[derive(Component)]
/// #[require(EnforceUniqueComponent::<Self>)]
/// pub struct MyUniqueComponent;
/// ```
#[derive(Component, Debug)]
#[component(storage = "SparseSet", on_insert = enforce_unique_hook::<T>)]
pub struct EnforceUniqueComponent<T: Component> {
    _marker: PhantomData<T>,
}

impl <T: Component> Default for EnforceUniqueComponent<T> {
    fn default() -> Self {
        Self { _marker: default() }
    }
}

/// This component is added to an entity after a unique component is properly inserted
/// via a [`Unique`] wrapper variant
///
/// When a component marked with [`EnforceUniqueComponent`] is inserted incorrectly
/// The absence of this component will trigger the hook to send an error and fallback
/// remove duplicates. In that case a short period of duplicates is possible which will
/// cause `Single<MyUniqueComponent>` in observers to fail.
///
/// This Component is purposely not public and is only inserted in [`Unique::unique_insert_hook`].
#[derive(Component, Debug)]
#[component(storage = "SparseSet")]
struct VerifyUniqueComponent<T: Component> {
    _marker: PhantomData<T>,
}

impl <T: Component> Default for VerifyUniqueComponent<T> {
    fn default() -> Self {
        Self { _marker: default() }
    }
}

/// This component is added to an entity after a unique component is properly inserted
/// via the [`Unique::IgnoreAdditions`] or [`Unique::DespawnAdditions`] wrapper variants.
///
/// When present when [`EnforceUniqueComponent`] is inserted incorrectly the new component
/// will be removed and the old unique instance will remain.
///
/// This Component is purposely not public and is only inserted in [`Unique::unique_insert_hook`].
#[derive(Component, Debug)]
#[component(storage = "SparseSet")]
struct IgnoreUniqueComponent<T: Component> {
    _marker: PhantomData<T>,
}

impl <T: Component> Default for IgnoreUniqueComponent<T> {
    fn default() -> Self {
        Self { _marker: default() }
    }
}

/// This component is added to an entity after a unique component is properly inserted
/// via the [`Unique::DespawnAdditions`] wrapper variant.
///
/// When present when [`EnforceUniqueComponent`] is inserted incorrectly the entity the
/// new component is inserted on will be despawned entirely and the old instance will remain.
///
/// This Component is purposely not public and is only inserted in [`Unique::unique_insert_hook`].
#[derive(Component, Debug)]
#[component(storage = "SparseSet")]
#[require(IgnoreUniqueComponent<T>)]
struct DespawnUniqueComponent<T: Component> {
    _marker: PhantomData<T>,
}

impl <T: Component> Default for DespawnUniqueComponent<T> {
    fn default() -> Self {
        Self { _marker: default() }
    }
}