diff --git a/crates/bevy_ecs/src/entity/mod.rs b/crates/bevy_ecs/src/entity/mod.rs index f4a06af5d4084..773fa94d34adc 100644 --- a/crates/bevy_ecs/src/entity/mod.rs +++ b/crates/bevy_ecs/src/entity/mod.rs @@ -121,7 +121,7 @@ use crate::{ storage::{SparseSetIndex, TableId, TableRow}, }; use alloc::vec::Vec; -use core::{fmt, hash::Hash, mem, num::NonZero, panic::Location}; +use core::{fmt, hash::Hash, mem, num::NonZero, ops::Range, panic::Location}; use derive_more::derive::Display; use log::warn; use nonmax::NonMaxU32; @@ -708,9 +708,21 @@ pub struct EntityAllocator { } impl EntityAllocator { + /// Creates a new `EntityAllocator` with a given range + /// + /// # Warning + /// - Any two entity allocators must not have overlapping ranges, otherwise an entity can be allocated twice by different allocators. + /// - Do not free entities with an [`EntityIndex`] outside of the given range. This ensures that the allocator won't allocate out-of-range entities. + pub(crate) fn new(range: Range) -> Self { + Self { + inner: remote_allocator::Allocator::new(range), + } + } + /// Restarts the allocator. pub(crate) fn restart(&mut self) { - self.inner = remote_allocator::Allocator::new(); + let range = self.inner.range().clone(); + self.inner = remote_allocator::Allocator::new(range); } /// Builds a new remote allocator that hooks into this [`EntityAllocator`]. @@ -729,6 +741,10 @@ impl EntityAllocator { /// Freeing an [`Entity`] such that one [`EntityIndex`] is in the allocator in multiple places can cause panics when spawning the allocated entity. /// Additionally, to differentiate versions of an [`Entity`], updating the [`EntityGeneration`] before freeing is a good idea /// (but not strictly necessary if you don't mind [`Entity`] id aliasing.) + /// + /// # Warning + /// Freeing an [`Entity`] does not check that its [`EntityIndex`] is in range. If an out-of-range entity is freed, it may be handed out again. + /// If you rely on entities being in range, do not free out-of-range entities. pub fn free(&mut self, freed: Entity) { self.inner.free(freed); } @@ -779,7 +795,13 @@ impl EntityAllocator { /// More generally, manually spawning and [`despawn_no_free`](crate::world::World::despawn_no_free)ing entities allows you to skip Bevy's default entity allocator. /// This is useful if you want to enforce properties about the [`EntityIndex`]s of a group of entities, make a custom allocator, etc. pub fn alloc(&self) -> Entity { - self.inner.alloc() + self.inner.try_alloc().expect("out of entities") + } + + /// Allocates some [`Entity`]. + /// Returns `None` if no entities are available. This is a non-`panic`ing version of `alloc`. + pub fn try_alloc(&self) -> Option { + self.inner.try_alloc() } /// A more efficient way of calling [`alloc`](Self::alloc) repeatedly `count` times. @@ -1499,7 +1521,7 @@ mod tests { #[test] fn allocator() { - let mut allocator = EntityAllocator::default(); + let mut allocator = EntityAllocator::new(0..u32::MAX); let mut entities = allocator.alloc_many(2048).collect::>(); for _ in 0..2048 { entities.push(allocator.alloc()); diff --git a/crates/bevy_ecs/src/entity/remote_allocator.rs b/crates/bevy_ecs/src/entity/remote_allocator.rs index 15a3239b30b03..df6b75be68abf 100644 --- a/crates/bevy_ecs/src/entity/remote_allocator.rs +++ b/crates/bevy_ecs/src/entity/remote_allocator.rs @@ -39,7 +39,7 @@ use bevy_platform::{ Arc, }, }; -use core::mem::ManuallyDrop; +use core::{mem::ManuallyDrop, ops::Range}; use log::warn; use nonmax::NonMaxU32; @@ -296,7 +296,7 @@ impl FreeBuffer { /// making safety for other operations afterward need careful justification. /// Otherwise, the compiler will make unsound optimizations. #[inline] - unsafe fn iter(&self, indices: core::ops::Range) -> FreeBufferIterator<'_> { + unsafe fn iter(&self, indices: Range) -> FreeBufferIterator<'_> { FreeBufferIterator { buffer: self, future_buffer_indices: indices, @@ -325,7 +325,7 @@ struct FreeBufferIterator<'a> { /// The part of the buffer we are iterating at the moment. current_chunk_slice: core::slice::Iter<'a, Slot>, /// The indices in the buffer that are not yet in `current_chunk_slice`. - future_buffer_indices: core::ops::Range, + future_buffer_indices: Range, } impl<'a> Iterator for FreeBufferIterator<'a> { @@ -737,12 +737,16 @@ impl FreeList { struct FreshAllocator { /// The next value of [`Entity::index`] to give out if needed. next_entity_index: AtomicU32, + max_index: u32, } impl FreshAllocator { - /// This exists because it may possibly change depending on platform. - /// Ex: We may want this to be smaller on 32 bit platforms at some point. - const MAX_ENTITIES: u32 = u32::MAX; + pub(crate) fn new(range: Range) -> Self { + Self { + next_entity_index: AtomicU32::new(range.start), + max_index: range.end, + } + } /// The total number of indices given out. #[inline] @@ -759,15 +763,24 @@ impl FreshAllocator { } /// Allocates a fresh [`EntityIndex`]. - /// This row has never been given out before. + /// This index has never been given out before. + /// If no index is available (out of range), than it returns None #[inline] - fn alloc(&self) -> Entity { + fn alloc(&self) -> Option { let index = self.next_entity_index.fetch_add(1, Ordering::Relaxed); - if index == Self::MAX_ENTITIES { - Self::on_overflow(); + if index >= self.max_index { + self.next_entity_index + .store(self.max_index, Ordering::Relaxed); + if index == u32::MAX { + Self::on_overflow(); + } + return None; } + // SAFETY: We just checked that this was not max and we only added 1, so we can't have missed it. - Entity::from_index(unsafe { EntityIndex::new(NonMaxU32::new_unchecked(index)) }) + Some(Entity::from_index(unsafe { + EntityIndex::new(NonMaxU32::new_unchecked(index)) + })) } /// Allocates `count` [`EntityIndex`]s. @@ -777,7 +790,7 @@ impl FreshAllocator { let start_new = self.next_entity_index.fetch_add(count, Ordering::Relaxed); let new = match start_new .checked_add(count) - .filter(|new| *new < Self::MAX_ENTITIES) + .filter(|new| *new < self.max_index) { Some(new_next_entity_index) => start_new..new_next_entity_index, None => Self::on_overflow(), @@ -790,7 +803,7 @@ impl FreshAllocator { /// These rows have never been given out before. /// /// **NOTE:** Dropping will leak the remaining entity rows! -pub(super) struct AllocUniqueEntityIndexIterator(core::ops::Range); +pub(super) struct AllocUniqueEntityIndexIterator(Range); impl Iterator for AllocUniqueEntityIndexIterator { type Item = Entity; @@ -824,25 +837,24 @@ struct SharedAllocator { impl SharedAllocator { /// Constructs a [`SharedAllocator`] - fn new() -> Self { + fn new(range: Range) -> Self { Self { free: FreeList::new(), - fresh: FreshAllocator { - next_entity_index: AtomicU32::new(0), - }, + fresh: FreshAllocator::new(range), is_closed: AtomicBool::new(false), } } /// Allocates a new [`Entity`], reusing a freed index if one exists. + /// If no more entities can be allocated, this returns None /// /// # Safety /// /// This must not conflict with [`FreeList::free`] calls. #[inline] - unsafe fn alloc(&self) -> Entity { + unsafe fn try_alloc(&self) -> Option { // SAFETY: assured by caller - unsafe { self.free.alloc() }.unwrap_or_else(|| self.fresh.alloc()) + unsafe { self.free.alloc() }.or_else(|| self.fresh.alloc()) } /// Allocates a `count` [`Entity`]s, reusing freed indices if they exist. @@ -862,10 +874,8 @@ impl SharedAllocator { /// Allocates a new [`Entity`]. /// This will only try to reuse a freed index if it is safe to do so. #[inline] - fn remote_alloc(&self) -> Entity { - self.free - .remote_alloc() - .unwrap_or_else(|| self.fresh.alloc()) + fn try_remote_alloc(&self) -> Option { + self.free.remote_alloc().or_else(|| self.fresh.alloc()) } /// Marks the allocator as closed, but it will still function normally. @@ -891,28 +901,43 @@ pub(crate) struct Allocator { /// The local free list. /// We use this to amortize the cost of freeing to the shared allocator since that is expensive. local_free: Box>, + /// The index range this allocator operates on + range: Range, } impl Default for Allocator { fn default() -> Self { - Self::new() + Self::new(0..u32::MAX) } } impl Allocator { /// Constructs a new [`Allocator`] - pub(super) fn new() -> Self { + pub(super) fn new(range: Range) -> Self { Self { - shared: Arc::new(SharedAllocator::new()), + shared: Arc::new(SharedAllocator::new(range.clone())), local_free: Box::new(ArrayVec::new()), + range, } } + /// Returns the range this allocator operates on + pub(super) fn range(&self) -> &Range { + &self.range + } + /// Allocates a new [`Entity`], reusing a freed index if one exists. + #[cfg(test)] + fn alloc(&self) -> Entity { + self.try_alloc().expect("out of entities") + } + + /// Allocates a new [`Entity`], reusing a freed index if one exists. + /// Returns None if no entities are available within the range #[inline] - pub(super) fn alloc(&self) -> Entity { + pub(super) fn try_alloc(&self) -> Option { // SAFETY: violating safety requires a `&mut self` to exist, but rust does not allow that. - unsafe { self.shared.alloc() } + unsafe { self.shared.try_alloc() } } /// The total number of indices given out. @@ -1049,7 +1074,7 @@ impl RemoteAllocator { Arc::ptr_eq(&self.shared, &source.shared) } - /// Allocates an entity remotely. + /// Allocates an entity remotely, panicking if no more entities are available. /// /// This comes with a major downside: /// Because this does not hold reference to the world, the world may be cleared or destroyed before you get a chance to use the result. @@ -1058,7 +1083,19 @@ impl RemoteAllocator { /// Before using the returned values in the world, first check that it is ok with [`EntityAllocator::has_remote_allocator`](super::EntityAllocator::has_remote_allocator). #[inline] pub fn alloc(&self) -> Entity { - self.shared.remote_alloc() + self.shared.try_remote_alloc().expect("out of entities") + } + + /// Allocates an entity remotely, returning `None` if no more entities are available. + /// + /// This comes with a major downside: + /// Because this does not hold reference to the world, the world may be cleared or destroyed before you get a chance to use the result. + /// If that happens, these entities will be garbage! + /// They will not be unique in the world anymore and you should not spawn them! + /// Before using the returned values in the world, first check that it is ok with [`EntityAllocator::has_remote_allocator`](super::EntityAllocator::has_remote_allocator). + #[inline] + pub fn try_alloc(&self) -> Option { + self.shared.try_remote_alloc() } /// Returns whether or not this [`RemoteAllocator`] is still connected to its source [`EntityAllocator`](super::EntityAllocator). @@ -1133,7 +1170,7 @@ mod tests { #[test] fn uniqueness() { let mut entities = Vec::with_capacity(2000); - let mut allocator = Allocator::new(); + let mut allocator = Allocator::default(); entities.extend(allocator.alloc_many(1000)); let pre_len = entities.len(); @@ -1159,7 +1196,7 @@ mod tests { /// This test just exists to make sure allocations don't step on each other's toes. #[test] fn allocation_order_correctness() { - let mut allocator = Allocator::new(); + let mut allocator = Allocator::default(); let e0 = allocator.alloc(); let e1 = allocator.alloc(); let e2 = allocator.alloc(); diff --git a/crates/bevy_ecs/src/world/mod.rs b/crates/bevy_ecs/src/world/mod.rs index d8f6868d35fd7..b8bbf3a5696ec 100644 --- a/crates/bevy_ecs/src/world/mod.rs +++ b/crates/bevy_ecs/src/world/mod.rs @@ -119,7 +119,7 @@ impl Default for World { let mut world = Self { id: WorldId::new().expect("More `bevy` `World`s have been created than is supported"), entities: Entities::new(), - entity_allocator: EntityAllocator::default(), + entity_allocator: EntityAllocator::new(0..u32::MAX), components: Default::default(), resource_entities: Default::default(), archetypes: Archetypes::new(),