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 fetch_sub(&self, value: T) -> T {
26        T::fetch_sub(&self.0, value)
27    }
28
29    #[inline(always)]
30    pub fn fetch_max(&self, value: T) -> T {
31        T::fetch_max(&self.0, value)
32    }
33
34    #[inline(always)]
35    pub fn get_mut(&mut self) -> &mut T {
36        T::get_mut(&mut self.0)
37    }
38}
39
40impl<T: AtomicNative> From<T> for RelaxedCell<T> {
41    #[inline(always)]
42    fn from(value: T) -> Self {
43        RelaxedCell(T::Atomic::from(value))
44    }
45}
46
47impl<T: AtomicNative> Clone for RelaxedCell<T> {
48    fn clone(&self) -> Self {
49        Self(T::Atomic::from(self.load()))
50    }
51}
52
53impl<T: AtomicNative> fmt::Debug for RelaxedCell<T> {
54    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
55        f.debug_tuple("RelaxedCell").field(&self.load()).finish()
56    }
57}
58
59pub trait AtomicNative: Sized + Default + fmt::Debug {
60    type Atomic: From<Self>;
61
62    fn load(atomic: &Self::Atomic) -> Self;
63    fn store(atomic: &Self::Atomic, val: Self);
64    fn fetch_add(atomic: &Self::Atomic, val: Self) -> Self;
65    fn fetch_sub(atomic: &Self::Atomic, val: Self) -> Self;
66    fn fetch_max(atomic: &Self::Atomic, val: Self) -> Self;
67    fn get_mut(atomic: &mut Self::Atomic) -> &mut Self;
68}
69
70macro_rules! impl_relaxed_cell {
71    ($T:ty, $new:ident, $A:ty) => {
72        impl RelaxedCell<$T> {
73            // Not part of the trait as it should be const.
74            pub const fn $new(value: $T) -> Self {
75                Self(<$A>::new(value))
76            }
77        }
78
79        impl AtomicNative for $T {
80            type Atomic = $A;
81
82            #[inline(always)]
83            fn load(atomic: &Self::Atomic) -> Self {
84                atomic.load(Ordering::Relaxed)
85            }
86
87            #[inline(always)]
88            fn store(atomic: &Self::Atomic, val: Self) {
89                atomic.store(val, Ordering::Relaxed);
90            }
91
92            #[inline(always)]
93            fn fetch_add(atomic: &Self::Atomic, val: Self) -> Self {
94                atomic.fetch_add(val, Ordering::Relaxed)
95            }
96
97            #[inline(always)]
98            fn fetch_sub(atomic: &Self::Atomic, val: Self) -> Self {
99                atomic.fetch_sub(val, Ordering::Relaxed)
100            }
101
102            #[inline(always)]
103            fn fetch_max(atomic: &Self::Atomic, val: Self) -> Self {
104                atomic.fetch_max(val, Ordering::Relaxed)
105            }
106
107            #[inline(always)]
108            fn get_mut(atomic: &mut Self::Atomic) -> &mut Self {
109                atomic.get_mut()
110            }
111        }
112    };
113}
114
115impl_relaxed_cell!(u8, new_u8, AtomicU8);
116impl_relaxed_cell!(u32, new_u32, AtomicU32);
117impl_relaxed_cell!(u64, new_u64, AtomicU64);
118impl_relaxed_cell!(usize, new_usize, AtomicUsize);
119
120impl RelaxedCell<bool> {
121    // Not part of the trait as it should be const.
122    pub const fn new_bool(value: bool) -> Self {
123        Self(AtomicBool::new(value))
124    }
125
126    #[inline(always)]
127    pub fn fetch_or(&self, val: bool) -> bool {
128        self.0.fetch_or(val, Ordering::Relaxed)
129    }
130}
131
132impl AtomicNative for bool {
133    type Atomic = AtomicBool;
134
135    #[inline(always)]
136    fn load(atomic: &Self::Atomic) -> Self {
137        atomic.load(Ordering::Relaxed)
138    }
139
140    #[inline(always)]
141    fn store(atomic: &Self::Atomic, val: Self) {
142        atomic.store(val, Ordering::Relaxed);
143    }
144
145    #[inline(always)]
146    fn fetch_add(_atomic: &Self::Atomic, _val: Self) -> Self {
147        unimplemented!()
148    }
149
150    #[inline(always)]
151    fn fetch_sub(_atomic: &Self::Atomic, _val: Self) -> Self {
152        unimplemented!()
153    }
154
155    #[inline(always)]
156    fn fetch_max(atomic: &Self::Atomic, val: Self) -> Self {
157        atomic.fetch_or(val, Ordering::Relaxed)
158    }
159
160    #[inline(always)]
161    fn get_mut(atomic: &mut Self::Atomic) -> &mut Self {
162        atomic.get_mut()
163    }
164}
polars_utils/relaxed_cell.rs

polars_utils/
relaxed_cell.rs
