Skip to main content

rustpython_common/
static_cell.rs

1#[cfg(feature = "threading")]
2mod threading {
3    use crate::lock::OnceCell;
4
5    pub struct StaticCell<T: 'static> {
6        inner: OnceCell<T>,
7    }
8
9    impl<T> StaticCell<T> {
10        #[doc(hidden)]
11        pub const fn _from_once_cell(inner: OnceCell<T>) -> Self {
12            Self { inner }
13        }
14
15        pub fn get(&'static self) -> Option<&'static T> {
16            self.inner.get()
17        }
18
19        pub fn set(&'static self, value: T) -> Result<(), T> {
20            self.inner.set(value)
21        }
22
23        pub fn get_or_init<F>(&'static self, f: F) -> &'static T
24        where
25            F: FnOnce() -> T,
26        {
27            self.inner.get_or_init(f)
28        }
29
30        pub fn get_or_try_init<F, E>(&'static self, f: F) -> Result<&'static T, E>
31        where
32            F: FnOnce() -> Result<T, E>,
33        {
34            if let Some(val) = self.inner.get() {
35                return Ok(val);
36            }
37            let val = f()?;
38            let _ = self.inner.set(val);
39            Ok(self.inner.get().unwrap())
40        }
41    }
42
43    #[macro_export]
44    macro_rules! static_cell {
45        ($($(#[$attr:meta])* $vis:vis static $name:ident: $t:ty;)+) => {
46            $($(#[$attr])*
47            $vis static $name: $crate::static_cell::StaticCell<$t> =
48                $crate::static_cell::StaticCell::_from_once_cell($crate::lock::OnceCell::new());)+
49        };
50    }
51}
52#[cfg(feature = "threading")]
53pub use threading::*;
54
55#[cfg(all(not(feature = "threading"), feature = "std"))]
56mod non_threading {
57    use crate::lock::OnceCell;
58    use std::thread::LocalKey;
59
60    pub struct StaticCell<T: 'static> {
61        inner: &'static LocalKey<OnceCell<&'static T>>,
62    }
63
64    fn leak<T>(x: T) -> &'static T {
65        Box::leak(Box::new(x))
66    }
67
68    impl<T> StaticCell<T> {
69        #[doc(hidden)]
70        pub const fn _from_local_key(inner: &'static LocalKey<OnceCell<&'static T>>) -> Self {
71            Self { inner }
72        }
73
74        pub fn get(&'static self) -> Option<&'static T> {
75            self.inner.with(|x| x.get().copied())
76        }
77
78        pub fn set(&'static self, value: T) -> Result<(), T> {
79            self.inner.with(|x| {
80                if x.get().is_some() {
81                    Err(value)
82                } else {
83                    let _ = x.set(leak(value));
84                    Ok(())
85                }
86            })
87        }
88
89        pub fn get_or_init<F>(&'static self, f: F) -> &'static T
90        where
91            F: FnOnce() -> T,
92        {
93            self.inner.with(|x| *x.get_or_init(|| leak(f())))
94        }
95
96        pub fn get_or_try_init<F, E>(&'static self, f: F) -> Result<&'static T, E>
97        where
98            F: FnOnce() -> Result<T, E>,
99        {
100            self.inner.with(|x| {
101                if let Some(val) = x.get() {
102                    Ok(*val)
103                } else {
104                    let val = leak(f()?);
105                    let _ = x.set(val);
106                    Ok(val)
107                }
108            })
109        }
110    }
111
112    #[macro_export]
113    macro_rules! static_cell {
114        ($($(#[$attr:meta])* $vis:vis static $name:ident: $t:ty;)+) => {
115            $($(#[$attr])*
116            $vis static $name: $crate::static_cell::StaticCell<$t> = {
117                ::std::thread_local! {
118                     $vis static $name: $crate::lock::OnceCell<&'static $t> = const {
119                         $crate::lock::OnceCell::new()
120                     };
121                }
122                $crate::static_cell::StaticCell::_from_local_key(&$name)
123            };)+
124        };
125    }
126}
127#[cfg(all(not(feature = "threading"), feature = "std"))]
128pub use non_threading::*;
129
130// Same as `threading` variant, but wraps unsync::OnceCell with Sync.
131#[cfg(all(not(feature = "threading"), not(feature = "std")))]
132mod no_std {
133    use crate::lock::OnceCell;
134
135    // unsync::OnceCell is !Sync, but without std there can be no threads.
136    struct SyncOnceCell<T>(OnceCell<T>);
137    // SAFETY: Without std, threading is impossible.
138    unsafe impl<T> Sync for SyncOnceCell<T> {}
139
140    pub struct StaticCell<T: 'static> {
141        inner: SyncOnceCell<T>,
142    }
143
144    impl<T> StaticCell<T> {
145        #[doc(hidden)]
146        pub const fn _from_once_cell(inner: OnceCell<T>) -> Self {
147            Self {
148                inner: SyncOnceCell(inner),
149            }
150        }
151
152        pub fn get(&'static self) -> Option<&'static T> {
153            self.inner.0.get()
154        }
155
156        pub fn set(&'static self, value: T) -> Result<(), T> {
157            self.inner.0.set(value)
158        }
159
160        pub fn get_or_init<F>(&'static self, f: F) -> &'static T
161        where
162            F: FnOnce() -> T,
163        {
164            self.inner.0.get_or_init(f)
165        }
166
167        pub fn get_or_try_init<F, E>(&'static self, f: F) -> Result<&'static T, E>
168        where
169            F: FnOnce() -> Result<T, E>,
170        {
171            if let Some(val) = self.inner.0.get() {
172                return Ok(val);
173            }
174            let val = f()?;
175            let _ = self.inner.0.set(val);
176            Ok(self.inner.0.get().unwrap())
177        }
178    }
179
180    #[macro_export]
181    macro_rules! static_cell {
182        ($($(#[$attr:meta])* $vis:vis static $name:ident: $t:ty;)+) => {
183            $($(#[$attr])*
184            $vis static $name: $crate::static_cell::StaticCell<$t> =
185                $crate::static_cell::StaticCell::_from_once_cell($crate::lock::OnceCell::new());)+
186        };
187    }
188}
189#[cfg(all(not(feature = "threading"), not(feature = "std")))]
190pub use no_std::*;