orx_parallel/use_var/
use_vec.rs

1use super::r#use::Use;
2use alloc::vec::Vec;
3use orx_concurrent_ordered_bag::ConcurrentOrderedBag;
4
5/// Owned worker-local mutable state.
6///
7/// `UseVec` stores one value per worker thread and lets parallel operations
8/// mutate those values independently. It is typically used with
9/// [`Par::use_vec`](crate::Par::use_vec).
10///
11/// # Examples
12///
13/// ```
14/// use orx_parallel::*;
15///
16/// let n = 10_000usize;
17/// let mut partial_sums = UseVec::new(|_| 0usize);
18///
19/// (0..n)
20///     .into_par()
21///     .use_vec(&mut partial_sums)
22///     .for_each(|thread_sum, x| *thread_sum += x);
23///
24/// let total: usize = partial_sums.into_vec().into_iter().sum();
25/// assert_eq!(total, (n - 1) * n / 2);
26/// ```
27pub struct UseVec<T: Send, F: Fn(usize) -> T + Sync> {
28    init: F,
29    cache: ConcurrentOrderedBag<T>,
30}
31
32impl<T: Send, F: Fn(usize) -> T + Sync> UseVec<T, F> {
33    /// Creates a `UseVec` with per-thread initialization logic.
34    ///
35    /// The `init` function receives the worker `thread_idx` and is invoked on
36    /// first access to create that thread's local value.
37    ///
38    /// # Example
39    ///
40    /// ```
41    /// use orx_parallel::*;
42    ///
43    /// let mut calls = UseVec::new(|_| 0usize);
44    ///
45    /// let values: Vec<_> = (0..128usize)
46    ///     .into_par()
47    ///     .use_vec(&mut calls)
48    ///     .map(|count, x| {
49    ///         *count += 1;
50    ///         x * 2
51    ///     })
52    ///     .collect();
53    ///
54    /// assert_eq!(values.len(), 128);
55    /// assert_eq!(calls.into_vec().into_iter().sum::<usize>(), 128);
56    /// ```
57    pub fn new(init: F) -> Self {
58        let cache = ConcurrentOrderedBag::new();
59        Self { init, cache }
60    }
61
62    /// Consumes the `UseVec` and returns the per-thread values as a `Vec`.
63    ///
64    /// This is typically used after a parallel computation to aggregate thread-local state.
65    ///
66    /// # Example
67    ///
68    /// ```
69    /// use orx_parallel::*;
70    ///
71    /// let mut partial_sums = UseVec::new(|_| 0usize);
72    ///
73    /// (0..100usize)
74    ///     .into_par()
75    ///     .num_threads(4)
76    ///     .use_vec(&mut partial_sums)
77    ///     .for_each(|thread_sum, x| *thread_sum += x);
78    ///
79    /// let partials = partial_sums.into_vec();
80    /// assert_eq!(partials.into_iter().sum::<usize>(), 4950);
81    /// ```
82    pub fn into_vec(self) -> Vec<T> {
83        let vec = unsafe { self.cache.into_inner().unwrap_only_if_counts_match() };
84        vec.into_iter().collect()
85    }
86}
87
88impl<T: Send, F: Fn(usize) -> T + Sync> Use for UseVec<T, F> {
89    type Item = T;
90
91    unsafe fn init_get(&self, thread_idx: usize) -> &mut Self::Item {
92        let use_var = (self.init)(thread_idx);
93        unsafe { self.cache.set_value(thread_idx, use_var) };
94
95        // SAFETY: it is safe to access to the index as it is
96        // pushed / initialized just above. Further, `get` will
97        // be called exactly once by the corresponding thread,
98        // and hence, there will be no race condition.
99        unsafe { &mut *self.cache.ptr_mut(thread_idx) }
100    }
101
102    #[inline]
103    fn get(&mut self, thread_idx: usize) -> &mut Self::Item {
104        assert!(self.cache.len() > thread_idx);
105        unsafe { &mut *self.cache.ptr_mut(thread_idx) }
106    }
107
108    fn max_threads(&self) -> Option<usize> {
109        None
110    }
111}
112
113impl<T: Send, F: Fn(usize) -> T + Sync> Use for &mut UseVec<T, F> {
114    type Item = T;
115
116    unsafe fn init_get(&self, thread_idx: usize) -> &mut Self::Item {
117        let use_var = (self.init)(thread_idx);
118        unsafe { self.cache.set_value(thread_idx, use_var) };
119
120        // SAFETY: it is safe to access to the index as it is
121        // pushed / initialized just above. Further, `get` will
122        // be called exactly once by the corresponding thread,
123        // and hence, there will be no race condition.
124        unsafe { &mut *self.cache.ptr_mut(thread_idx) }
125    }
126
127    #[inline]
128    fn get(&mut self, thread_idx: usize) -> &mut Self::Item {
129        assert!(self.cache.len() > thread_idx);
130        unsafe { &mut *self.cache.ptr_mut(thread_idx) }
131    }
132
133    fn max_threads(&self) -> Option<usize> {
134        None
135    }
136}
orx_parallel/use_var/use_vec.rs

orx_parallel/use_var/
use_vec.rs
