bevy/crates/bevy_ecs/src/component/register.rs
2025-07-17 18:56:12 +02:00

684 lines
25 KiB
Rust

use alloc::{boxed::Box, vec::Vec};
use bevy_platform::sync::PoisonError;
use bevy_utils::TypeIdMap;
use core::any::Any;
use core::ops::DerefMut;
use core::{any::TypeId, fmt::Debug, ops::Deref};
use crate::component::{enforce_no_required_components_recursion, RequiredComponentsRegistrator};
use crate::query::DebugCheckedUnwrap as _;
use crate::{
component::{
Component, ComponentDescriptor, ComponentId, Components, RequiredComponents, StorageType,
},
resource::Resource,
};
/// Generates [`ComponentId`]s.
#[derive(Debug, Default)]
pub struct ComponentIds {
next: bevy_platform::sync::atomic::AtomicUsize,
}
impl ComponentIds {
/// Peeks the next [`ComponentId`] to be generated without generating it.
pub fn peek(&self) -> ComponentId {
ComponentId(
self.next
.load(bevy_platform::sync::atomic::Ordering::Relaxed),
)
}
/// Generates and returns the next [`ComponentId`].
pub fn next(&self) -> ComponentId {
ComponentId(
self.next
.fetch_add(1, bevy_platform::sync::atomic::Ordering::Relaxed),
)
}
/// Peeks the next [`ComponentId`] to be generated without generating it.
pub fn peek_mut(&mut self) -> ComponentId {
ComponentId(*self.next.get_mut())
}
/// Generates and returns the next [`ComponentId`].
pub fn next_mut(&mut self) -> ComponentId {
let id = self.next.get_mut();
let result = ComponentId(*id);
*id += 1;
result
}
/// Returns the number of [`ComponentId`]s generated.
pub fn len(&self) -> usize {
self.peek().0
}
/// Returns true if and only if no ids have been generated.
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
/// A [`Components`] wrapper that enables additional features, like registration.
pub struct ComponentsRegistrator<'w> {
components: &'w mut Components,
ids: &'w mut ComponentIds,
pub(super) recursion_check_stack: Vec<ComponentId>,
}
impl Deref for ComponentsRegistrator<'_> {
type Target = Components;
fn deref(&self) -> &Self::Target {
self.components
}
}
impl DerefMut for ComponentsRegistrator<'_> {
fn deref_mut(&mut self) -> &mut Self::Target {
self.components
}
}
impl<'w> ComponentsRegistrator<'w> {
/// Constructs a new [`ComponentsRegistrator`].
///
/// # Safety
///
/// The [`Components`] and [`ComponentIds`] must match.
/// For example, they must be from the same world.
pub unsafe fn new(components: &'w mut Components, ids: &'w mut ComponentIds) -> Self {
Self {
components,
ids,
recursion_check_stack: Vec::new(),
}
}
/// Converts this [`ComponentsRegistrator`] into a [`ComponentsQueuedRegistrator`].
/// This is intended for use to pass this value to a function that requires [`ComponentsQueuedRegistrator`].
/// It is generally not a good idea to queue a registration when you can instead register directly on this type.
pub fn as_queued(&self) -> ComponentsQueuedRegistrator<'_> {
// SAFETY: ensured by the caller that created self.
unsafe { ComponentsQueuedRegistrator::new(self.components, self.ids) }
}
/// Applies every queued registration.
/// This ensures that every valid [`ComponentId`] is registered,
/// enabling retrieving [`ComponentInfo`](super::ComponentInfo), etc.
pub fn apply_queued_registrations(&mut self) {
if !self.any_queued_mut() {
return;
}
// Note:
//
// This is not just draining the queue. We need to empty the queue without removing the information from `Components`.
// If we drained directly, we could break invariance.
//
// For example, say `ComponentA` and `ComponentB` are queued, and `ComponentA` requires `ComponentB`.
// If we drain directly, and `ComponentA` was the first to be registered, then, when `ComponentA`
// registers `ComponentB` in `Component::register_required_components`,
// `Components` will not know that `ComponentB` was queued
// (since it will have been drained from the queue.)
// If that happened, `Components` would assign a new `ComponentId` to `ComponentB`
// which would be *different* than the id it was assigned in the queue.
// Then, when the drain iterator gets to `ComponentB`,
// it would be unsafely registering `ComponentB`, which is already registered.
//
// As a result, we need to pop from each queue one by one instead of draining.
// components
while let Some(registrator) = {
let queued = self
.components
.queued
.get_mut()
.unwrap_or_else(PoisonError::into_inner);
queued.components.keys().next().copied().map(|type_id| {
// SAFETY: the id just came from a valid iterator.
unsafe { queued.components.remove(&type_id).debug_checked_unwrap() }
})
} {
registrator.register(self);
}
// resources
while let Some(registrator) = {
let queued = self
.components
.queued
.get_mut()
.unwrap_or_else(PoisonError::into_inner);
queued.resources.keys().next().copied().map(|type_id| {
// SAFETY: the id just came from a valid iterator.
unsafe { queued.resources.remove(&type_id).debug_checked_unwrap() }
})
} {
registrator.register(self);
}
// dynamic
let queued = &mut self
.components
.queued
.get_mut()
.unwrap_or_else(PoisonError::into_inner);
if !queued.dynamic_registrations.is_empty() {
for registrator in core::mem::take(&mut queued.dynamic_registrations) {
registrator.register(self);
}
}
}
/// Registers a [`Component`] of type `T` with this instance.
/// If a component of this type has already been registered, this will return
/// the ID of the pre-existing component.
///
/// # See also
///
/// * [`Components::component_id()`]
/// * [`ComponentsRegistrator::register_component_with_descriptor()`]
#[inline]
pub fn register_component<T: Component>(&mut self) -> ComponentId {
self.register_component_checked::<T>()
}
/// Same as [`Self::register_component_unchecked`] but keeps a checks for safety.
#[inline]
pub(super) fn register_component_checked<T: Component>(&mut self) -> ComponentId {
let type_id = TypeId::of::<T>();
if let Some(&id) = self.indices.get(&type_id) {
enforce_no_required_components_recursion(self, &self.recursion_check_stack, id);
return id;
}
if let Some(registrator) = self
.components
.queued
.get_mut()
.unwrap_or_else(PoisonError::into_inner)
.components
.remove(&type_id)
{
// If we are trying to register something that has already been queued, we respect the queue.
// Just like if we are trying to register something that already is, we respect the first registration.
return registrator.register(self);
}
let id = self.ids.next_mut();
// SAFETY: The component is not currently registered, and the id is fresh.
unsafe {
self.register_component_unchecked::<T>(id);
}
id
}
/// # Safety
///
/// Neither this component, nor its id may be registered or queued. This must be a new registration.
#[inline]
unsafe fn register_component_unchecked<T: Component>(&mut self, id: ComponentId) {
// SAFETY: ensured by caller.
unsafe {
self.register_component_inner(id, ComponentDescriptor::new::<T>());
}
let type_id = TypeId::of::<T>();
let prev = self.indices.insert(type_id, id);
debug_assert!(prev.is_none());
self.recursion_check_stack.push(id);
let mut required_components = RequiredComponents::default();
// SAFETY: `required_components` is empty
let mut required_components_registrator =
unsafe { RequiredComponentsRegistrator::new(self, &mut required_components) };
T::register_required_components(id, &mut required_components_registrator);
// SAFETY:
// - `id` was just registered in `self`
// - RequiredComponentsRegistrator guarantees that only components from `self` are included in `required_components`.
unsafe { self.register_required_by(id, &required_components) };
self.recursion_check_stack.pop();
// SAFETY: we just inserted it in `register_component_inner`
let info = unsafe {
&mut self
.components
.components
.get_mut(id.0)
.debug_checked_unwrap()
.as_mut()
.debug_checked_unwrap()
};
info.hooks.update_from_component::<T>();
info.required_components = required_components;
}
/// Registers a component described by `descriptor`.
///
/// # Note
///
/// If this method is called multiple times with identical descriptors, a distinct [`ComponentId`]
/// will be created for each one.
///
/// # See also
///
/// * [`Components::component_id()`]
/// * [`ComponentsRegistrator::register_component()`]
#[inline]
pub fn register_component_with_descriptor(
&mut self,
descriptor: ComponentDescriptor,
) -> ComponentId {
let id = self.ids.next_mut();
// SAFETY: The id is fresh.
unsafe {
self.register_component_inner(id, descriptor);
}
id
}
/// Registers a [`Resource`] of type `T` with this instance.
/// If a resource of this type has already been registered, this will return
/// the ID of the pre-existing resource.
///
/// # See also
///
/// * [`Components::resource_id()`]
/// * [`ComponentsRegistrator::register_resource_with_descriptor()`]
#[inline]
pub fn register_resource<T: Resource>(&mut self) -> ComponentId {
// SAFETY: The [`ComponentDescriptor`] matches the [`TypeId`]
unsafe {
self.register_resource_with(TypeId::of::<T>(), || {
ComponentDescriptor::new_resource::<T>()
})
}
}
/// Registers a [non-send resource](crate::system::NonSend) of type `T` with this instance.
/// If a resource of this type has already been registered, this will return
/// the ID of the pre-existing resource.
#[inline]
pub fn register_non_send<T: Any>(&mut self) -> ComponentId {
// SAFETY: The [`ComponentDescriptor`] matches the [`TypeId`]
unsafe {
self.register_resource_with(TypeId::of::<T>(), || {
ComponentDescriptor::new_non_send::<T>(StorageType::default())
})
}
}
/// Same as [`Components::register_resource_unchecked`] but handles safety.
///
/// # Safety
///
/// The [`ComponentDescriptor`] must match the [`TypeId`].
#[inline]
unsafe fn register_resource_with(
&mut self,
type_id: TypeId,
descriptor: impl FnOnce() -> ComponentDescriptor,
) -> ComponentId {
if let Some(id) = self.resource_indices.get(&type_id) {
return *id;
}
if let Some(registrator) = self
.components
.queued
.get_mut()
.unwrap_or_else(PoisonError::into_inner)
.resources
.remove(&type_id)
{
// If we are trying to register something that has already been queued, we respect the queue.
// Just like if we are trying to register something that already is, we respect the first registration.
return registrator.register(self);
}
let id = self.ids.next_mut();
// SAFETY: The resource is not currently registered, the id is fresh, and the [`ComponentDescriptor`] matches the [`TypeId`]
unsafe {
self.register_resource_unchecked(type_id, id, descriptor());
}
id
}
/// Registers a [`Resource`] described by `descriptor`.
///
/// # Note
///
/// If this method is called multiple times with identical descriptors, a distinct [`ComponentId`]
/// will be created for each one.
///
/// # See also
///
/// * [`Components::resource_id()`]
/// * [`ComponentsRegistrator::register_resource()`]
#[inline]
pub fn register_resource_with_descriptor(
&mut self,
descriptor: ComponentDescriptor,
) -> ComponentId {
let id = self.ids.next_mut();
// SAFETY: The id is fresh.
unsafe {
self.register_component_inner(id, descriptor);
}
id
}
}
/// A queued component registration.
pub(super) struct QueuedRegistration {
pub(super) registrator:
Box<dyn FnOnce(&mut ComponentsRegistrator, ComponentId, ComponentDescriptor)>,
pub(super) id: ComponentId,
pub(super) descriptor: ComponentDescriptor,
}
impl QueuedRegistration {
/// Creates the [`QueuedRegistration`].
///
/// # Safety
///
/// [`ComponentId`] must be unique.
unsafe fn new(
id: ComponentId,
descriptor: ComponentDescriptor,
func: impl FnOnce(&mut ComponentsRegistrator, ComponentId, ComponentDescriptor) + 'static,
) -> Self {
Self {
registrator: Box::new(func),
id,
descriptor,
}
}
/// Performs the registration, returning the now valid [`ComponentId`].
pub(super) fn register(self, registrator: &mut ComponentsRegistrator) -> ComponentId {
(self.registrator)(registrator, self.id, self.descriptor);
self.id
}
}
/// Allows queuing components to be registered.
#[derive(Default)]
pub struct QueuedComponents {
pub(super) components: TypeIdMap<QueuedRegistration>,
pub(super) resources: TypeIdMap<QueuedRegistration>,
pub(super) dynamic_registrations: Vec<QueuedRegistration>,
}
impl Debug for QueuedComponents {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
let components = self
.components
.iter()
.map(|(type_id, queued)| (type_id, queued.id))
.collect::<Vec<_>>();
let resources = self
.resources
.iter()
.map(|(type_id, queued)| (type_id, queued.id))
.collect::<Vec<_>>();
let dynamic_registrations = self
.dynamic_registrations
.iter()
.map(|queued| queued.id)
.collect::<Vec<_>>();
write!(f, "components: {components:?}, resources: {resources:?}, dynamic_registrations: {dynamic_registrations:?}")
}
}
/// A type that enables queuing registration in [`Components`].
///
/// # Note
///
/// These queued registrations return [`ComponentId`]s.
/// These ids are not yet valid, but they will become valid
/// when either [`ComponentsRegistrator::apply_queued_registrations`] is called or the same registration is made directly.
/// In either case, the returned [`ComponentId`]s will be correct, but they are not correct yet.
///
/// Generally, that means these [`ComponentId`]s can be safely used for read-only purposes.
/// Modifying the contents of the world through these [`ComponentId`]s directly without waiting for them to be fully registered
/// and without then confirming that they have been fully registered is not supported.
/// Hence, extra care is needed with these [`ComponentId`]s to ensure all safety rules are followed.
///
/// As a rule of thumb, if you have mutable access to [`ComponentsRegistrator`], prefer to use that instead.
/// Use this only if you need to know the id of a component but do not need to modify the contents of the world based on that id.
#[derive(Clone, Copy)]
pub struct ComponentsQueuedRegistrator<'w> {
components: &'w Components,
ids: &'w ComponentIds,
}
impl Deref for ComponentsQueuedRegistrator<'_> {
type Target = Components;
fn deref(&self) -> &Self::Target {
self.components
}
}
impl<'w> ComponentsQueuedRegistrator<'w> {
/// Constructs a new [`ComponentsQueuedRegistrator`].
///
/// # Safety
///
/// The [`Components`] and [`ComponentIds`] must match.
/// For example, they must be from the same world.
pub unsafe fn new(components: &'w Components, ids: &'w ComponentIds) -> Self {
Self { components, ids }
}
/// Queues this function to run as a component registrator.
///
/// # Safety
///
/// The [`TypeId`] must not already be registered or queued as a component.
unsafe fn force_register_arbitrary_component(
&self,
type_id: TypeId,
descriptor: ComponentDescriptor,
func: impl FnOnce(&mut ComponentsRegistrator, ComponentId, ComponentDescriptor) + 'static,
) -> ComponentId {
let id = self.ids.next();
self.components
.queued
.write()
.unwrap_or_else(PoisonError::into_inner)
.components
.insert(
type_id,
// SAFETY: The id was just generated.
unsafe { QueuedRegistration::new(id, descriptor, func) },
);
id
}
/// Queues this function to run as a resource registrator.
///
/// # Safety
///
/// The [`TypeId`] must not already be registered or queued as a resource.
unsafe fn force_register_arbitrary_resource(
&self,
type_id: TypeId,
descriptor: ComponentDescriptor,
func: impl FnOnce(&mut ComponentsRegistrator, ComponentId, ComponentDescriptor) + 'static,
) -> ComponentId {
let id = self.ids.next();
self.components
.queued
.write()
.unwrap_or_else(PoisonError::into_inner)
.resources
.insert(
type_id,
// SAFETY: The id was just generated.
unsafe { QueuedRegistration::new(id, descriptor, func) },
);
id
}
/// Queues this function to run as a dynamic registrator.
fn force_register_arbitrary_dynamic(
&self,
descriptor: ComponentDescriptor,
func: impl FnOnce(&mut ComponentsRegistrator, ComponentId, ComponentDescriptor) + 'static,
) -> ComponentId {
let id = self.ids.next();
self.components
.queued
.write()
.unwrap_or_else(PoisonError::into_inner)
.dynamic_registrations
.push(
// SAFETY: The id was just generated.
unsafe { QueuedRegistration::new(id, descriptor, func) },
);
id
}
/// This is a queued version of [`ComponentsRegistrator::register_component`].
/// This will reserve an id and queue the registration.
/// These registrations will be carried out at the next opportunity.
///
/// If this has already been registered or queued, this returns the previous [`ComponentId`].
///
/// # Note
///
/// Technically speaking, the returned [`ComponentId`] is not valid, but it will become valid later.
/// See type level docs for details.
#[inline]
pub fn queue_register_component<T: Component>(&self) -> ComponentId {
self.component_id::<T>().unwrap_or_else(|| {
// SAFETY: We just checked that this type was not in the queue.
unsafe {
self.force_register_arbitrary_component(
TypeId::of::<T>(),
ComponentDescriptor::new::<T>(),
|registrator, id, _descriptor| {
// SAFETY: We just checked that this is not currently registered or queued, and if it was registered since, this would have been dropped from the queue.
#[expect(unused_unsafe, reason = "More precise to specify.")]
unsafe {
registrator.register_component_unchecked::<T>(id);
}
},
)
}
})
}
/// This is a queued version of [`ComponentsRegistrator::register_component_with_descriptor`].
/// This will reserve an id and queue the registration.
/// These registrations will be carried out at the next opportunity.
///
/// # Note
///
/// Technically speaking, the returned [`ComponentId`] is not valid, but it will become valid later.
/// See type level docs for details.
#[inline]
pub fn queue_register_component_with_descriptor(
&self,
descriptor: ComponentDescriptor,
) -> ComponentId {
self.force_register_arbitrary_dynamic(descriptor, |registrator, id, descriptor| {
// SAFETY: Id uniqueness handled by caller.
unsafe {
registrator.register_component_inner(id, descriptor);
}
})
}
/// This is a queued version of [`ComponentsRegistrator::register_resource`].
/// This will reserve an id and queue the registration.
/// These registrations will be carried out at the next opportunity.
///
/// If this has already been registered or queued, this returns the previous [`ComponentId`].
///
/// # Note
///
/// Technically speaking, the returned [`ComponentId`] is not valid, but it will become valid later.
/// See type level docs for details.
#[inline]
pub fn queue_register_resource<T: Resource>(&self) -> ComponentId {
let type_id = TypeId::of::<T>();
self.get_resource_id(type_id).unwrap_or_else(|| {
// SAFETY: We just checked that this type was not in the queue.
unsafe {
self.force_register_arbitrary_resource(
type_id,
ComponentDescriptor::new_resource::<T>(),
move |registrator, id, descriptor| {
// SAFETY: We just checked that this is not currently registered or queued, and if it was registered since, this would have been dropped from the queue.
// SAFETY: Id uniqueness handled by caller, and the type_id matches descriptor.
#[expect(unused_unsafe, reason = "More precise to specify.")]
unsafe {
registrator.register_resource_unchecked(type_id, id, descriptor);
}
},
)
}
})
}
/// This is a queued version of [`ComponentsRegistrator::register_non_send`].
/// This will reserve an id and queue the registration.
/// These registrations will be carried out at the next opportunity.
///
/// If this has already been registered or queued, this returns the previous [`ComponentId`].
///
/// # Note
///
/// Technically speaking, the returned [`ComponentId`] is not valid, but it will become valid later.
/// See type level docs for details.
#[inline]
pub fn queue_register_non_send<T: Any>(&self) -> ComponentId {
let type_id = TypeId::of::<T>();
self.get_resource_id(type_id).unwrap_or_else(|| {
// SAFETY: We just checked that this type was not in the queue.
unsafe {
self.force_register_arbitrary_resource(
type_id,
ComponentDescriptor::new_non_send::<T>(StorageType::default()),
move |registrator, id, descriptor| {
// SAFETY: We just checked that this is not currently registered or queued, and if it was registered since, this would have been dropped from the queue.
// SAFETY: Id uniqueness handled by caller, and the type_id matches descriptor.
#[expect(unused_unsafe, reason = "More precise to specify.")]
unsafe {
registrator.register_resource_unchecked(type_id, id, descriptor);
}
},
)
}
})
}
/// This is a queued version of [`ComponentsRegistrator::register_resource_with_descriptor`].
/// This will reserve an id and queue the registration.
/// These registrations will be carried out at the next opportunity.
///
/// # Note
///
/// Technically speaking, the returned [`ComponentId`] is not valid, but it will become valid later.
/// See type level docs for details.
#[inline]
pub fn queue_register_resource_with_descriptor(
&self,
descriptor: ComponentDescriptor,
) -> ComponentId {
self.force_register_arbitrary_dynamic(descriptor, |registrator, id, descriptor| {
// SAFETY: Id uniqueness handled by caller.
unsafe {
registrator.register_component_inner(id, descriptor);
}
})
}
}