swc_parallel/
lib.rs

1#![cfg_attr(not(feature = "chili"), allow(unused_variables))]
2
3use std::{cell::RefCell, mem::transmute};
4
5pub mod items;
6pub mod merge;
7
8#[cfg(all(not(feature = "chili"), not(feature = "rayon"), feature = "parallel"))]
9compile_error!("You must enable `chili` or `rayon` feature if you want to use `parallel` feature");
10
11#[cfg(all(feature = "chili", feature = "rayon"))]
12compile_error!("You must enable `chili` or `rayon` feature, not both");
13
14#[derive(Default)]
15pub struct MaybeScope<'a>(ScopeLike<'a>);
16
17enum ScopeLike<'a> {
18    Scope(Scope<'a>),
19    #[cfg(feature = "chili")]
20    Global(Option<chili::Scope<'a>>),
21}
22
23impl Default for ScopeLike<'_> {
24    fn default() -> Self {
25        #[cfg(feature = "chili")]
26        {
27            ScopeLike::Global(None)
28        }
29
30        #[cfg(not(feature = "chili"))]
31        {
32            ScopeLike::Scope(Scope(std::marker::PhantomData))
33        }
34    }
35}
36
37impl<'a> From<Scope<'a>> for MaybeScope<'a> {
38    fn from(value: Scope<'a>) -> Self {
39        MaybeScope(ScopeLike::Scope(value))
40    }
41}
42
43impl<'a> MaybeScope<'a> {
44    #[allow(clippy::redundant_closure)]
45    pub fn with<F, R>(&mut self, f: F) -> R
46    where
47        F: FnOnce(Scope<'a>) -> R,
48    {
49        #[cfg(feature = "chili")]
50        let scope: &mut chili::Scope = match &mut self.0 {
51            ScopeLike::Scope(scope) => unsafe {
52                // Safety: chili Scope will be alive until the end of the function, because it's
53                // contract of 'a lifetime in the type.
54                transmute::<&mut chili::Scope, &mut chili::Scope>(&mut scope.0)
55            },
56            #[cfg(feature = "chili")]
57            ScopeLike::Global(global_scope) => {
58                // Initialize global scope lazily, and only once.
59                let scope = global_scope.get_or_insert_with(|| chili::Scope::global());
60
61                unsafe {
62                    // Safety: Global scope is not dropped until the end of the program, and no one
63                    // can access this **instance** of the global scope in the same time.
64                    transmute::<&mut chili::Scope, &mut chili::Scope>(scope)
65                }
66            }
67        };
68
69        #[cfg(feature = "chili")]
70        let scope = Scope(scope);
71
72        #[cfg(not(feature = "chili"))]
73        let scope = Scope(std::marker::PhantomData);
74
75        f(scope)
76    }
77}
78
79#[cfg(not(feature = "chili"))]
80pub struct Scope<'a>(std::marker::PhantomData<&'a ()>);
81
82#[cfg(feature = "chili")]
83pub struct Scope<'a>(&'a mut chili::Scope<'a>);
84
85#[inline]
86pub fn join<A, B, RA, RB>(oper_a: A, oper_b: B) -> (RA, RB)
87where
88    A: Send + FnOnce() -> RA,
89    B: Send + FnOnce() -> RB,
90    RA: Send,
91    RB: Send,
92{
93    thread_local! {
94        static SCOPE: RefCell<Option<MaybeScope<'static>>> = Default::default();
95    }
96
97    struct RemoveScopeGuard;
98
99    impl Drop for RemoveScopeGuard {
100        fn drop(&mut self) {
101            SCOPE.set(None);
102        }
103    }
104
105    let mut scope = SCOPE.take().unwrap_or_default();
106
107    let (ra, rb) = join_maybe_scoped(
108        &mut scope,
109        |scope| {
110            let scope = unsafe {
111                // Safety: inner scope cannot outlive the outer scope
112                transmute::<Scope, Scope>(scope)
113            };
114            let _guard = RemoveScopeGuard;
115            SCOPE.set(Some(MaybeScope(ScopeLike::Scope(scope))));
116
117            oper_a()
118        },
119        |scope| {
120            let scope = unsafe {
121                // Safety: inner scope cannot outlive the outer scope
122                transmute::<Scope, Scope>(scope)
123            };
124            let _guard = RemoveScopeGuard;
125            SCOPE.set(Some(MaybeScope(ScopeLike::Scope(scope))));
126
127            oper_b()
128        },
129    );
130
131    // In case of panic, we does not restore the scope so it will be None.
132    SCOPE.set(Some(scope));
133
134    (ra, rb)
135}
136
137#[inline]
138pub fn join_maybe_scoped<'a, A, B, RA, RB>(
139    scope: &mut MaybeScope<'a>,
140    oper_a: A,
141    oper_b: B,
142) -> (RA, RB)
143where
144    A: Send + FnOnce(Scope<'a>) -> RA,
145    B: Send + FnOnce(Scope<'a>) -> RB,
146    RA: Send,
147    RB: Send,
148{
149    scope.with(|scope| join_scoped(scope, oper_a, oper_b))
150}
151
152#[inline]
153pub fn join_scoped<'a, A, B, RA, RB>(scope: Scope<'a>, oper_a: A, oper_b: B) -> (RA, RB)
154where
155    A: Send + FnOnce(Scope<'a>) -> RA,
156    B: Send + FnOnce(Scope<'a>) -> RB,
157    RA: Send,
158    RB: Send,
159{
160    #[cfg(feature = "chili")]
161    let (ra, rb) = scope.0.join(
162        |scope| {
163            let scope = Scope(unsafe {
164                // Safety: This can be dangerous if the user do transmute on the scope, but it's
165                // not our fault if the user uses transmute.
166                transmute::<&mut chili::Scope, &mut chili::Scope>(scope)
167            });
168
169            oper_a(scope)
170        },
171        |scope| {
172            let scope = Scope(unsafe {
173                // Safety: This can be dangerous if the user do transmute on the scope, but it's
174                // not our fault if the user uses transmute.
175                transmute::<&mut chili::Scope, &mut chili::Scope>(scope)
176            });
177
178            oper_b(scope)
179        },
180    );
181
182    #[cfg(feature = "rayon")]
183    let (ra, rb) = rayon::join(
184        || oper_a(Scope(std::marker::PhantomData)),
185        || oper_b(Scope(std::marker::PhantomData)),
186    );
187
188    #[cfg(not(feature = "parallel"))]
189    let (ra, rb) = (
190        oper_a(Scope(std::marker::PhantomData)),
191        oper_b(Scope(std::marker::PhantomData)),
192    );
193
194    (ra, rb)
195}