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
85        impl AtomicNative for $T {
86            type Atomic = $A;
87
88            #[inline(always)]
89            fn load(atomic: &Self::Atomic) -> Self {
90                atomic.load(Ordering::Relaxed)
91            }
92
93            #[inline(always)]
94            fn store(atomic: &Self::Atomic, val: Self) {
95                atomic.store(val, Ordering::Relaxed);
96            }
97
98            #[inline(always)]
99            fn fetch_add(atomic: &Self::Atomic, val: Self) -> Self {
100                atomic.fetch_add(val, Ordering::Relaxed)
101            }
102
103            #[inline(always)]
104            fn fetch_sub(atomic: &Self::Atomic, val: Self) -> Self {
105                atomic.fetch_sub(val, Ordering::Relaxed)
106            }
107
108            #[inline(always)]
109            fn fetch_max(atomic: &Self::Atomic, val: Self) -> Self {
110                atomic.fetch_max(val, Ordering::Relaxed)
111            }
112
113            #[inline(always)]
114            fn get_mut(atomic: &mut Self::Atomic) -> &mut Self {
115                atomic.get_mut()
116            }
117
118            #[inline(always)]
119            fn swap(atomic: &Self::Atomic, val: Self) -> Self {
120                atomic.swap(val, Ordering::Relaxed)
121            }
122        }
123    };
124}
125
126impl_relaxed_cell!(u8, new_u8, AtomicU8);
127impl_relaxed_cell!(u32, new_u32, AtomicU32);
128impl_relaxed_cell!(u64, new_u64, AtomicU64);
129impl_relaxed_cell!(i64, new_i64, AtomicI64);
130impl_relaxed_cell!(usize, new_usize, AtomicUsize);
131
132impl RelaxedCell<bool> {
133    // Not part of the trait as it should be const.
134    pub const fn new_bool(value: bool) -> Self {
135        Self(AtomicBool::new(value))
136    }
137
138    #[inline(always)]
139    pub fn fetch_or(&self, val: bool) -> bool {
140        self.0.fetch_or(val, Ordering::Relaxed)
141    }
142}
143
144impl AtomicNative for bool {
145    type Atomic = AtomicBool;
146
147    #[inline(always)]
148    fn load(atomic: &Self::Atomic) -> Self {
149        atomic.load(Ordering::Relaxed)
150    }
151
152    #[inline(always)]
153    fn store(atomic: &Self::Atomic, val: Self) {
154        atomic.store(val, Ordering::Relaxed);
155    }
156
157    #[inline(always)]
158    fn fetch_add(_atomic: &Self::Atomic, _val: Self) -> Self {
159        unimplemented!()
160    }
161
162    #[inline(always)]
163    fn fetch_sub(_atomic: &Self::Atomic, _val: Self) -> Self {
164        unimplemented!()
165    }
166
167    #[inline(always)]
168    fn fetch_max(atomic: &Self::Atomic, val: Self) -> Self {
169        atomic.fetch_or(val, Ordering::Relaxed)
170    }
171
172    #[inline(always)]
173    fn get_mut(atomic: &mut Self::Atomic) -> &mut Self {
174        atomic.get_mut()
175    }
176
177    #[inline(always)]
178    fn swap(atomic: &Self::Atomic, val: Self) -> Self {
179        atomic.swap(val, Ordering::Relaxed)
180    }
181}
polars_utils/relaxed_cell.rs

polars_utils/
relaxed_cell.rs
