bevy/crates/bevy_utils/src/parallel_queue.rs
Zachary Harrold bd489068c6
Allow bevy_utils in no_std Contexts (#15279)
# Objective

- Adjust `bevy_utils` to make it `no_std` compatible
- Partially replaces #6581
- Contributes to #8161
- Contributes to #6370

## Solution

Added `alloc` and `std` features to `bevy_utils` (`std` is enabled by
default), allowing the crate's use in `no_std` contexts.

## Testing

- CI passed locally.
- Used `bevy_utils` in a `no_std` crate as an experiment and compiled
successfully.

## Migration Guide

If you were importing `bevy_utils` and setting `default_features` to
`false`, but relying on elements which are now gated behind the `std` or
`alloc` features, include the relevant feature in your `Cargo.toml`.

## Notes

- Bevy already includes a single `no_std` crate, `bevy_ptr`, so there is
precedent for this change.
- As `bevy_utils` is widely used across the rest of Bevy, further work
to make Bevy `no_std` compatible would be blocked on this crate, if such
work was to be undertaken.
- Most of the changes in this PR are just the removal of an unnecessary
call to `to_string()` within unit tests.
2024-09-18 16:00:03 +00:00

79 lines
2.6 KiB
Rust

#[cfg(all(feature = "alloc", not(feature = "std")))]
use alloc::vec::Vec;
use core::{cell::RefCell, ops::DerefMut};
use thread_local::ThreadLocal;
/// A cohesive set of thread-local values of a given type.
///
/// Mutable references can be fetched if `T: Default` via [`Parallel::scope`].
#[derive(Default)]
pub struct Parallel<T: Send> {
locals: ThreadLocal<RefCell<T>>,
}
/// A scope guard of a `Parallel`, when this struct is dropped ,the value will writeback to its `Parallel`
impl<T: Send> Parallel<T> {
/// Gets a mutable iterator over all of the per-thread queues.
pub fn iter_mut(&mut self) -> impl Iterator<Item = &'_ mut T> {
self.locals.iter_mut().map(RefCell::get_mut)
}
/// Clears all of the stored thread local values.
pub fn clear(&mut self) {
self.locals.clear();
}
}
impl<T: Default + Send> Parallel<T> {
/// Retrieves the thread-local value for the current thread and runs `f` on it.
///
/// If there is no thread-local value, it will be initialized to its default.
pub fn scope<R>(&self, f: impl FnOnce(&mut T) -> R) -> R {
let mut cell = self.locals.get_or_default().borrow_mut();
let ret = f(cell.deref_mut());
ret
}
/// Mutably borrows the thread-local value.
///
/// If there is no thread-local value, it will be initialized to it's default.
pub fn borrow_local_mut(&self) -> impl DerefMut<Target = T> + '_ {
self.locals.get_or_default().borrow_mut()
}
}
impl<T, I> Parallel<I>
where
I: IntoIterator<Item = T> + Default + Send + 'static,
{
/// Drains all enqueued items from all threads and returns an iterator over them.
///
/// Unlike [`Vec::drain`], this will piecemeal remove chunks of the data stored.
/// If iteration is terminated part way, the rest of the enqueued items in the same
/// chunk will be dropped, and the rest of the undrained elements will remain.
///
/// The ordering is not guaranteed.
pub fn drain(&mut self) -> impl Iterator<Item = T> + '_ {
self.locals.iter_mut().flat_map(|item| item.take())
}
}
#[cfg(feature = "alloc")]
impl<T: Send> Parallel<Vec<T>> {
/// Collect all enqueued items from all threads and appends them to the end of a
/// single Vec.
///
/// The ordering is not guaranteed.
pub fn drain_into(&mut self, out: &mut Vec<T>) {
let size = self
.locals
.iter_mut()
.map(|queue| queue.get_mut().len())
.sum();
out.reserve(size);
for queue in self.locals.iter_mut() {
out.append(queue.get_mut());
}
}
}