rustpython_common/
atomic.rs

1use core::ptr::{self, NonNull};
2pub use core::sync::atomic::*;
3pub use radium::Radium;
4
5mod sealed {
6    pub trait Sealed {}
7}
8pub trait PyAtomicScalar: sealed::Sealed {
9    type Radium: Radium<Item = Self>;
10}
11
12pub type PyAtomic<T> = <T as PyAtomicScalar>::Radium;
13
14#[cfg(feature = "threading")]
15macro_rules! atomic_ty {
16    ($i:ty, $atomic:ty) => {
17        $atomic
18    };
19}
20#[cfg(not(feature = "threading"))]
21macro_rules! atomic_ty {
22    ($i:ty, $atomic:ty) => {
23        core::cell::Cell<$i>
24    };
25}
26macro_rules! impl_atomic_scalar {
27    ($(($i:ty, $atomic:ty),)*) => {
28        $(
29            impl sealed::Sealed for $i {}
30            impl PyAtomicScalar for $i {
31                type Radium = atomic_ty!($i, $atomic);
32            }
33        )*
34    };
35}
36impl_atomic_scalar!(
37    (u8, AtomicU8),
38    (i8, AtomicI8),
39    (u16, AtomicU16),
40    (i16, AtomicI16),
41    (u32, AtomicU32),
42    (i32, AtomicI32),
43    (u64, AtomicU64),
44    (i64, AtomicI64),
45    (usize, AtomicUsize),
46    (isize, AtomicIsize),
47    (bool, AtomicBool),
48);
49
50impl<T> sealed::Sealed for *mut T {}
51impl<T> PyAtomicScalar for *mut T {
52    type Radium = atomic_ty!(*mut T, AtomicPtr<T>);
53}
54
55/// Atomic `f64`. The value is stored as its bits in an `AtomicU64`.
56#[repr(transparent)]
57pub struct AtomicF64 {
58    bits: AtomicU64,
59}
60
61impl AtomicF64 {
62    #[inline]
63    #[must_use]
64    pub const fn new(value: f64) -> Self {
65        Self {
66            bits: AtomicU64::new(value.to_bits()),
67        }
68    }
69
70    #[inline]
71    #[must_use]
72    pub fn load(&self, order: Ordering) -> f64 {
73        f64::from_bits(self.bits.load(order))
74    }
75
76    #[inline]
77    pub fn store(&self, value: f64, order: Ordering) {
78        self.bits.store(value.to_bits(), order);
79    }
80}
81
82pub struct OncePtr<T> {
83    inner: PyAtomic<*mut T>,
84}
85
86impl<T> Default for OncePtr<T> {
87    fn default() -> Self {
88        Self::new()
89    }
90}
91
92impl<T> OncePtr<T> {
93    #[inline]
94    #[must_use]
95    pub fn new() -> Self {
96        Self {
97            inner: Radium::new(ptr::null_mut()),
98        }
99    }
100
101    pub fn get(&self) -> Option<NonNull<T>> {
102        NonNull::new(self.inner.load(Ordering::Acquire))
103    }
104
105    pub fn set(&self, value: NonNull<T>) -> Result<(), NonNull<T>> {
106        let exchange = self.inner.compare_exchange(
107            ptr::null_mut(),
108            value.as_ptr(),
109            Ordering::AcqRel,
110            Ordering::Acquire,
111        );
112        match exchange {
113            Ok(_) => Ok(()),
114            Err(_) => Err(value),
115        }
116    }
117
118    pub fn get_or_init<F>(&self, f: F) -> NonNull<T>
119    where
120        F: FnOnce() -> Box<T>,
121    {
122        enum Void {}
123        match self.get_or_try_init(|| Ok::<_, Void>(f())) {
124            Ok(val) => val,
125            Err(void) => match void {},
126        }
127    }
128
129    pub fn get_or_try_init<F, E>(&self, f: F) -> Result<NonNull<T>, E>
130    where
131        F: FnOnce() -> Result<Box<T>, E>,
132    {
133        if let Some(val) = self.get() {
134            return Ok(val);
135        }
136
137        Ok(self.initialize(f()?))
138    }
139
140    #[cold]
141    fn initialize(&self, val: Box<T>) -> NonNull<T> {
142        let ptr = Box::into_raw(val);
143        let exchange =
144            self.inner
145                .compare_exchange(ptr::null_mut(), ptr, Ordering::AcqRel, Ordering::Acquire);
146        let ptr = match exchange {
147            Ok(_) => ptr,
148            Err(winner) => {
149                drop(unsafe { Box::from_raw(ptr) });
150                winner
151            }
152        };
153        unsafe { NonNull::new_unchecked(ptr) }
154    }
155}
rustpython_common/atomic.rs

rustpython_common/
atomic.rs
