1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

use atomic::Atomic as BasicAtomic;
use core::sync::atomic::Ordering;
use futures::task::AtomicWaker;

/// Atomic value that also contains [`Waker`](`core::task::Waker`) to notify subscriber asynchronously.
#[derive(Default, Debug)]
pub struct Atomic<T: Copy> {
    pub(crate) value: BasicAtomic<T>,
    pub(crate) waker: AtomicWaker,
}

impl<T: Copy> Atomic<T> {
    pub const fn new(value: T) -> Self {
        Self {
            value: BasicAtomic::new(value),
            waker: AtomicWaker::new(),
        }
    }

    pub fn load(&self) -> T {
        self.value.load(Ordering::Acquire)
    }
    pub fn store(&self, val: T) {
        self.value.store(val, Ordering::Release);
        self.waker.wake();
    }
    pub fn swap(&self, val: T) -> T {
        let old = self.value.swap(val, Ordering::AcqRel);
        self.waker.wake();
        old
    }
    pub fn compare_exchange(&self, current: T, new: T) -> Result<T, T> {
        self.value
            .compare_exchange(current, new, Ordering::AcqRel, Ordering::Acquire)
            .map(|x| {
                self.waker.wake();
                x
            })
    }
    pub fn fetch_update<F: FnMut(T) -> Option<T>>(&self, f: F) -> Result<T, T> {
        self.value
            .fetch_update(Ordering::AcqRel, Ordering::Acquire, f)
            .map(|x| {
                self.waker.wake();
                x
            })
    }
}

macro_rules! impl_atomic_bitwise {
    ($T:ty) => {
        impl Atomic<$T> {
            pub fn fetch_and(&self, val: $T) -> $T {
                let old = self.value.fetch_and(val, Ordering::AcqRel);
                self.waker.wake();
                old
            }
            pub fn fetch_or(&self, val: $T) -> $T {
                let old = self.value.fetch_or(val, Ordering::AcqRel);
                self.waker.wake();
                old
            }
            pub fn fetch_xor(&self, val: $T) -> $T {
                let old = self.value.fetch_xor(val, Ordering::AcqRel);
                self.waker.wake();
                old
            }
        }
    };
}

macro_rules! impl_atomic_num {
    ($T:ty) => {
        impl_atomic_bitwise!($T);

        impl Atomic<$T> {
            pub fn fetch_add(&self, val: $T) -> $T {
                let old = self.value.fetch_add(val, Ordering::AcqRel);
                self.waker.wake();
                old
            }
            pub fn fetch_sub(&self, val: $T) -> $T {
                let old = self.value.fetch_sub(val, Ordering::AcqRel);
                self.waker.wake();
                old
            }
            pub fn fetch_max(&self, val: $T) -> $T {
                let old = self.value.fetch_max(val, Ordering::AcqRel);
                self.waker.wake();
                old
            }
            pub fn fetch_min(&self, val: $T) -> $T {
                let old = self.value.fetch_min(val, Ordering::AcqRel);
                self.waker.wake();
                old
            }
        }
    };
}

impl_atomic_bitwise!(bool);

impl_atomic_num!(u8);
impl_atomic_num!(u16);
impl_atomic_num!(u32);
impl_atomic_num!(u64);
impl_atomic_num!(usize);

impl_atomic_num!(i8);
impl_atomic_num!(i16);
impl_atomic_num!(i32);
impl_atomic_num!(i64);
impl_atomic_num!(isize);

impl<T: Copy> AsRef<Atomic<T>> for Atomic<T> {
    fn as_ref(&self) -> &Atomic<T> {
        self
    }
}