bevy_ecs/world/
spawn_batch.rs

1use bevy_ptr::move_as_ptr;
2
3use crate::{
4    bundle::{Bundle, BundleSpawner, NoBundleEffect},
5    change_detection::MaybeLocation,
6    entity::{Entity, EntitySetIterator},
7    world::World,
8};
9use core::iter::FusedIterator;
10
11/// An iterator that spawns a series of entities and returns the [ID](Entity) of
12/// each spawned entity.
13///
14/// If this iterator is not fully exhausted, any remaining entities will be spawned when this type is dropped.
15pub struct SpawnBatchIter<'w, I>
16where
17    I: Iterator,
18    I::Item: Bundle<Effect: NoBundleEffect>,
19{
20    inner: I,
21    spawner: BundleSpawner<'w>,
22    caller: MaybeLocation,
23}
24
25impl<'w, I> SpawnBatchIter<'w, I>
26where
27    I: Iterator,
28    I::Item: Bundle<Effect: NoBundleEffect>,
29{
30    #[inline]
31    #[track_caller]
32    pub(crate) fn new(world: &'w mut World, iter: I, caller: MaybeLocation) -> Self {
33        // Ensure all entity allocations are accounted for so `self.entities` can realloc if
34        // necessary
35        world.flush();
36
37        let change_tick = world.change_tick();
38
39        let (lower, upper) = iter.size_hint();
40        let length = upper.unwrap_or(lower);
41        world.entities.reserve(length as u32);
42
43        let mut spawner = BundleSpawner::new::<I::Item>(world, change_tick);
44        spawner.reserve_storage(length);
45
46        Self {
47            inner: iter,
48            spawner,
49            caller,
50        }
51    }
52}
53
54impl<I> Drop for SpawnBatchIter<'_, I>
55where
56    I: Iterator,
57    I::Item: Bundle<Effect: NoBundleEffect>,
58{
59    fn drop(&mut self) {
60        // Iterate through self in order to spawn remaining bundles.
61        for _ in &mut *self {}
62        // Apply any commands from those operations.
63        // SAFETY: `self.spawner` will be dropped immediately after this call.
64        unsafe { self.spawner.flush_commands() };
65    }
66}
67
68impl<I> Iterator for SpawnBatchIter<'_, I>
69where
70    I: Iterator,
71    I::Item: Bundle<Effect: NoBundleEffect>,
72{
73    type Item = Entity;
74
75    fn next(&mut self) -> Option<Entity> {
76        let bundle = self.inner.next()?;
77        move_as_ptr!(bundle);
78        // SAFETY:
79        // - The spawner matches `I::Item`'s type.
80        // - `I::Item::Effect: NoBundleEffect`, thus [`apply_effect`] does not need to be called.
81        // - `bundle` is not accessed or dropped after this function call.
82        unsafe { Some(self.spawner.spawn::<I::Item>(bundle, self.caller)) }
83    }
84
85    fn size_hint(&self) -> (usize, Option<usize>) {
86        self.inner.size_hint()
87    }
88}
89
90impl<I, T> ExactSizeIterator for SpawnBatchIter<'_, I>
91where
92    I: ExactSizeIterator<Item = T>,
93    T: Bundle<Effect: NoBundleEffect>,
94{
95    fn len(&self) -> usize {
96        self.inner.len()
97    }
98}
99
100impl<I, T> FusedIterator for SpawnBatchIter<'_, I>
101where
102    I: FusedIterator<Item = T>,
103    T: Bundle<Effect: NoBundleEffect>,
104{
105}
106
107// SAFETY: Newly spawned entities are unique.
108unsafe impl<I: Iterator, T> EntitySetIterator for SpawnBatchIter<'_, I>
109where
110    I: FusedIterator<Item = T>,
111    T: Bundle<Effect: NoBundleEffect>,
112{
113}