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!(usize, new_usize, AtomicUsize);
130
131impl RelaxedCell<bool> {
132    // Not part of the trait as it should be const.
133    pub const fn new_bool(value: bool) -> Self {
134        Self(AtomicBool::new(value))
135    }
136
137    #[inline(always)]
138    pub fn fetch_or(&self, val: bool) -> bool {
139        self.0.fetch_or(val, Ordering::Relaxed)
140    }
141}
142
143impl AtomicNative for bool {
144    type Atomic = AtomicBool;
145
146    #[inline(always)]
147    fn load(atomic: &Self::Atomic) -> Self {
148        atomic.load(Ordering::Relaxed)
149    }
150
151    #[inline(always)]
152    fn store(atomic: &Self::Atomic, val: Self) {
153        atomic.store(val, Ordering::Relaxed);
154    }
155
156    #[inline(always)]
157    fn fetch_add(_atomic: &Self::Atomic, _val: Self) -> Self {
158        unimplemented!()
159    }
160
161    #[inline(always)]
162    fn fetch_sub(_atomic: &Self::Atomic, _val: Self) -> Self {
163        unimplemented!()
164    }
165
166    #[inline(always)]
167    fn fetch_max(atomic: &Self::Atomic, val: Self) -> Self {
168        atomic.fetch_or(val, Ordering::Relaxed)
169    }
170
171    #[inline(always)]
172    fn get_mut(atomic: &mut Self::Atomic) -> &mut Self {
173        atomic.get_mut()
174    }
175
176    #[inline(always)]
177    fn swap(atomic: &Self::Atomic, val: Self) -> Self {
178        atomic.swap(val, Ordering::Relaxed)
179    }
180}
polars_utils/relaxed_cell.rs

polars_utils/
relaxed_cell.rs
