diff --git a/crates/bevy_ecs/src/entity/mod.rs b/crates/bevy_ecs/src/entity/mod.rs index 82f8c7a64419e..1b1d0280e7765 100644 --- a/crates/bevy_ecs/src/entity/mod.rs +++ b/crates/bevy_ecs/src/entity/mod.rs @@ -94,6 +94,9 @@ use derive_more::derive::Display; pub use entity_set::*; pub use map_entities::*; +mod remote_allocator; +pub use remote_allocator::RemoteAllocator; + mod hash; pub use hash::*; @@ -124,7 +127,6 @@ use crate::{ storage::{SparseSetIndex, TableId, TableRow}, }; use alloc::vec::Vec; -use bevy_platform::sync::atomic::{AtomicU32, AtomicUsize, Ordering}; use core::{fmt, hash::Hash, mem, num::NonZero, panic::Location}; use log::warn; @@ -698,25 +700,25 @@ impl SparseSetIndex for Entity { /// See the module docs for how these ids and this allocator participate in the life cycle of an entity. #[derive(Default, Debug)] pub struct EntityAllocator { - /// All the entities to reuse. - /// This is a buffer, which contains an array of [`Entity`] ids to hand out. - /// The next id to hand out is tracked by `free_len`. - free: Vec, - /// This is continually subtracted from. - /// If it wraps to a very large number, it will be outside the bounds of `free`, - /// and a new index will be needed. - free_len: AtomicUsize, - /// This is the next "fresh" index to hand out. - /// If there are no indices to reuse, this index, which has a generation of 0, is the next to return. - next_index: AtomicU32, + inner: remote_allocator::Allocator, } impl EntityAllocator { /// Restarts the allocator. pub(crate) fn restart(&mut self) { - self.free.clear(); - *self.free_len.get_mut() = 0; - *self.next_index.get_mut() = 0; + self.inner = remote_allocator::Allocator::new(); + } + + /// Builds a new remote allocator that hooks into this [`EntityAllocator`]. + /// This is useful when you need to allocate entities without holding a reference to the world (like in async). + pub fn build_remote_allocator(&mut self) -> RemoteAllocator { + RemoteAllocator::new(&self.inner) + } + + /// Returns `true` when the `allocator` is connected to this [`EntityAllocator`] + /// and its allocated [`Entity`] values can still be used in this world. + pub fn has_remote_allocator(&self, allocator: &RemoteAllocator) -> bool { + allocator.is_connected_to(&self.inner) } /// This allows `freed` to be retrieved from [`alloc`](Self::alloc), etc. @@ -724,14 +726,7 @@ impl EntityAllocator { /// 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.) pub fn free(&mut self, freed: Entity) { - let expected_len = *self.free_len.get_mut(); - if expected_len > self.free.len() { - self.free.clear(); - } else { - self.free.truncate(expected_len); - } - self.free.push(freed); - *self.free_len.get_mut() = self.free.len(); + self.inner.free(freed); } /// Allocates some [`Entity`]. @@ -772,15 +767,7 @@ 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 { - let index = self - .free_len - .fetch_sub(1, Ordering::Relaxed) - .wrapping_sub(1); - self.free.get(index).copied().unwrap_or_else(|| { - let index = self.next_index.fetch_add(1, Ordering::Relaxed); - let index = NonMaxU32::new(index).expect("too many entities"); - Entity::from_index(EntityIndex::new(index)) - }) + self.inner.alloc() } /// A more efficient way of calling [`alloc`](Self::alloc) repeatedly `count` times. @@ -790,27 +777,8 @@ impl EntityAllocator { /// If the iterator is not exhausted, its remaining entities are forgotten. /// See [`AllocEntitiesIterator`] docs for more. pub fn alloc_many(&self, count: u32) -> AllocEntitiesIterator<'_> { - let current_len = self.free_len.fetch_sub(count as usize, Ordering::Relaxed); - let current_len = if current_len < self.free.len() { - current_len - } else { - 0 - }; - let start = current_len.saturating_sub(count as usize); - let reuse = start..current_len; - let still_need = (count as usize - reuse.len()) as u32; - let new = if still_need > 0 { - let start_new = self.next_index.fetch_add(still_need, Ordering::Relaxed); - let end_new = start_new - .checked_add(still_need) - .expect("too many entities"); - start_new..end_new - } else { - 0..0 - }; AllocEntitiesIterator { - reuse: self.free[reuse].iter(), - new, + inner: self.inner.alloc_many(count), } } } @@ -819,26 +787,18 @@ impl EntityAllocator { /// Dropping this will still retain the entities as allocated; this is effectively a leak. /// To prevent this, ensure the iterator is exhausted before dropping it. pub struct AllocEntitiesIterator<'a> { - reuse: core::slice::Iter<'a, Entity>, - new: core::ops::Range, + inner: remote_allocator::AllocEntitiesIterator<'a>, } impl<'a> Iterator for AllocEntitiesIterator<'a> { type Item = Entity; fn next(&mut self) -> Option { - self.reuse.next().copied().or_else(|| { - self.new.next().map(|index| { - // SAFETY: This came from an exclusive range so the max can't be hit. - let index = unsafe { EntityIndex::new(NonMaxU32::new_unchecked(index)) }; - Entity::from_index(index) - }) - }) + self.inner.next() } fn size_hint(&self) -> (usize, Option) { - let len = self.reuse.len() + self.new.len(); - (len, Some(len)) + self.inner.size_hint() } } diff --git a/crates/bevy_ecs/src/entity/remote_allocator.rs b/crates/bevy_ecs/src/entity/remote_allocator.rs new file mode 100644 index 0000000000000..a5b1cd29691b7 --- /dev/null +++ b/crates/bevy_ecs/src/entity/remote_allocator.rs @@ -0,0 +1,1052 @@ +//! This module contains the guts of Bevy's entity allocator. +//! +//! Entity allocation needs to work concurrently and remotely. +//! Remote allocations (where no reference to the world is held) is needed for async primarily for use. +//! Non-remote, "normal" allocation needs to be as fast as possible while still supporting remote allocation. +//! +//! The allocator fundamentally is made of a cursor for the next fresh, never used [`EntityIndex`] and a free list. +//! The free list is some collection that hold [`Entity`] values that were used and can be reused; they are "free"/available. +//! If the free list is empty, it's really simple to just increment the fresh index cursor. +//! The tricky part is implementing a remotely accessible free list. +//! +//! A naive free list could just a concurrent queue. +//! That would probably be fine for remote allocation but for non-remote, we can go much faster. +//! In particular, a concurrent queue must do additional work to handle cases where something is added concurrently with being removed. +//! But for non-remote allocation, we can guarantee that no free will happen during an allocation since `free` needs mutably access to the world already. +//! That means, we can skip a lot of those safety checks. +//! Plus, we know the maximum size of the free list ahead of time, since we can assume there are no duplicates. +//! That means, we can have a much more efficient allocation scheme, far better than a linked list. +//! +//! For the free list, the list needs to be pinned in memory and yet grow-able. +//! That's quire the pickle, but by splitting the growth over multiple arrays, this isn't so bad. +//! When the list needs to grow, we just *add* on another array to the buffer (instead of *replacing* the old one with a bigger one). +//! These arrays are called [`Chunk`]s. +//! This keeps everything pinned, and since we know the maximum size ahead of time, we can make this mapping very fast. +//! +//! Similar to how `Vec` is implemented, the free list is implemented as a [`FreeBuffer`] (handling allocations and implicit capacity) +//! and the [`FreeCount`] manages the length of the free list. +//! The free list's item is a [`Slot`], which manages accessing each item concurrently. +//! +//! These types are summed up in [`SharedAllocator`], which is highly unsafe. +//! The interfaces [`Allocator`] and [`RemoteAllocator`] provide safe interfaces to them. + +use bevy_platform::{ + cell::SyncUnsafeCell, + prelude::Vec, + sync::{ + atomic::{AtomicBool, AtomicPtr, AtomicU32, AtomicU64, Ordering}, + Arc, + }, +}; +use core::mem::ManuallyDrop; +use log::warn; +use nonmax::NonMaxU32; + +use crate::query::DebugCheckedUnwrap; + +use super::{Entity, EntityIndex, EntitySetIterator}; + +/// This is the item we store in the free list. +/// Effectively, this is a `MaybeUninit` where uninit is represented by `Entity::PLACEHOLDER`. +/// This is +struct Slot { + inner: SyncUnsafeCell, +} + +impl Slot { + /// Produces a meaningless empty value. This is a valid but incorrect `Entity`. + /// It's valid because the bits do represent a valid bit pattern of an `Entity`. + /// It's incorrect because this is in the free buffer even though the entity was never freed. + /// Importantly, [`FreeCount`] determines which part of the free buffer is the free list. + /// An empty slot may be in the free buffer, but should not be in the free list. + /// This can be thought of as the `MaybeUninit` uninit in `Vec`'s excess capacity. + const fn empty() -> Self { + let source = Entity::PLACEHOLDER; + Self { + inner: SyncUnsafeCell::new(source), + } + } + + /// Sets the entity at this slot. + /// + /// # Safety + /// + /// There must be a clear, strict order between this call and the previous uses of this [`Slot`]. + /// Otherwise, the compiler will make unsound optimizations. + #[inline] + const unsafe fn set_entity(&self, entity: Entity) { + // SAFETY: Ensured by caller. + unsafe { + self.inner.get().write(entity); + } + } + + /// Gets the stored entity. The result will be [`Entity::PLACEHOLDER`] unless [`set_entity`](Self::set_entity) has been called. + /// + /// # Safety + /// + /// There must be a clear, strict order between this call and the previous uses of this [`Slot`]. + /// Otherwise, the compiler will make unsound optimizations. + #[inline] + const unsafe fn get_entity(&self) -> Entity { + // SAFETY: Ensured by caller. + unsafe { self.inner.get().read() } + } +} + +/// Each chunk stores a buffer of [`Slot`]s at a fixed capacity. +struct Chunk { + /// Points to the first slot. If this is null, we need to allocate it. + first: AtomicPtr, +} + +impl Chunk { + /// Constructs a null [`Chunk`]. + const fn new() -> Self { + Self { + first: AtomicPtr::new(core::ptr::null_mut()), + } + } + + /// Gets the entity at the index within this chunk. + /// + /// # Safety + /// + /// [`Self::set`] must have been called on this index before, ensuring it is in bounds and the chunk is initialized. + /// There must be a clear, strict order between this call and the previous uses of this `index`. + /// Otherwise, the compiler will make unsound optimizations. + #[inline] + unsafe fn get(&self, index: u32) -> Entity { + // Relaxed is fine since caller has already assured memory ordering is satisfied. + let head = self.first.load(Ordering::Relaxed); + // SAFETY: caller ensures we are in bounds and init (because `set` must be in bounds) + let target = unsafe { &*head.add(index as usize) }; + // SAFETY: Caller ensures ordering. + unsafe { target.get_entity() } + } + + /// Gets a slice of indices. + /// + /// # Safety + /// + /// [`Self::set`] must have been called on these indices before, ensuring it is in bounds and the chunk is initialized. + #[inline] + unsafe fn get_slice(&self, index: u32, ideal_len: u32, chunk_capacity: u32) -> &[Slot] { + let after_index_slice_len = chunk_capacity - index; + let len = after_index_slice_len.min(ideal_len) as usize; + + // Relaxed is fine since caller ensures we are initialized already. + // In order for the caller to guarantee that, they must have an ordering that orders this `get` after the required `set`. + let head = self.first.load(Ordering::Relaxed); + + // SAFETY: Caller ensures we are init, so the chunk was allocated via a `Vec` and the index is within the capacity. + unsafe { core::slice::from_raw_parts(head, len) } + } + + /// Sets this entity at this index. + /// + /// # Safety + /// + /// Index must be in bounds. + /// There must be a clear, strict order between this call and the previous uses of this `index`. + /// Otherwise, the compiler will make unsound optimizations. + #[inline] + unsafe fn set(&self, index: u32, entity: Entity, chunk_capacity: u32) { + // Relaxed is fine here since the caller ensures memory ordering. + let ptr = self.first.load(Ordering::Relaxed); + let head = if ptr.is_null() { + self.init(chunk_capacity) + } else { + ptr + }; + + // SAFETY: caller ensures it is in bounds and we are not fighting with other `set` calls or `get` calls. + // A race condition is therefore impossible. + // The address can't wrap or pass isize max since this addition is within an allocation. + // For that to happen, you would first run out of memory in practice. + let target = unsafe { &*head.add(index as usize) }; + + target.set_entity(entity); + } + + /// Initializes the chunk to be valid, returning the pointer. + /// + /// # Safety + /// + /// This must not be called concurrently with itself. + #[cold] + unsafe fn init(&self, chunk_capacity: u32) -> *mut Slot { + let mut buff = ManuallyDrop::new(Vec::new()); + buff.reserve_exact(chunk_capacity as usize); + buff.resize_with(chunk_capacity as usize, Slot::empty); + let ptr = buff.as_mut_ptr(); + // Relaxed is fine here since this is not called concurrently. + self.first.store(ptr, Ordering::Relaxed); + ptr + } + + /// Frees memory + /// + /// # Safety + /// + /// `chunk_capacity` must be the same as it was initialized with. + unsafe fn dealloc(&mut self, chunk_capacity: u32) { + let to_drop = *self.first.get_mut(); + if !to_drop.is_null() { + // SAFETY: This was created in [`Self::init`] from a standard Vec. + unsafe { + Vec::from_raw_parts(to_drop, chunk_capacity as usize, chunk_capacity as usize); + } + } + } +} + +/// This is a buffer that has been split into power-of-two sized chunks, so that each chunk is pinned in memory. +/// Conceptually, each chunk is put end-to-end to form the buffer. This ultimately avoids copying elements on resize, +/// while allowing it to expand in capacity as needed. A separate system must track the length of the list in the buffer. +/// Each chunk is twice as large as the last, except for the first two which have a capacity of 512. +struct FreeBuffer([Chunk; Self::NUM_CHUNKS as usize]); + +impl FreeBuffer { + const NUM_CHUNKS: u32 = 24; + const NUM_SKIPPED: u32 = u32::BITS - Self::NUM_CHUNKS; + + /// Constructs an empty [`FreeBuffer`]. + const fn new() -> Self { + Self([const { Chunk::new() }; Self::NUM_CHUNKS as usize]) + } + + /// Computes the capacity of the chunk at this index within [`Self::NUM_CHUNKS`]. + /// The first 2 have length 512 (2^9) and the last has length (2^31) + #[inline] + const fn capacity_of_chunk(chunk_index: u32) -> u32 { + // We do this because we're skipping the first `NUM_SKIPPED` powers, so we need to make up for them by doubling the first index. + // This is why the first 2 indices both have a capacity of 512. + let corrected = if chunk_index == 0 { 1 } else { chunk_index }; + // We add NUM_SKIPPED because the total capacity should be as if [`Self::NUM_CHUNKS`] were 32. + // This skips the first NUM_SKIPPED powers. + let corrected = corrected + Self::NUM_SKIPPED; + // This bit shift is just 2^corrected. + 1 << corrected + } + + /// For this index in the whole buffer, returns the index of the [`Chunk`], the index within that chunk, and the capacity of that chunk. + #[inline] + const fn index_info(full_index: u32) -> (u32, u32, u32) { + // We do a `saturating_sub` because we skip the first `NUM_SKIPPED` powers to make space for the first chunk's entity count. + // The -1 is because this is the number of chunks, but we want the index in the end. + // We store chunks in smallest to biggest order, so we need to reverse it. + let chunk_index = (Self::NUM_CHUNKS - 1).saturating_sub(full_index.leading_zeros()); + let chunk_capacity = Self::capacity_of_chunk(chunk_index); + // We only need to cut off this particular bit. + // The capacity is only one bit, and if other bits needed to be dropped, `leading` would have been greater + let index_in_chunk = full_index & !chunk_capacity; + + (chunk_index, index_in_chunk, chunk_capacity) + } + + /// For this index in the whole buffer, returns the [`Chunk`], the index within that chunk, and the capacity of that chunk. + #[inline] + fn index_in_chunk(&self, full_index: u32) -> (&Chunk, u32, u32) { + let (chunk_index, index_in_chunk, chunk_capacity) = Self::index_info(full_index); + // SAFETY: The `index_info` is correct. + let chunk = unsafe { self.0.get_unchecked(chunk_index as usize) }; + (chunk, index_in_chunk, chunk_capacity) + } + + /// Gets the entity at an index. + /// + /// # Safety + /// + /// [`set`](Self::set) must have been called on this index to initialize the its memory. + /// There must be a clear, strict order between this call and the previous uses of this `full_index`. + /// Otherwise, the compiler will make unsound optimizations. + unsafe fn get(&self, full_index: u32) -> Entity { + let (chunk, index, _) = self.index_in_chunk(full_index); + // SAFETY: Ensured by caller. + unsafe { chunk.get(index) } + } + + /// Sets an entity at an index. + /// + /// # Safety + /// + /// There must be a clear, strict order between this call and the previous uses of this `full_index`. + /// Otherwise, the compiler will make unsound optimizations. + #[inline] + unsafe fn set(&self, full_index: u32, entity: Entity) { + let (chunk, index, chunk_capacity) = self.index_in_chunk(full_index); + // SAFETY: Ensured by caller and that the index is correct. + unsafe { chunk.set(index, entity, chunk_capacity) } + } + + /// Iterates the entities in these indices. + /// + /// # Safety + /// + /// [`Self::set`] must have been called on these indices before to initialize memory. + /// There must be a clear, strict order between this call and the previous uses of these `indices`. + /// Note that until the returned value is dropped, these `indices` are still being accessed, + /// 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<'_> { + FreeBufferIterator { + buffer: self, + future_buffer_indices: indices, + current_chunk_slice: [].iter(), + } + } +} + +impl Drop for FreeBuffer { + fn drop(&mut self) { + for index in 0..Self::NUM_CHUNKS { + let capacity = Self::capacity_of_chunk(index); + // SAFETY: we have `&mut` and the capacity is correct. + unsafe { self.0[index as usize].dealloc(capacity) }; + } + } +} + +/// An iterator over a [`FreeBuffer`]. +/// +/// # Safety +/// +/// [`FreeBuffer::set`] must have been called on these indices beforehand to initialize memory. +struct FreeBufferIterator<'a> { + buffer: &'a FreeBuffer, + /// 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, +} + +impl<'a> Iterator for FreeBufferIterator<'a> { + type Item = Entity; + + #[inline] + fn next(&mut self) -> Option { + if let Some(found) = self.current_chunk_slice.next() { + // SAFETY: We have `&mut self`, so that memory order is certain. + // The caller of `FreeBuffer::iter` ensures the memory order of this value's lifetime. + return Some(unsafe { found.get_entity() }); + } + + let still_need = self.future_buffer_indices.len() as u32; + let next_index = self.future_buffer_indices.next()?; + let (chunk, index, chunk_capacity) = self.buffer.index_in_chunk(next_index); + + // SAFETY: Assured by `FreeBuffer::iter` + let slice = unsafe { chunk.get_slice(index, still_need, chunk_capacity) }; + self.future_buffer_indices.start += slice.len() as u32; + self.current_chunk_slice = slice.iter(); + + // SAFETY: Constructor ensures these indices are valid in the buffer; the buffer is not sparse, and we just got the next slice. + // So the only way for the slice to be empty is if the constructor did not uphold safety. + let next = unsafe { self.current_chunk_slice.next().debug_checked_unwrap() }; + // SAFETY: We have `&mut self`, so that memory order is certain. + // The caller of `FreeBuffer::iter` ensures the memory order of this value's lifetime. + Some(unsafe { next.get_entity() }) + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + let len = self.future_buffer_indices.len() + self.current_chunk_slice.len(); + (len, Some(len)) + } +} + +impl<'a> ExactSizeIterator for FreeBufferIterator<'a> {} +impl<'a> core::iter::FusedIterator for FreeBufferIterator<'a> {} + +/// This tracks the state of a [`FreeCount`], which has lots of information packed into it. +/// +/// - The least 33 bits store a signed 33 bit integer. This behaves like a u33, but we define `1 << 32` as 0. +/// - The 34th bit stores a flag that indicates if the count has been disabled/suspended. +/// - The remaining 30 bits are the generation. The generation just differentiates different versions of the state that happen to encode the same length. +#[derive(Clone, Copy)] +struct FreeCountState(u64); + +impl FreeCountState { + /// When this bit is on, the count is disabled. + /// This is used to prevent remote allocations from running at the same time as a free operation. + const DISABLING_BIT: u64 = 1 << 33; + /// This is the mask for the length bits. + const LENGTH_MASK: u64 = (1 << 32) | u32::MAX as u64; + /// This is the value of the length mask we consider to be 0. + const LENGTH_0: u64 = 1 << 32; + /// This is the lowest bit in the u30 generation. + const GENERATION_LEAST_BIT: u64 = 1 << 34; + + /// Constructs a length of 0. + const fn new_zero_len() -> Self { + Self(Self::LENGTH_0) + } + + /// Gets the encoded length. + #[inline] + const fn length(self) -> u32 { + let unsigned_length = self.0 & Self::LENGTH_MASK; + unsigned_length.saturating_sub(Self::LENGTH_0) as u32 + } + + /// Returns whether or not the count is disabled. + #[inline] + const fn is_disabled(self) -> bool { + (self.0 & Self::DISABLING_BIT) > 0 + } + + /// Changes only the length of this count to `length`. + #[inline] + const fn with_length(self, length: u32) -> Self { + // Just turns on the "considered zero" bit since this is non-negative. + let length = length as u64 | Self::LENGTH_0; + Self(self.0 & !Self::LENGTH_MASK | length) + } + + /// For popping `num` off the count, subtract the resulting u64. + #[inline] + const fn encode_pop(num: u32) -> u64 { + let subtract_length = num as u64; + // Also subtract one from the generation bit. + subtract_length | Self::GENERATION_LEAST_BIT + } + + /// Returns the count after popping off `num` elements. + #[inline] + const fn pop(self, num: u32) -> Self { + Self(self.0.wrapping_sub(Self::encode_pop(num))) + } +} + +/// This is an atomic interface to [`FreeCountState`]. +struct FreeCount(AtomicU64); + +impl FreeCount { + /// Constructs a length of 0. + const fn new_zero_len() -> Self { + Self(AtomicU64::new(FreeCountState::new_zero_len().0)) + } + + /// Gets the current state of the buffer. + #[inline] + fn state(&self, order: Ordering) -> FreeCountState { + FreeCountState(self.0.load(order)) + } + + /// Subtracts `num` from the length, returning the previous state. + /// + /// **NOTE:** Caller should be careful that changing the state is allowed and that the state is not disabled. + #[inline] + fn pop_for_state(&self, num: u32, order: Ordering) -> FreeCountState { + let to_sub = FreeCountState::encode_pop(num); + let raw = self.0.fetch_sub(to_sub, order); + FreeCountState(raw) + } + + /// Marks the state as disabled, returning the previous state + /// When the length is disabled, [`try_set_state`](Self::try_set_state) will fail. + /// This is used to prevent remote allocation during a free. + #[inline] + fn disable_len_for_state(&self, order: Ordering) -> FreeCountState { + // We don't care about the generation here since this changes the value anyway. + FreeCountState(self.0.fetch_or(FreeCountState::DISABLING_BIT, order)) + } + + /// Sets the state explicitly. + /// Caller must be careful that the state has not changed since getting the state and setting it. + /// If that happens, the state may not properly reflect the length of the free list or its generation, + /// causing entities to be skipped or given out twice. + /// This is not a safety concern, but it is a major correctness concern. + #[inline] + fn set_state_risky(&self, state: FreeCountState, order: Ordering) { + self.0.store(state.0, order); + } + + /// Attempts to update the state, returning the new [`FreeCountState`] if it fails. + #[inline] + fn try_set_state( + &self, + expected_current_state: FreeCountState, + target_state: FreeCountState, + success: Ordering, + failure: Ordering, + ) -> Result<(), FreeCountState> { + self.0 + .compare_exchange(expected_current_state.0, target_state.0, success, failure) + .map(|_| ()) + .map_err(FreeCountState) + } +} + +/// This is conceptually like a `Vec` that stores entities pending reuse. +struct FreeList { + /// The actual buffer of [`Slot`]s. + /// Conceptually, this is like the `RawVec` for this `Vec`. + buffer: FreeBuffer, + /// The length of the free buffer + len: FreeCount, +} + +impl FreeList { + /// Constructs a empty [`FreeList`]. + fn new() -> Self { + Self { + buffer: FreeBuffer::new(), + len: FreeCount::new_zero_len(), + } + } + + /// Gets the number of free entities. + /// + /// # Risk + /// + /// For this to be accurate, this must not be called during a [`Self::free`]. + #[inline] + fn num_free(&self) -> u32 { + // Relaxed ordering is fine since this doesn't act on the length value in memory. + self.len.state(Ordering::Relaxed).length() + } + + /// Frees the `entity` allowing it to be reused. + /// + /// # Safety + /// + /// There must be a clear, strict order between this call and calls to [`Self::free`], [`Self::alloc_many`], and [`Self::alloc`]. + /// Otherwise, the compiler will make unsound optimizations. + #[inline] + unsafe fn free(&self, entity: Entity) { + // Disable remote allocation. + // We don't need to acquire the most recent memory from remote threads because we never read it. + // We do not need to release to remote threads because we only changed the disabled bit, + // which the remote allocator would with relaxed ordering. + let state = self.len.disable_len_for_state(Ordering::Relaxed); + + // Push onto the buffer + let len = state.length(); + // SAFETY: Caller ensures this does not conflict with `free` or `alloc` calls, + // and we just disabled remote allocation with a strict memory ordering. + unsafe { + self.buffer.set(len, entity); + } + + // Update length + let new_state = state.with_length(len + 1); + // This is safe because `alloc` is not being called and `remote_alloc` checks that it is not disabled. + // We don't need to change the generation since this will change the length, which changes the value anyway. + // If, from a `remote_alloc` perspective, this does not change the length (i.e. this changes it *back* to what it was), + // then `alloc` must have been called, which changes the generation. + self.len.set_state_risky(new_state, Ordering::Release); + } + + /// Allocates an [`Entity`] from the free list if one is available. + /// + /// # Safety + /// + /// There must be a clear, strict order between this call and calls to [`Self::free`]. + /// Otherwise, the compiler will make unsound optimizations. + #[inline] + unsafe fn alloc(&self) -> Option { + // SAFETY: This will get a valid index because caller ensures there is no way for `free` to be done at the same time. + // Relaxed is ok here since `free` is the only time memory is changed, and relaxed still gets the most recent state. + // The memory ordering to ensure we read the most recent value at the index is ensured by the caller. + let len = self.len.pop_for_state(1, Ordering::Relaxed).length(); + let index = len.checked_sub(1)?; + + // SAFETY: This was less then `len`, so it must have been `set` via `free` before. + // There is a strict memory ordering of this use of the index because the length is only decreasing. + // That means there is only one use of this index since the last call to `free`. + // The only time it the length increases is during `free`, which the caller ensures has a "happened before" relationship with this call. + Some(unsafe { self.buffer.get(index) }) + } + + /// Allocates as many [`Entity`]s from the free list as are available, up to `count`. + /// + /// # Safety + /// + /// There must be a clear, strict order between this call and calls to [`Self::free`]. + /// Otherwise, the compiler will make unsound optimizations. + /// + /// Note that this allocation call doesn't end until the returned value is dropped. + /// So, calling [`Self::free`] while the returned value is live is unsound. + #[inline] + unsafe fn alloc_many(&self, count: u32) -> FreeBufferIterator<'_> { + // SAFETY: This will get a valid index because there is no way for `free` to be done at the same time. + // Relaxed is ok here since `free` is the only time memory is changed, and relaxed still gets the most recent state. + // The memory ordering to ensure we read the most recent value at the index is ensured by the caller. + let len = self.len.pop_for_state(count, Ordering::Relaxed).length(); + let index = len.saturating_sub(count); + + // SAFETY: The iterator's items are all less than the length, so they are in bounds and have been previously set. + // There is a strict memory ordering of this use of the indices because the length is only decreasing. + // That means there is only one use of these indices since the last call to `free`. + // The only time it the length increases is during `free`, which the caller ensures has a "happened before" relationship with this call. + unsafe { self.buffer.iter(index..len) } + } + + /// Allocates an [`Entity`] from the free list if one is available and it is safe to do so. + #[inline] + fn remote_alloc(&self) -> Option { + // The goal is the same as `alloc`, so what's the difference? + // `alloc` knows `free` is not being called, but this does not. + // What if we `len.fetch_sub(1)` but then `free` overwrites the entity before we could read it? + // That would mean we would leak an entity and give another entity out twice. + // We get around this by only updating `len` after the read is complete. + // But that means something else could be trying to allocate the same index! + // So we need a `len.compare_exchange` loop to ensure the index is unique. + // Because we keep a generation value in the `FreeCount`, if any of these things happen, we simply try again. + // We also need to prevent this from conflicting with a `free` call, so we check to ensure the state is not disabled. + + // We keep track of the attempts so we can yield the thread on std after a few fails. + #[cfg(feature = "std")] + let mut attempts = 1u32; + // We need an acquire ordering to acquire the most recent memory of `free` calls. + let mut state = self.len.state(Ordering::Acquire); + loop { + // The state is only disabled when freeing. + // If a free is happening, we need to wait for the new entity to be ready on the free buffer. + // That means we will also need to re-fetch the state and acquire the new memory. + // Then, we can allocate it. + if state.is_disabled() { + // Spin 64 times before yielding. + #[cfg(feature = "std")] + { + attempts += 1; + if attempts % 64 == 0 { + // scheduler probably isn't running the thread doing the `free` call, so yield so it can finish. + std::thread::yield_now(); + } else { + core::hint::spin_loop(); + } + } + + #[cfg(not(feature = "std"))] + core::hint::spin_loop(); + + // Retry with the fresh state and acquired memory order. + state = self.len.state(Ordering::Acquire); + continue; + } + + // At this point, we know a `free` was not happening when we started. + + let len = state.length(); + let index = len.checked_sub(1)?; + + // SAFETY: + // + // If no `free` call has started, this safety follows the same logic as in non-remote `alloc`. + // That is, the len always counts down, so this is the only use of this index since the last `free`, + // and another `free` hasn't happened. + // + // But if a `free` did start at this point, it would be operating on indices greater than `index`. + // We haven't updated the `FreeCount` yet, so the `free` call would be adding to it, while we've been subtracting from it. + // That means this is still the only time this index is used since the last `free`! + // So, even though we can't guarantee when the concurrent `free` is happening in memory order, it doesn't matter since that `free` doesn't use this index. + // We can still establish a clear, strict ordering for this slot because 1) any concurrent `free` doesn't use this index and 2) we have an `Acquire` relationship with the `free` before it. + // + // So yeah, we could be reading from outdated memory (the free buffer), but the part that we are reading, hasn't changed, so that's ok. + // That satisfies safety but not correctness. + // We still need to double check that a free didn't happen, and retry if it did. + // Otherwise, this entity might be given out twice. + let entity = unsafe { self.buffer.get(index) }; + + let ideal_state = state.pop(1); + // If we fail, we need to acquire the new state. + // If we succeed, we are finished, and we haven't changed any memory, so we can used relaxed ordering. + match self + .len + .try_set_state(state, ideal_state, Ordering::Relaxed, Ordering::Acquire) + { + Ok(_) => return Some(entity), + Err(new_state) => state = new_state, + } + } + } +} + +struct FreshAllocator { + /// The next value of [`Entity::index`] to give out if needed. + next_entity_index: AtomicU32, +} + +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; + + /// The total number of indices given out. + #[inline] + fn total_entity_indices(&self) -> u32 { + self.next_entity_index.load(Ordering::Relaxed) + } + + /// This just panics. + /// It is included to help with branch prediction, and put the panic message in one spot. + #[cold] + #[inline] + fn on_overflow() -> ! { + panic!("too many entities") + } + + /// Allocates a fresh [`EntityIndex`]. + /// This row has never been given out before. + #[inline] + fn alloc(&self) -> Entity { + let index = self.next_entity_index.fetch_add(1, Ordering::Relaxed); + if index == Self::MAX_ENTITIES { + Self::on_overflow(); + } + // 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)) }) + } + + /// Allocates `count` [`EntityIndex`]s. + /// These rows will be fresh. + /// They have never been given out before. + fn alloc_many(&self, count: u32) -> AllocUniqueEntityIndexIterator { + 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) + { + Some(new_next_entity_index) => start_new..new_next_entity_index, + None => Self::on_overflow(), + }; + AllocUniqueEntityIndexIterator(new) + } +} + +/// An [`Iterator`] returning a sequence of [`EntityIndex`] values from an [`Allocator`] that are never aliased. +/// These rows have never been given out before. +/// +/// **NOTE:** Dropping will leak the remaining entity rows! +pub(super) struct AllocUniqueEntityIndexIterator(core::ops::Range); + +impl Iterator for AllocUniqueEntityIndexIterator { + type Item = Entity; + + #[inline] + fn next(&mut self) -> Option { + self.0 + .next() + // SAFETY: This came from an *exclusive* range. It can never be max. + .map(|idx| unsafe { EntityIndex::new(NonMaxU32::new_unchecked(idx)) }) + .map(Entity::from_index) + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + self.0.size_hint() + } +} + +impl ExactSizeIterator for AllocUniqueEntityIndexIterator {} +impl core::iter::FusedIterator for AllocUniqueEntityIndexIterator {} + +/// This stores allocation data shared by all entity allocators. +struct SharedAllocator { + /// The entities pending reuse + free: FreeList, + fresh: FreshAllocator, + /// Tracks whether or not the primary [`Allocator`] has been closed or not. + is_closed: AtomicBool, +} + +impl SharedAllocator { + /// Constructs a [`SharedAllocator`] + fn new() -> Self { + Self { + free: FreeList::new(), + fresh: FreshAllocator { + next_entity_index: AtomicU32::new(0), + }, + is_closed: AtomicBool::new(false), + } + } + + /// Allocates a new [`Entity`], reusing a freed index if one exists. + /// + /// # Safety + /// + /// This must not conflict with [`FreeList::free`] calls. + #[inline] + unsafe fn alloc(&self) -> Entity { + // SAFETY: assured by caller + unsafe { self.free.alloc() }.unwrap_or_else(|| self.fresh.alloc()) + } + + /// Allocates a `count` [`Entity`]s, reusing freed indices if they exist. + /// + /// # Safety + /// + /// This must not conflict with [`FreeList::free`] calls for the duration of the iterator. + #[inline] + unsafe fn alloc_many(&self, count: u32) -> AllocEntitiesIterator<'_> { + let reused = self.free.alloc_many(count); + let still_need = count - reused.len() as u32; + let new = self.fresh.alloc_many(still_need); + AllocEntitiesIterator { new, reused } + } + + /// 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()) + } + + /// Marks the allocator as closed, but it will still function normally. + fn close(&self) { + self.is_closed.store(true, Ordering::Release); + } + + /// Returns true if [`Self::close`] has been called. + fn is_closed(&self) -> bool { + self.is_closed.load(Ordering::Acquire) + } +} + +/// This keeps track of freed entities and allows the allocation of new ones. +pub(super) struct Allocator { + shared: Arc, +} + +impl Default for Allocator { + fn default() -> Self { + Self::new() + } +} + +impl Allocator { + /// Constructs a new [`Allocator`] + pub(super) fn new() -> Self { + Self { + shared: Arc::new(SharedAllocator::new()), + } + } + + /// Allocates a new [`Entity`], reusing a freed index if one exists. + #[inline] + pub(super) fn alloc(&self) -> Entity { + // SAFETY: violating safety requires a `&mut self` to exist, but rust does not allow that. + unsafe { self.shared.alloc() } + } + + /// The total number of indices given out. + #[inline] + fn total_entity_indices(&self) -> u32 { + self.shared.fresh.total_entity_indices() + } + + /// The number of free entities. + #[inline] + fn num_free(&self) -> u32 { + // RISK: `free` requires mutable access. + self.shared.free.num_free() + } + + /// Frees the entity allowing it to be reused. + #[inline] + pub(super) fn free(&mut self, entity: Entity) { + // SAFETY: We have `&mut self`. + unsafe { + self.shared.free.free(entity); + } + } + + /// Allocates `count` entities in an iterator. + #[inline] + pub(super) fn alloc_many(&self, count: u32) -> AllocEntitiesIterator<'_> { + // SAFETY: `free` takes `&mut self`, and this lifetime is captured by the iterator. + unsafe { self.shared.alloc_many(count) } + } +} + +impl Drop for Allocator { + fn drop(&mut self) { + self.shared.close(); + } +} + +impl core::fmt::Debug for Allocator { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct(core::any::type_name::()) + .field("total_indices", &self.total_entity_indices()) + .field("total_free", &self.num_free()) + .finish() + } +} + +/// An [`Iterator`] returning a sequence of [`Entity`] values from an [`Allocator`]. +/// +/// **NOTE:** Dropping will leak the remaining entities! +pub(super) struct AllocEntitiesIterator<'a> { + new: AllocUniqueEntityIndexIterator, + reused: FreeBufferIterator<'a>, +} + +impl<'a> Iterator for AllocEntitiesIterator<'a> { + type Item = Entity; + + fn next(&mut self) -> Option { + self.reused.next().or_else(|| self.new.next()) + } + + fn size_hint(&self) -> (usize, Option) { + let len = self.reused.len() + self.new.len(); + (len, Some(len)) + } +} + +impl<'a> ExactSizeIterator for AllocEntitiesIterator<'a> {} +impl<'a> core::iter::FusedIterator for AllocEntitiesIterator<'a> {} + +// SAFETY: Newly reserved entity values are unique. +unsafe impl EntitySetIterator for AllocEntitiesIterator<'_> {} + +impl Drop for AllocEntitiesIterator<'_> { + fn drop(&mut self) { + let leaking = self.len(); + if leaking > 0 { + warn!( + "{} entities being leaked via unfinished `AllocEntitiesIterator`", + leaking + ); + } + } +} + +/// This is a stripped down entity allocator that operates on fewer assumptions than [`EntityAllocator`](super::EntityAllocator). +/// As a result, using this will be slower than than the main allocator but this offers additional freedoms. +/// In particular, this type is fully owned, allowing you to allocate entities for a world without locking or holding reference to the world. +/// This is especially useful in async contexts. +#[derive(Clone)] +pub struct RemoteAllocator { + shared: Arc, +} + +impl RemoteAllocator { + /// Creates a new [`RemoteAllocator`] with the provided [`Allocator`] source. + /// If the source is ever destroyed, [`Self::alloc`] will yield garbage values. + /// Be sure to use [`Self::is_closed`] to determine if it is safe to use these entities. + pub(super) fn new(source: &Allocator) -> Self { + Self { + shared: source.shared.clone(), + } + } + + /// Returns whether or not this [`RemoteAllocator`] is connected to this source [`Allocator`]. + pub(super) fn is_connected_to(&self, source: &Allocator) -> bool { + Arc::ptr_eq(&self.shared, &source.shared) + } + + /// Allocates an entity remotely. + /// + /// 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 alloc(&self) -> Entity { + self.shared.remote_alloc() + } + + /// Returns whether or not this [`RemoteAllocator`] is still connected to its source [`EntityAllocator`](super::EntityAllocator). + /// + /// Note that this could close immediately after the function returns false, so be careful. + /// The best way to ensure that does not happen is to only trust the returned value while holding a reference to the world + /// and to ensure it is the right world through [`EntityAllocator::has_remote_allocator`](super::EntityAllocator::has_remote_allocator). + /// + /// This is generally best used as a diagnostic. + /// [`EntityAllocator::has_remote_allocator`](super::EntityAllocator::has_remote_allocator) is a better check for correctness. + pub fn is_closed(&self) -> bool { + self.shared.is_closed() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use alloc::vec; + + /// Ensure the total capacity of [`OwnedBuffer`] is `u32::MAX + 1`. + #[test] + fn chunk_capacity_sums() { + let total: u64 = (0..FreeBuffer::NUM_CHUNKS) + .map(FreeBuffer::capacity_of_chunk) + .map(|x| x as u64) + .sum(); + // The last 2 won't be used, but that's ok. + // Keeping them powers of 2 makes things faster. + let expected = u32::MAX as u64 + 1; + assert_eq!(total, expected); + } + + /// Ensure [`OwnedBuffer`] can be properly indexed + #[test] + fn chunk_indexing() { + let to_test = vec![ + (0, (0, 0, 512)), // index 0 cap = 512 + (1, (0, 1, 512)), + (256, (0, 256, 512)), + (511, (0, 511, 512)), + (512, (1, 0, 512)), // index 1 cap = 512 + (1023, (1, 511, 512)), + (1024, (2, 0, 1024)), // index 2 cap = 1024 + (1025, (2, 1, 1024)), + (2047, (2, 1023, 1024)), + (2048, (3, 0, 2048)), // index 3 cap = 2048 + (4095, (3, 2047, 2048)), + (4096, (4, 0, 4096)), // index 3 cap = 4096 + ]; + + for (input, output) in to_test { + assert_eq!(FreeBuffer::index_info(input), output); + } + } + + #[test] + fn buffer_len_encoding() { + let len = FreeCount::new_zero_len(); + assert_eq!(len.state(Ordering::Relaxed).length(), 0); + assert_eq!(len.pop_for_state(200, Ordering::Relaxed).length(), 0); + len.set_state_risky( + FreeCountState::new_zero_len().with_length(5), + Ordering::Relaxed, + ); + assert_eq!(len.pop_for_state(2, Ordering::Relaxed).length(), 5); + assert_eq!(len.pop_for_state(2, Ordering::Relaxed).length(), 3); + assert_eq!(len.pop_for_state(2, Ordering::Relaxed).length(), 1); + assert_eq!(len.pop_for_state(2, Ordering::Relaxed).length(), 0); + } + + #[test] + fn uniqueness() { + let mut entities = Vec::with_capacity(2000); + let mut allocator = Allocator::new(); + entities.extend(allocator.alloc_many(1000)); + + let pre_len = entities.len(); + entities.dedup(); + assert_eq!(pre_len, entities.len()); + + for e in entities.drain(..) { + allocator.free(e); + } + + entities.extend(allocator.alloc_many(500)); + for _ in 0..1000 { + entities.push(allocator.alloc()); + } + entities.extend(allocator.alloc_many(500)); + + let pre_len = entities.len(); + entities.dedup(); + assert_eq!(pre_len, entities.len()); + } +}