1pub 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
22pub 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
48pub 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#[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 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#[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}