bevy_utils/
parallel_queue.rs

1#[cfg(all(feature = "alloc", not(feature = "std")))]
2use alloc::vec::Vec;
3
4use core::{cell::RefCell, ops::DerefMut};
5use thread_local::ThreadLocal;
6
7/// A cohesive set of thread-local values of a given type.
8///
9/// Mutable references can be fetched if `T: Default` via [`Parallel::scope`].
10#[derive(Default)]
11pub struct Parallel<T: Send> {
12    locals: ThreadLocal<RefCell<T>>,
13}
14
15/// A scope guard of a `Parallel`, when this struct is dropped ,the value will writeback to its `Parallel`
16impl<T: Send> Parallel<T> {
17    /// Gets a mutable iterator over all of the per-thread queues.
18    pub fn iter_mut(&mut self) -> impl Iterator<Item = &'_ mut T> {
19        self.locals.iter_mut().map(RefCell::get_mut)
20    }
21
22    /// Clears all of the stored thread local values.
23    pub fn clear(&mut self) {
24        self.locals.clear();
25    }
26}
27
28impl<T: Default + Send> Parallel<T> {
29    /// Retrieves the thread-local value for the current thread and runs `f` on it.
30    ///
31    /// If there is no thread-local value, it will be initialized to its default.
32    pub fn scope<R>(&self, f: impl FnOnce(&mut T) -> R) -> R {
33        let mut cell = self.locals.get_or_default().borrow_mut();
34        let ret = f(cell.deref_mut());
35        ret
36    }
37
38    /// Mutably borrows the thread-local value.
39    ///
40    /// If there is no thread-local value, it will be initialized to it's default.
41    pub fn borrow_local_mut(&self) -> impl DerefMut<Target = T> + '_ {
42        self.locals.get_or_default().borrow_mut()
43    }
44}
45
46impl<T, I> Parallel<I>
47where
48    I: IntoIterator<Item = T> + Default + Send + 'static,
49{
50    /// Drains all enqueued items from all threads and returns an iterator over them.
51    ///
52    /// Unlike [`Vec::drain`], this will piecemeal remove chunks of the data stored.
53    /// If iteration is terminated part way, the rest of the enqueued items in the same
54    /// chunk will be dropped, and the rest of the undrained elements will remain.
55    ///
56    /// The ordering is not guaranteed.
57    pub fn drain(&mut self) -> impl Iterator<Item = T> + '_ {
58        self.locals.iter_mut().flat_map(|item| item.take())
59    }
60}
61
62#[cfg(feature = "alloc")]
63impl<T: Send> Parallel<Vec<T>> {
64    /// Collect all enqueued items from all threads and appends them to the end of a
65    /// single Vec.
66    ///
67    /// The ordering is not guaranteed.
68    pub fn drain_into(&mut self, out: &mut Vec<T>) {
69        let size = self
70            .locals
71            .iter_mut()
72            .map(|queue| queue.get_mut().len())
73            .sum();
74        out.reserve(size);
75        for queue in self.locals.iter_mut() {
76            out.append(queue.get_mut());
77        }
78    }
79}
bevy_utils/parallel_queue.rs

bevy_utils/
parallel_queue.rs
