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
55pub struct OncePtr<T> {
56    inner: PyAtomic<*mut T>,
57}
58
59impl<T> Default for OncePtr<T> {
60    fn default() -> Self {
61        Self::new()
62    }
63}
64
65impl<T> OncePtr<T> {
66    #[inline]
67    pub fn new() -> Self {
68        OncePtr {
69            inner: Radium::new(ptr::null_mut()),
70        }
71    }
72
73    pub fn get(&self) -> Option<NonNull<T>> {
74        NonNull::new(self.inner.load(Ordering::Acquire))
75    }
76
77    pub fn set(&self, value: NonNull<T>) -> Result<(), NonNull<T>> {
78        let exchange = self.inner.compare_exchange(
79            ptr::null_mut(),
80            value.as_ptr(),
81            Ordering::AcqRel,
82            Ordering::Acquire,
83        );
84        match exchange {
85            Ok(_) => Ok(()),
86            Err(_) => Err(value),
87        }
88    }
89
90    pub fn get_or_init<F>(&self, f: F) -> NonNull<T>
91    where
92        F: FnOnce() -> Box<T>,
93    {
94        enum Void {}
95        match self.get_or_try_init(|| Ok::<_, Void>(f())) {
96            Ok(val) => val,
97            Err(void) => match void {},
98        }
99    }
100
101    pub fn get_or_try_init<F, E>(&self, f: F) -> Result<NonNull<T>, E>
102    where
103        F: FnOnce() -> Result<Box<T>, E>,
104    {
105        if let Some(val) = self.get() {
106            return Ok(val);
107        }
108
109        Ok(self.initialize(f()?))
110    }
111
112    #[cold]
113    fn initialize(&self, val: Box<T>) -> NonNull<T> {
114        let ptr = Box::into_raw(val);
115        let exchange =
116            self.inner
117                .compare_exchange(ptr::null_mut(), ptr, Ordering::AcqRel, Ordering::Acquire);
118        let ptr = match exchange {
119            Ok(_) => ptr,
120            Err(winner) => {
121                drop(unsafe { Box::from_raw(ptr) });
122                winner
123            }
124        };
125        unsafe { NonNull::new_unchecked(ptr) }
126    }
127}
rustpython_common/atomic.rs

rustpython_common/
atomic.rs
