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

polars_utils/
relaxed_cell.rs
