shared_vec/
counter.rs

1//! Counter trait and its implementations for various integer types.
2//!
3//! See [alloc::sync::Arc].
4
5use core::cell::Cell;
6use core::sync::atomic;
7use core::sync::atomic::{AtomicU16, AtomicU32, AtomicU64, AtomicU8, AtomicUsize, Ordering};
8
9/// Counter trait for uniformed reference counting.
10///
11/// # Safety
12///
13/// The implementation must hold arithmetic invariants and respect synchronization.
14pub unsafe trait Counter<T>: Sized {
15    /// Increment the counter by 1.
16    ///
17    /// # Safety
18    ///
19    /// The caller must ensure that the counter will not overflow.
20    unsafe fn increment(&self);
21    /// Decrement the counter by 1.
22    ///
23    /// # Returns
24    ///
25    /// `true` if the counter reaches zero.
26    ///
27    /// # Safety
28    ///
29    /// The caller must ensure that the counter will not underflow.
30    ///
31    /// [`Counter::fence_acquire`] must be called after this if it returns `true`.
32    unsafe fn decrement(&self) -> bool;
33    /// Issue an acquire fence.
34    fn fence_acquire();
35    /// Create a counter with initial value of one.
36    fn one() -> Self;
37}
38
39macro_rules! impl_cell {
40    ($ty:ty) => {
41        unsafe impl Counter<$ty> for Cell<$ty> {
42            unsafe fn increment(&self) {
43                self.set(self.get() + 1);
44            }
45
46            unsafe fn decrement(&self) -> bool {
47                let new = self.get() - 1;
48                self.set(new);
49                new == 0
50            }
51
52            fn fence_acquire() {}
53
54            fn one() -> Self {
55                Cell::new(1)
56            }
57        }
58    };
59}
60
61impl_cell!(u8);
62impl_cell!(u16);
63impl_cell!(u32);
64impl_cell!(u64);
65impl_cell!(usize);
66
67macro_rules! impl_atomic {
68    ($ty:ty, $equiv:ty) => {
69        unsafe impl Counter<$equiv> for $ty {
70            unsafe fn increment(&self) {
71                self.fetch_add(1, Ordering::Relaxed);
72            }
73
74            unsafe fn decrement(&self) -> bool {
75                let old = self.fetch_sub(1, Ordering::Release);
76                debug_assert!(old != 0);
77                old == 1
78            }
79
80            fn fence_acquire() {
81                atomic::fence(Ordering::Acquire);
82            }
83
84            fn one() -> Self {
85                Self::new(1)
86            }
87        }
88    };
89}
90
91impl_atomic!(AtomicU8, u8);
92impl_atomic!(AtomicU16, u16);
93impl_atomic!(AtomicU32, u32);
94impl_atomic!(AtomicU64, u64);
95impl_atomic!(AtomicUsize, usize);