bevy_ecs/component/
register.rs

1use alloc::{boxed::Box, vec::Vec};
2use bevy_platform::sync::PoisonError;
3use bevy_utils::TypeIdMap;
4use core::any::Any;
5use core::{any::TypeId, fmt::Debug, ops::Deref};
6
7use crate::component::{enforce_no_required_components_recursion, RequiredComponentsRegistrator};
8use crate::{
9    component::{
10        Component, ComponentDescriptor, ComponentId, Components, RequiredComponents, StorageType,
11    },
12    query::DebugCheckedUnwrap as _,
13    resource::Resource,
14};
15
16/// Generates [`ComponentId`]s.
17#[derive(Debug, Default)]
18pub struct ComponentIds {
19    next: bevy_platform::sync::atomic::AtomicUsize,
20}
21
22impl ComponentIds {
23    /// Peeks the next [`ComponentId`] to be generated without generating it.
24    pub fn peek(&self) -> ComponentId {
25        ComponentId(
26            self.next
27                .load(bevy_platform::sync::atomic::Ordering::Relaxed),
28        )
29    }
30
31    /// Generates and returns the next [`ComponentId`].
32    pub fn next(&self) -> ComponentId {
33        ComponentId(
34            self.next
35                .fetch_add(1, bevy_platform::sync::atomic::Ordering::Relaxed),
36        )
37    }
38
39    /// Peeks the next [`ComponentId`] to be generated without generating it.
40    pub fn peek_mut(&mut self) -> ComponentId {
41        ComponentId(*self.next.get_mut())
42    }
43
44    /// Generates and returns the next [`ComponentId`].
45    pub fn next_mut(&mut self) -> ComponentId {
46        let id = self.next.get_mut();
47        let result = ComponentId(*id);
48        *id += 1;
49        result
50    }
51
52    /// Returns the number of [`ComponentId`]s generated.
53    pub fn len(&self) -> usize {
54        self.peek().0
55    }
56
57    /// Returns true if and only if no ids have been generated.
58    pub fn is_empty(&self) -> bool {
59        self.len() == 0
60    }
61}
62
63/// A [`Components`] wrapper that enables additional features, like registration.
64pub struct ComponentsRegistrator<'w> {
65    pub(super) components: &'w mut Components,
66    pub(super) ids: &'w mut ComponentIds,
67    pub(super) recursion_check_stack: Vec<ComponentId>,
68}
69
70impl Deref for ComponentsRegistrator<'_> {
71    type Target = Components;
72
73    fn deref(&self) -> &Self::Target {
74        self.components
75    }
76}
77
78impl<'w> ComponentsRegistrator<'w> {
79    /// Constructs a new [`ComponentsRegistrator`].
80    ///
81    /// # Safety
82    ///
83    /// The [`Components`] and [`ComponentIds`] must match.
84    /// For example, they must be from the same world.
85    pub unsafe fn new(components: &'w mut Components, ids: &'w mut ComponentIds) -> Self {
86        Self {
87            components,
88            ids,
89            recursion_check_stack: Vec::new(),
90        }
91    }
92
93    /// Converts this [`ComponentsRegistrator`] into a [`ComponentsQueuedRegistrator`].
94    /// This is intended for use to pass this value to a function that requires [`ComponentsQueuedRegistrator`].
95    /// It is generally not a good idea to queue a registration when you can instead register directly on this type.
96    pub fn as_queued(&self) -> ComponentsQueuedRegistrator<'_> {
97        // SAFETY: ensured by the caller that created self.
98        unsafe { ComponentsQueuedRegistrator::new(self.components, self.ids) }
99    }
100
101    /// Applies every queued registration.
102    /// This ensures that every valid [`ComponentId`] is registered,
103    /// enabling retrieving [`ComponentInfo`](super::ComponentInfo), etc.
104    pub fn apply_queued_registrations(&mut self) {
105        if !self.any_queued_mut() {
106            return;
107        }
108
109        // Note:
110        //
111        // This is not just draining the queue. We need to empty the queue without removing the information from `Components`.
112        // If we drained directly, we could break invariance.
113        //
114        // For example, say `ComponentA` and `ComponentB` are queued, and `ComponentA` requires `ComponentB`.
115        // If we drain directly, and `ComponentA` was the first to be registered, then, when `ComponentA`
116        // registers `ComponentB` in `Component::register_required_components`,
117        // `Components` will not know that `ComponentB` was queued
118        // (since it will have been drained from the queue.)
119        // If that happened, `Components` would assign a new `ComponentId` to `ComponentB`
120        // which would be *different* than the id it was assigned in the queue.
121        // Then, when the drain iterator gets to `ComponentB`,
122        // it would be unsafely registering `ComponentB`, which is already registered.
123        //
124        // As a result, we need to pop from each queue one by one instead of draining.
125
126        // components
127        while let Some(registrator) = {
128            let queued = self
129                .components
130                .queued
131                .get_mut()
132                .unwrap_or_else(PoisonError::into_inner);
133            queued.components.keys().next().copied().map(|type_id| {
134                // SAFETY: the id just came from a valid iterator.
135                unsafe { queued.components.remove(&type_id).debug_checked_unwrap() }
136            })
137        } {
138            registrator.register(self);
139        }
140
141        // resources
142        while let Some(registrator) = {
143            let queued = self
144                .components
145                .queued
146                .get_mut()
147                .unwrap_or_else(PoisonError::into_inner);
148            queued.resources.keys().next().copied().map(|type_id| {
149                // SAFETY: the id just came from a valid iterator.
150                unsafe { queued.resources.remove(&type_id).debug_checked_unwrap() }
151            })
152        } {
153            registrator.register(self);
154        }
155
156        // dynamic
157        let queued = &mut self
158            .components
159            .queued
160            .get_mut()
161            .unwrap_or_else(PoisonError::into_inner);
162        if !queued.dynamic_registrations.is_empty() {
163            for registrator in core::mem::take(&mut queued.dynamic_registrations) {
164                registrator.register(self);
165            }
166        }
167    }
168
169    /// Registers a [`Component`] of type `T` with this instance.
170    /// If a component of this type has already been registered, this will return
171    /// the ID of the pre-existing component.
172    ///
173    /// # See also
174    ///
175    /// * [`Components::component_id()`]
176    /// * [`ComponentsRegistrator::register_component_with_descriptor()`]
177    #[inline]
178    pub fn register_component<T: Component>(&mut self) -> ComponentId {
179        self.register_component_checked::<T>()
180    }
181
182    /// Same as [`Self::register_component_unchecked`] but keeps a checks for safety.
183    #[inline]
184    pub(super) fn register_component_checked<T: Component>(&mut self) -> ComponentId {
185        let type_id = TypeId::of::<T>();
186        if let Some(&id) = self.indices.get(&type_id) {
187            enforce_no_required_components_recursion(self, &self.recursion_check_stack, id);
188            return id;
189        }
190
191        if let Some(registrator) = self
192            .components
193            .queued
194            .get_mut()
195            .unwrap_or_else(PoisonError::into_inner)
196            .components
197            .remove(&type_id)
198        {
199            // If we are trying to register something that has already been queued, we respect the queue.
200            // Just like if we are trying to register something that already is, we respect the first registration.
201            return registrator.register(self);
202        }
203
204        let id = self.ids.next_mut();
205        // SAFETY: The component is not currently registered, and the id is fresh.
206        unsafe {
207            self.register_component_unchecked::<T>(id);
208        }
209        id
210    }
211
212    /// # Safety
213    ///
214    /// Neither this component, nor its id may be registered or queued. This must be a new registration.
215    #[inline]
216    unsafe fn register_component_unchecked<T: Component>(&mut self, id: ComponentId) {
217        // SAFETY: ensured by caller.
218        unsafe {
219            self.components
220                .register_component_inner(id, ComponentDescriptor::new::<T>());
221        }
222        let type_id = TypeId::of::<T>();
223        let prev = self.components.indices.insert(type_id, id);
224        debug_assert!(prev.is_none());
225
226        self.recursion_check_stack.push(id);
227        let mut required_components = RequiredComponents::default();
228        // SAFETY: `required_components` is empty
229        let mut required_components_registrator =
230            unsafe { RequiredComponentsRegistrator::new(self, &mut required_components) };
231        T::register_required_components(id, &mut required_components_registrator);
232        // SAFETY:
233        // - `id` was just registered in `self`
234        // - RequiredComponentsRegistrator guarantees that only components from `self` are included in `required_components`;
235        // - we just initialized the component with id `id` so no component requiring it can exist yet.
236        unsafe {
237            self.components
238                .register_required_by(id, &required_components);
239        }
240        self.recursion_check_stack.pop();
241
242        // SAFETY: we just inserted it in `register_component_inner`
243        let info = unsafe {
244            &mut self
245                .components
246                .components
247                .get_mut(id.0)
248                .debug_checked_unwrap()
249                .as_mut()
250                .debug_checked_unwrap()
251        };
252
253        info.hooks.update_from_component::<T>();
254
255        info.required_components = required_components;
256    }
257
258    /// Registers a component described by `descriptor`.
259    ///
260    /// # Note
261    ///
262    /// If this method is called multiple times with identical descriptors, a distinct [`ComponentId`]
263    /// will be created for each one.
264    ///
265    /// # See also
266    ///
267    /// * [`Components::component_id()`]
268    /// * [`ComponentsRegistrator::register_component()`]
269    #[inline]
270    pub fn register_component_with_descriptor(
271        &mut self,
272        descriptor: ComponentDescriptor,
273    ) -> ComponentId {
274        let id = self.ids.next_mut();
275        // SAFETY: The id is fresh.
276        unsafe {
277            self.components.register_component_inner(id, descriptor);
278        }
279        id
280    }
281
282    /// Registers a [`Resource`] of type `T` with this instance.
283    /// If a resource of this type has already been registered, this will return
284    /// the ID of the pre-existing resource.
285    ///
286    /// # See also
287    ///
288    /// * [`Components::resource_id()`]
289    /// * [`ComponentsRegistrator::register_resource_with_descriptor()`]
290    #[inline]
291    pub fn register_resource<T: Resource>(&mut self) -> ComponentId {
292        // SAFETY: The [`ComponentDescriptor`] matches the [`TypeId`]
293        unsafe {
294            self.register_resource_with(TypeId::of::<T>(), || {
295                ComponentDescriptor::new_resource::<T>()
296            })
297        }
298    }
299
300    /// Registers a [non-send resource](crate::system::NonSend) of type `T` with this instance.
301    /// If a resource of this type has already been registered, this will return
302    /// the ID of the pre-existing resource.
303    #[inline]
304    pub fn register_non_send<T: Any>(&mut self) -> ComponentId {
305        // SAFETY: The [`ComponentDescriptor`] matches the [`TypeId`]
306        unsafe {
307            self.register_resource_with(TypeId::of::<T>(), || {
308                ComponentDescriptor::new_non_send::<T>(StorageType::default())
309            })
310        }
311    }
312
313    /// Same as [`Components::register_resource_unchecked`] but handles safety.
314    ///
315    /// # Safety
316    ///
317    /// The [`ComponentDescriptor`] must match the [`TypeId`].
318    #[inline]
319    unsafe fn register_resource_with(
320        &mut self,
321        type_id: TypeId,
322        descriptor: impl FnOnce() -> ComponentDescriptor,
323    ) -> ComponentId {
324        if let Some(id) = self.resource_indices.get(&type_id) {
325            return *id;
326        }
327
328        if let Some(registrator) = self
329            .components
330            .queued
331            .get_mut()
332            .unwrap_or_else(PoisonError::into_inner)
333            .resources
334            .remove(&type_id)
335        {
336            // If we are trying to register something that has already been queued, we respect the queue.
337            // Just like if we are trying to register something that already is, we respect the first registration.
338            return registrator.register(self);
339        }
340
341        let id = self.ids.next_mut();
342        // SAFETY: The resource is not currently registered, the id is fresh, and the [`ComponentDescriptor`] matches the [`TypeId`]
343        unsafe {
344            self.components
345                .register_resource_unchecked(type_id, id, descriptor());
346        }
347        id
348    }
349
350    /// Registers a [`Resource`] described by `descriptor`.
351    ///
352    /// # Note
353    ///
354    /// If this method is called multiple times with identical descriptors, a distinct [`ComponentId`]
355    /// will be created for each one.
356    ///
357    /// # See also
358    ///
359    /// * [`Components::resource_id()`]
360    /// * [`ComponentsRegistrator::register_resource()`]
361    #[inline]
362    pub fn register_resource_with_descriptor(
363        &mut self,
364        descriptor: ComponentDescriptor,
365    ) -> ComponentId {
366        let id = self.ids.next_mut();
367        // SAFETY: The id is fresh.
368        unsafe {
369            self.components.register_component_inner(id, descriptor);
370        }
371        id
372    }
373
374    /// Equivalent of `Components::any_queued_mut`
375    pub fn any_queued_mut(&mut self) -> bool {
376        self.components.any_queued_mut()
377    }
378
379    /// Equivalent of `Components::any_queued_mut`
380    pub fn num_queued_mut(&mut self) -> usize {
381        self.components.num_queued_mut()
382    }
383}
384
385/// A queued component registration.
386pub(super) struct QueuedRegistration {
387    pub(super) registrator:
388        Box<dyn FnOnce(&mut ComponentsRegistrator, ComponentId, ComponentDescriptor)>,
389    pub(super) id: ComponentId,
390    pub(super) descriptor: ComponentDescriptor,
391}
392
393impl QueuedRegistration {
394    /// Creates the [`QueuedRegistration`].
395    ///
396    /// # Safety
397    ///
398    /// [`ComponentId`] must be unique.
399    unsafe fn new(
400        id: ComponentId,
401        descriptor: ComponentDescriptor,
402        func: impl FnOnce(&mut ComponentsRegistrator, ComponentId, ComponentDescriptor) + 'static,
403    ) -> Self {
404        Self {
405            registrator: Box::new(func),
406            id,
407            descriptor,
408        }
409    }
410
411    /// Performs the registration, returning the now valid [`ComponentId`].
412    pub(super) fn register(self, registrator: &mut ComponentsRegistrator) -> ComponentId {
413        (self.registrator)(registrator, self.id, self.descriptor);
414        self.id
415    }
416}
417
418/// Allows queuing components to be registered.
419#[derive(Default)]
420pub struct QueuedComponents {
421    pub(super) components: TypeIdMap<QueuedRegistration>,
422    pub(super) resources: TypeIdMap<QueuedRegistration>,
423    pub(super) dynamic_registrations: Vec<QueuedRegistration>,
424}
425
426impl Debug for QueuedComponents {
427    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
428        let components = self
429            .components
430            .iter()
431            .map(|(type_id, queued)| (type_id, queued.id))
432            .collect::<Vec<_>>();
433        let resources = self
434            .resources
435            .iter()
436            .map(|(type_id, queued)| (type_id, queued.id))
437            .collect::<Vec<_>>();
438        let dynamic_registrations = self
439            .dynamic_registrations
440            .iter()
441            .map(|queued| queued.id)
442            .collect::<Vec<_>>();
443        write!(f, "components: {components:?}, resources: {resources:?}, dynamic_registrations: {dynamic_registrations:?}")
444    }
445}
446
447/// A type that enables queuing registration in [`Components`].
448///
449/// # Note
450///
451/// These queued registrations return [`ComponentId`]s.
452/// These ids are not yet valid, but they will become valid
453/// when either [`ComponentsRegistrator::apply_queued_registrations`] is called or the same registration is made directly.
454/// In either case, the returned [`ComponentId`]s will be correct, but they are not correct yet.
455///
456/// Generally, that means these [`ComponentId`]s can be safely used for read-only purposes.
457/// Modifying the contents of the world through these [`ComponentId`]s directly without waiting for them to be fully registered
458/// and without then confirming that they have been fully registered is not supported.
459/// Hence, extra care is needed with these [`ComponentId`]s to ensure all safety rules are followed.
460///
461/// As a rule of thumb, if you have mutable access to [`ComponentsRegistrator`], prefer to use that instead.
462/// 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.
463#[derive(Clone, Copy)]
464pub struct ComponentsQueuedRegistrator<'w> {
465    components: &'w Components,
466    ids: &'w ComponentIds,
467}
468
469impl Deref for ComponentsQueuedRegistrator<'_> {
470    type Target = Components;
471
472    fn deref(&self) -> &Self::Target {
473        self.components
474    }
475}
476
477impl<'w> ComponentsQueuedRegistrator<'w> {
478    /// Constructs a new [`ComponentsQueuedRegistrator`].
479    ///
480    /// # Safety
481    ///
482    /// The [`Components`] and [`ComponentIds`] must match.
483    /// For example, they must be from the same world.
484    pub unsafe fn new(components: &'w Components, ids: &'w ComponentIds) -> Self {
485        Self { components, ids }
486    }
487
488    /// Queues this function to run as a component registrator if the given
489    /// type is not already queued as a component.
490    ///
491    /// # Safety
492    ///
493    /// The [`TypeId`] must not already be registered as a component.
494    unsafe fn register_arbitrary_component(
495        &self,
496        type_id: TypeId,
497        descriptor: ComponentDescriptor,
498        func: impl FnOnce(&mut ComponentsRegistrator, ComponentId, ComponentDescriptor) + 'static,
499    ) -> ComponentId {
500        self.components
501            .queued
502            .write()
503            .unwrap_or_else(PoisonError::into_inner)
504            .components
505            .entry(type_id)
506            .or_insert_with(|| {
507                // SAFETY: The id was just generated.
508                unsafe { QueuedRegistration::new(self.ids.next(), descriptor, func) }
509            })
510            .id
511    }
512
513    /// Queues this function to run as a resource registrator if the given
514    /// type is not already queued as a resource.
515    ///
516    /// # Safety
517    ///
518    /// The [`TypeId`] must not already be registered as a resource.
519    unsafe fn register_arbitrary_resource(
520        &self,
521        type_id: TypeId,
522        descriptor: ComponentDescriptor,
523        func: impl FnOnce(&mut ComponentsRegistrator, ComponentId, ComponentDescriptor) + 'static,
524    ) -> ComponentId {
525        self.components
526            .queued
527            .write()
528            .unwrap_or_else(PoisonError::into_inner)
529            .resources
530            .entry(type_id)
531            .or_insert_with(|| {
532                // SAFETY: The id was just generated.
533                unsafe { QueuedRegistration::new(self.ids.next(), descriptor, func) }
534            })
535            .id
536    }
537
538    /// Queues this function to run as a dynamic registrator.
539    fn register_arbitrary_dynamic(
540        &self,
541        descriptor: ComponentDescriptor,
542        func: impl FnOnce(&mut ComponentsRegistrator, ComponentId, ComponentDescriptor) + 'static,
543    ) -> ComponentId {
544        let id = self.ids.next();
545        self.components
546            .queued
547            .write()
548            .unwrap_or_else(PoisonError::into_inner)
549            .dynamic_registrations
550            .push(
551                // SAFETY: The id was just generated.
552                unsafe { QueuedRegistration::new(id, descriptor, func) },
553            );
554        id
555    }
556
557    /// This is a queued version of [`ComponentsRegistrator::register_component`].
558    /// This will reserve an id and queue the registration.
559    /// These registrations will be carried out at the next opportunity.
560    ///
561    /// If this has already been registered or queued, this returns the previous [`ComponentId`].
562    ///
563    /// # Note
564    ///
565    /// Technically speaking, the returned [`ComponentId`] is not valid, but it will become valid later.
566    /// See type level docs for details.
567    #[inline]
568    pub fn queue_register_component<T: Component>(&self) -> ComponentId {
569        self.component_id::<T>().unwrap_or_else(|| {
570            // SAFETY: We just checked that this type was not already registered.
571            unsafe {
572                self.register_arbitrary_component(
573                    TypeId::of::<T>(),
574                    ComponentDescriptor::new::<T>(),
575                    |registrator, id, _descriptor| {
576                        // 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.
577                        #[expect(unused_unsafe, reason = "More precise to specify.")]
578                        unsafe {
579                            registrator.register_component_unchecked::<T>(id);
580                        }
581                    },
582                )
583            }
584        })
585    }
586
587    /// This is a queued version of [`ComponentsRegistrator::register_component_with_descriptor`].
588    /// This will reserve an id and queue the registration.
589    /// These registrations will be carried out at the next opportunity.
590    ///
591    /// # Note
592    ///
593    /// Technically speaking, the returned [`ComponentId`] is not valid, but it will become valid later.
594    /// See type level docs for details.
595    #[inline]
596    pub fn queue_register_component_with_descriptor(
597        &self,
598        descriptor: ComponentDescriptor,
599    ) -> ComponentId {
600        self.register_arbitrary_dynamic(descriptor, |registrator, id, descriptor| {
601            // SAFETY: Id uniqueness handled by caller.
602            unsafe {
603                registrator
604                    .components
605                    .register_component_inner(id, descriptor);
606            }
607        })
608    }
609
610    /// This is a queued version of [`ComponentsRegistrator::register_resource`].
611    /// This will reserve an id and queue the registration.
612    /// These registrations will be carried out at the next opportunity.
613    ///
614    /// If this has already been registered or queued, this returns the previous [`ComponentId`].
615    ///
616    /// # Note
617    ///
618    /// Technically speaking, the returned [`ComponentId`] is not valid, but it will become valid later.
619    /// See type level docs for details.
620    #[inline]
621    pub fn queue_register_resource<T: Resource>(&self) -> ComponentId {
622        let type_id = TypeId::of::<T>();
623        self.get_resource_id(type_id).unwrap_or_else(|| {
624            // SAFETY: We just checked that this type was not already registered.
625            unsafe {
626                self.register_arbitrary_resource(
627                    type_id,
628                    ComponentDescriptor::new_resource::<T>(),
629                    move |registrator, id, descriptor| {
630                        // 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.
631                        // SAFETY: Id uniqueness handled by caller, and the type_id matches descriptor.
632                        #[expect(unused_unsafe, reason = "More precise to specify.")]
633                        unsafe {
634                            registrator
635                                .components
636                                .register_resource_unchecked(type_id, id, descriptor);
637                        }
638                    },
639                )
640            }
641        })
642    }
643
644    /// This is a queued version of [`ComponentsRegistrator::register_non_send`].
645    /// This will reserve an id and queue the registration.
646    /// These registrations will be carried out at the next opportunity.
647    ///
648    /// If this has already been registered or queued, this returns the previous [`ComponentId`].
649    ///
650    /// # Note
651    ///
652    /// Technically speaking, the returned [`ComponentId`] is not valid, but it will become valid later.
653    /// See type level docs for details.
654    #[inline]
655    pub fn queue_register_non_send<T: Any>(&self) -> ComponentId {
656        let type_id = TypeId::of::<T>();
657        self.get_resource_id(type_id).unwrap_or_else(|| {
658            // SAFETY: We just checked that this type was not already registered.
659            unsafe {
660                self.register_arbitrary_resource(
661                    type_id,
662                    ComponentDescriptor::new_non_send::<T>(StorageType::default()),
663                    move |registrator, id, descriptor| {
664                        // 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.
665                        // SAFETY: Id uniqueness handled by caller, and the type_id matches descriptor.
666                        #[expect(unused_unsafe, reason = "More precise to specify.")]
667                        unsafe {
668                            registrator
669                                .components
670                                .register_resource_unchecked(type_id, id, descriptor);
671                        }
672                    },
673                )
674            }
675        })
676    }
677
678    /// This is a queued version of [`ComponentsRegistrator::register_resource_with_descriptor`].
679    /// This will reserve an id and queue the registration.
680    /// These registrations will be carried out at the next opportunity.
681    ///
682    /// # Note
683    ///
684    /// Technically speaking, the returned [`ComponentId`] is not valid, but it will become valid later.
685    /// See type level docs for details.
686    #[inline]
687    pub fn queue_register_resource_with_descriptor(
688        &self,
689        descriptor: ComponentDescriptor,
690    ) -> ComponentId {
691        self.register_arbitrary_dynamic(descriptor, |registrator, id, descriptor| {
692            // SAFETY: Id uniqueness handled by caller.
693            unsafe {
694                registrator
695                    .components
696                    .register_component_inner(id, descriptor);
697            }
698        })
699    }
700}