polars_utils/
relaxed_cell.rs

1use std::fmt;
2use std::sync::atomic::*;
3
4#[derive(Default)]
5#[repr(transparent)]
6pub struct RelaxedCell<T: AtomicNative>(T::Atomic);
7
8impl<T: AtomicNative> RelaxedCell<T> {
9    #[inline(always)]
10    pub fn load(&self) -> T {
11        T::load(&self.0)
12    }
13
14    #[inline(always)]
15    pub fn store(&self, value: T) {
16        T::store(&self.0, value)
17    }
18
19    #[inline(always)]
20    pub fn fetch_add(&self, value: T) -> T {
21        T::fetch_add(&self.0, value)
22    }
23
24    #[inline(always)]
25    pub fn get_mut(&mut self) -> &mut T {
26        T::get_mut(&mut self.0)
27    }
28}
29
30impl<T: AtomicNative> From<T> for RelaxedCell<T> {
31    #[inline(always)]
32    fn from(value: T) -> Self {
33        RelaxedCell(T::Atomic::from(value))
34    }
35}
36
37impl<T: AtomicNative> Clone for RelaxedCell<T> {
38    fn clone(&self) -> Self {
39        Self(T::Atomic::from(self.load()))
40    }
41}
42
43impl<T: AtomicNative> fmt::Debug for RelaxedCell<T> {
44    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
45        f.debug_tuple("RelaxedCell").field(&self.load()).finish()
46    }
47}
48
49pub trait AtomicNative: Sized + Default + fmt::Debug {
50    type Atomic: From<Self>;
51
52    fn load(atomic: &Self::Atomic) -> Self;
53    fn store(atomic: &Self::Atomic, val: Self);
54    fn fetch_add(atomic: &Self::Atomic, val: Self) -> Self;
55    fn get_mut(atomic: &mut Self::Atomic) -> &mut Self;
56}
57
58macro_rules! impl_relaxed_cell {
59    ($T:ty, $new:ident, $A:ty) => {
60        impl RelaxedCell<$T> {
61            // Not part of the trait as it should be const.
62            pub const fn $new(value: $T) -> Self {
63                Self(<$A>::new(value))
64            }
65        }
66
67        impl AtomicNative for $T {
68            type Atomic = $A;
69
70            #[inline(always)]
71            fn load(atomic: &Self::Atomic) -> Self {
72                atomic.load(Ordering::Relaxed)
73            }
74
75            #[inline(always)]
76            fn store(atomic: &Self::Atomic, val: Self) {
77                atomic.store(val, Ordering::Relaxed);
78            }
79
80            #[inline(always)]
81            fn fetch_add(atomic: &Self::Atomic, val: Self) -> Self {
82                atomic.fetch_add(val, Ordering::Relaxed)
83            }
84
85            #[inline(always)]
86            fn get_mut(atomic: &mut Self::Atomic) -> &mut Self {
87                atomic.get_mut()
88            }
89        }
90    };
91}
92
93impl_relaxed_cell!(u8, new_u8, AtomicU8);
94impl_relaxed_cell!(u32, new_u32, AtomicU32);
95impl_relaxed_cell!(u64, new_u64, AtomicU64);
96impl_relaxed_cell!(usize, new_usize, AtomicUsize);
97
98impl RelaxedCell<bool> {
99    // Not part of the trait as it should be const.
100    pub const fn new_bool(value: bool) -> Self {
101        Self(AtomicBool::new(value))
102    }
103}
104
105impl AtomicNative for bool {
106    type Atomic = AtomicBool;
107
108    #[inline(always)]
109    fn load(atomic: &Self::Atomic) -> Self {
110        atomic.load(Ordering::Relaxed)
111    }
112
113    #[inline(always)]
114    fn store(atomic: &Self::Atomic, val: Self) {
115        atomic.store(val, Ordering::Relaxed);
116    }
117
118    #[inline(always)]
119    fn fetch_add(_atomic: &Self::Atomic, _val: Self) -> Self {
120        unimplemented!()
121    }
122
123    #[inline(always)]
124    fn get_mut(atomic: &mut Self::Atomic) -> &mut Self {
125        atomic.get_mut()
126    }
127}
polars_utils/relaxed_cell.rs

polars_utils/
relaxed_cell.rs
