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

polars_utils/
relaxed_cell.rs
