Skip to main content

banished_formats/
par.rs

1//! Parallel map for loading work that is a pure function of one item (decoding a texture, a
2//! sound or a resource): the results come back in input order, so callers see exactly what a
3//! sequential `map` gives. Native builds spread the items over the CPU cores with scoped
4//! threads; wasm32 builds run them in order on the calling thread (wasm threads need
5//! cross-origin isolation, see `docs/PERFORMANCE.md`).
6//!
7//! Never use this for work that draws from the game's `Random` or depends on the order in
8//! which items are processed.
9
10/// `items.iter().map(f).collect()`, with the calls to `f` run on several threads (native).
11pub fn map<T: Sync, R: Send>(items: &[T], f: impl Fn(&T) -> R + Sync) -> Vec<R> {
12    #[cfg(not(target_arch = "wasm32"))]
13    {
14        let workers = threads().min(items.len());
15        if workers > 1 {
16            return threaded(items, &f, workers);
17        }
18    }
19    items.iter().map(f).collect()
20}
21
22/// `items.iter_mut().map(f).collect()` for disjoint parts of one buffer (for example bands
23/// of rows of an image). Each thread takes a run of neighbouring items (native).
24pub fn map_mut<T: Send, R: Send>(items: &mut [T], f: impl Fn(&mut T) -> R + Sync) -> Vec<R> {
25    #[cfg(not(target_arch = "wasm32"))]
26    {
27        let workers = threads().min(items.len());
28        if workers > 1 {
29            let run = items.len().div_ceil(workers);
30            let f = &f;
31            return std::thread::scope(|s| {
32                let mut parts = items.chunks_mut(run);
33                let first = parts.next();
34                let handles: Vec<_> = parts
35                    .map(|part| s.spawn(move || part.iter_mut().map(f).collect::<Vec<R>>()))
36                    .collect();
37                let mut out: Vec<R> = first.into_iter().flatten().map(f).collect();
38                for h in handles {
39                    out.extend(h.join().unwrap_or_else(|e| std::panic::resume_unwind(e)));
40                }
41                out
42            });
43        }
44    }
45    items.iter_mut().map(f).collect()
46}
47
48/// Run `background` on another thread while `foreground` runs on this one (native); on
49/// wasm32, `foreground` then `background`, in order. `background` must be a pure function of
50/// what it reads, so that the order makes no difference.
51pub fn join<A: Send, B>(
52    background: impl FnOnce() -> A + Send,
53    foreground: impl FnOnce() -> B,
54) -> (A, B) {
55    #[cfg(not(target_arch = "wasm32"))]
56    if threads() > 1 {
57        return std::thread::scope(|s| {
58            let a = s.spawn(background);
59            let b = foreground();
60            (a.join().unwrap_or_else(|e| std::panic::resume_unwind(e)), b)
61        });
62    }
63    let b = foreground();
64    (background(), b)
65}
66
67/// Each worker takes the next unclaimed item until none are left; the results are put back
68/// in input order.
69#[cfg(not(target_arch = "wasm32"))]
70fn threaded<T: Sync, R: Send>(
71    items: &[T],
72    f: &(impl Fn(&T) -> R + Sync),
73    workers: usize,
74) -> Vec<R> {
75    use std::sync::atomic::{AtomicUsize, Ordering};
76    let next = AtomicUsize::new(0);
77    let work = || {
78        let mut done = Vec::new();
79        loop {
80            let i = next.fetch_add(1, Ordering::Relaxed);
81            let Some(item) = items.get(i) else {
82                return done;
83            };
84            done.push((i, f(item)));
85        }
86    };
87    let mut out: Vec<Option<R>> = (0..items.len()).map(|_| None).collect();
88    std::thread::scope(|s| {
89        let handles: Vec<_> = (1..workers).map(|_| s.spawn(work)).collect();
90        let mut place = |done: Vec<(usize, R)>| {
91            for (i, r) in done {
92                out[i] = Some(r);
93            }
94        };
95        place(work());
96        for h in handles {
97            // A panic in `f` on a worker is raised again here, as it would be in order.
98            place(h.join().unwrap_or_else(|e| std::panic::resume_unwind(e)));
99        }
100    });
101    out.into_iter()
102        .map(|r| r.expect("every item mapped"))
103        .collect()
104}
105
106/// Threads to use: `BANISHED_THREADS` if set (1 runs everything in order), else the number
107/// of cores.
108#[cfg(not(target_arch = "wasm32"))]
109fn threads() -> usize {
110    static THREADS: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
111    *THREADS.get_or_init(|| {
112        std::env::var("BANISHED_THREADS")
113            .ok()
114            .and_then(|v| v.parse().ok())
115            .filter(|&t: &usize| t >= 1)
116            .or_else(|| std::thread::available_parallelism().ok().map(|n| n.get()))
117            .unwrap_or(1)
118    })
119}
120
121#[cfg(test)]
122mod tests {
123    #[test]
124    fn keeps_input_order() {
125        let items: Vec<u64> = (0..1000).collect();
126        let out = super::map(&items, |&i| i * i);
127        assert_eq!(out, items.iter().map(|&i| i * i).collect::<Vec<_>>());
128        assert!(super::map(&[] as &[u8], |&b| b).is_empty());
129        assert_eq!(super::map(&[7u8], |&b| b + 1), vec![8]);
130        assert_eq!(super::join(|| 1, || 2), (1, 2));
131        let mut buf: Vec<u32> = (0..1000).collect();
132        let mut parts: Vec<&mut [u32]> = buf.chunks_mut(7).collect();
133        let sums = super::map_mut(&mut parts, |p| {
134            p.iter_mut().for_each(|v| *v *= 2);
135            p.iter().sum::<u32>()
136        });
137        let expect: Vec<u32> = (0..1000u32)
138            .collect::<Vec<_>>()
139            .chunks(7)
140            .map(|c| c.iter().map(|v| v * 2).sum())
141            .collect();
142        assert_eq!(sums, expect);
143        assert_eq!(buf, (0..1000).map(|v| v * 2).collect::<Vec<u32>>());
144    }
145}