shared-vec 0.1.0

//! Counter trait and its implementations for various integer types.
//!
//! See [alloc::sync::Arc].

use core::cell::Cell;
use core::sync::atomic;
use core::sync::atomic::{AtomicU16, AtomicU32, AtomicU64, AtomicU8, AtomicUsize, Ordering};

/// Counter trait for uniformed reference counting.
///
/// # Safety
///
/// The implementation must hold arithmetic invariants and respect synchronization.
pub unsafe trait Counter<T>: Sized {
    /// Increment the counter by 1.
    ///
    /// # Safety
    ///
    /// The caller must ensure that the counter will not overflow.
    unsafe fn increment(&self);
    /// Decrement the counter by 1.
    ///
    /// # Returns
    ///
    /// `true` if the counter reaches zero.
    ///
    /// # Safety
    ///
    /// The caller must ensure that the counter will not underflow.
    ///
    /// [`Counter::fence_acquire`] must be called after this if it returns `true`.
    unsafe fn decrement(&self) -> bool;
    /// Issue an acquire fence.
    fn fence_acquire();
    /// Create a counter with initial value of one.
    fn one() -> Self;
}

macro_rules! impl_cell {
    ($ty:ty) => {
        unsafe impl Counter<$ty> for Cell<$ty> {
            unsafe fn increment(&self) {
                self.set(self.get() + 1);
            }

            unsafe fn decrement(&self) -> bool {
                let new = self.get() - 1;
                self.set(new);
                new == 0
            }

            fn fence_acquire() {}

            fn one() -> Self {
                Cell::new(1)
            }
        }
    };
}

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

macro_rules! impl_atomic {
    ($ty:ty, $equiv:ty) => {
        unsafe impl Counter<$equiv> for $ty {
            unsafe fn increment(&self) {
                self.fetch_add(1, Ordering::Relaxed);
            }

            unsafe fn decrement(&self) -> bool {
                let old = self.fetch_sub(1, Ordering::Release);
                debug_assert!(old != 0);
                old == 1
            }

            fn fence_acquire() {
                atomic::fence(Ordering::Acquire);
            }

            fn one() -> Self {
                Self::new(1)
            }
        }
    };
}

impl_atomic!(AtomicU8, u8);
impl_atomic!(AtomicU16, u16);
impl_atomic!(AtomicU32, u32);
impl_atomic!(AtomicU64, u64);
impl_atomic!(AtomicUsize, usize);