Documentation
//! Atomic enums.

#![no_std]

use core::{
    fmt,
    mem::{self, MaybeUninit},
    ptr,
    sync::atomic::*,
};

#[macro_export]
macro_rules! atomium {
    (
        #[repr($repr:ident)]
        $(#[$enum_meta:meta])*
        $vis:vis enum $ident:ident {
            $(
                $(#[$variant_meta:meta])*
                $variant:ident $(= $value:expr)?
            ),* $(,)?
        }
    ) => {
        #[repr($repr)]
        $(#[$enum_meta])*
        $vis enum $ident {
            $(
                $(#[$variant_meta])*
                $variant $(= $value)?
            ),*
        }

        unsafe impl $crate::Atomium for $ident {
            type Repr = ::core::primitive::$repr;
        }
    };
}

#[repr(transparent)]
pub struct Atomic<T: Atomium>(AtomicOf<T::Repr>);

impl<T: Atomium + Default + Copy> Default for Atomic<T> {
    fn default() -> Self {
        Self::new(T::default())
    }
}
impl<T: Atomium + fmt::Debug> fmt::Debug for Atomic<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.load(Ordering::Relaxed).fmt(f)
    }
}
impl<T: Atomium + Copy> From<T> for Atomic<T> {
    fn from(value: T) -> Self {
        Self::new(value)
    }
}

impl<T: Atomium> Atomic<T> {
    const fn to_repr(value: T) -> T::Repr {
        let mut slot = MaybeUninit::<T::Repr>::uninit();
        unsafe { ptr::write(slot.as_mut_ptr().cast(), value) };
        unsafe { slot.assume_init() }
    }
    const fn from_repr(repr: T::Repr) -> T {
        let mut slot = MaybeUninit::<T>::uninit();
        unsafe { ptr::write(slot.as_mut_ptr().cast(), repr) };
        unsafe { slot.assume_init() }
    }

    pub const fn new(value: T) -> Self
    where
        T: Copy,
    {
        let mut slot = MaybeUninit::<AtomicOf<T::Repr>>::uninit();
        unsafe { ptr::write(slot.as_mut_ptr().cast(), value) };
        Self(unsafe { slot.assume_init() })
    }

    /// # Safety
    /// - See [`AtomicU8::from_ptr`].
    pub const unsafe fn from_ptr<'a>(ptr: *mut T) -> &'a Self {
        #[expect(clippy::transmute_ptr_to_ref)]
        unsafe {
            mem::transmute::<*mut T, &Self>(ptr)
        }
    }

    pub fn get_mut(&mut self) -> &mut T {
        unsafe { mem::transmute::<&mut Self, &mut T>(self) }
    }

    pub const fn into_inner(self) -> T {
        let mut slot = MaybeUninit::<T>::uninit();
        unsafe { ptr::write(slot.as_mut_ptr().cast(), self) };
        unsafe { slot.assume_init() }
    }

    pub fn load(&self, order: Ordering) -> T {
        Self::from_repr(<AtomicOf<T::Repr> as AtomicPrimitive>::load(&self.0, order))
    }

    pub fn store(&self, value: T, order: Ordering) {
        <AtomicOf<T::Repr> as AtomicPrimitive>::store(&self.0, Self::to_repr(value), order)
    }

    pub fn swap(&self, value: T, order: Ordering) -> T {
        Self::from_repr(<AtomicOf<T::Repr> as AtomicPrimitive>::swap(
            &self.0,
            Self::to_repr(value),
            order,
        ))
    }

    pub fn compare_exchange(
        &self,
        current: T,
        new: T,
        success: Ordering,
        failure: Ordering,
    ) -> Result<T, T> {
        <AtomicOf<T::Repr> as AtomicPrimitive>::compare_exchange(
            &self.0,
            Self::to_repr(current),
            Self::to_repr(new),
            success,
            failure,
        )
        .map(Self::from_repr)
        .map_err(Self::from_repr)
    }

    pub fn compare_exchange_weak(
        &self,
        current: T,
        new: T,
        success: Ordering,
        failure: Ordering,
    ) -> Result<T, T> {
        <AtomicOf<T::Repr> as AtomicPrimitive>::compare_exchange_weak(
            &self.0,
            Self::to_repr(current),
            Self::to_repr(new),
            success,
            failure,
        )
        .map(Self::from_repr)
        .map_err(Self::from_repr)
    }

    pub fn try_update(
        &self,
        set_order: Ordering,
        fetch_order: Ordering,
        mut f: impl FnMut(T) -> Option<T>,
    ) -> Result<T, T> {
        <AtomicOf<T::Repr> as AtomicPrimitive>::try_update(
            &self.0,
            set_order,
            fetch_order,
            |repr| f(Self::from_repr(repr)).map(Self::to_repr),
        )
        .map(Self::from_repr)
        .map_err(Self::from_repr)
    }

    pub fn update(
        &self,
        set_order: Ordering,
        fetch_order: Ordering,
        mut f: impl FnMut(T) -> T,
    ) -> T {
        Self::from_repr(<AtomicOf<T::Repr> as AtomicPrimitive>::update(
            &self.0,
            set_order,
            fetch_order,
            |repr| Self::to_repr(f(Self::from_repr(repr))),
        ))
    }

    pub const fn as_ptr(&self) -> *mut T {
        self as *const Self as *mut Self as *mut T
    }
}

/// # Safety
/// - Must be a C-like enum with the given `#[repr(..)]`.
pub unsafe trait Atomium {
    type Repr: PrimitiveAtomic;
}

#[expect(clippy::missing_safety_doc)]
pub unsafe trait AtomicPrimitive {
    type Primitive: PrimitiveAtomic<Atomic = Self>;
    fn into_primitive(self) -> Self::Primitive;
    fn compare_exchange(
        &self,
        current: Self::Primitive,
        new: Self::Primitive,
        success: Ordering,
        failure: Ordering,
    ) -> Result<Self::Primitive, Self::Primitive>;
    fn compare_exchange_weak(
        &self,
        current: Self::Primitive,
        new: Self::Primitive,
        success: Ordering,
        failure: Ordering,
    ) -> Result<Self::Primitive, Self::Primitive>;
    fn load(&self, order: Ordering) -> Self::Primitive;
    fn store(&self, value: Self::Primitive, order: Ordering);
    fn swap(&self, value: Self::Primitive, order: Ordering) -> Self::Primitive;
    fn try_update(
        &self,
        set_order: Ordering,
        fetch_order: Ordering,
        f: impl FnMut(Self::Primitive) -> Option<Self::Primitive>,
    ) -> Result<Self::Primitive, Self::Primitive>;
    fn update(
        &self,
        set_order: Ordering,
        fetch_order: Ordering,
        f: impl FnMut(Self::Primitive) -> Self::Primitive,
    ) -> Self::Primitive;
}

#[expect(clippy::missing_safety_doc)]
pub unsafe trait PrimitiveAtomic {
    type Atomic: AtomicPrimitive<Primitive = Self>;
}

type AtomicOf<T> = <T as PrimitiveAtomic>::Atomic;

macro_rules! types {
    ($(
        $(#[$meta:meta])*
        $prim:ty: $atom:ty
    );* $(;)?) => {
        $(
            $(#[$meta])*
            unsafe impl PrimitiveAtomic for $prim {
                type Atomic = $atom;
            }
            $(#[$meta])*
            unsafe impl AtomicPrimitive for $atom {
                type Primitive = $prim;
                fn into_primitive(self) -> Self::Primitive {
                    self.into_inner()
                }
                fn compare_exchange(
                    &self,
                    current: Self::Primitive,
                    new: Self::Primitive,
                    success: Ordering,
                    failure: Ordering,
                ) -> Result<Self::Primitive, Self::Primitive> {
                    self.compare_exchange(current, new, success, failure)
                }
                fn compare_exchange_weak(
                    &self,
                    current: Self::Primitive,
                    new: Self::Primitive,
                    success: Ordering,
                    failure: Ordering,
                ) -> Result<Self::Primitive, Self::Primitive> {
                    self.compare_exchange_weak(current, new, success, failure)
                }
                fn load(&self, order: Ordering) -> Self::Primitive {
                    self.load(order)
                }
                fn store(&self, value: Self::Primitive, order: Ordering) {
                    self.store(value, order)
                }
                fn swap(&self, value: Self::Primitive, order: Ordering) -> Self::Primitive {
                    self.swap(value, order)
                }
                fn try_update(
                    &self,
                    set_order: Ordering,
                    fetch_order: Ordering,
                    f: impl FnMut(Self::Primitive) -> Option<Self::Primitive>,
                ) -> Result<Self::Primitive, Self::Primitive> {
                    self.try_update(set_order, fetch_order, f)
                }
                fn update(
                    &self,
                    set_order: Ordering,
                    fetch_order: Ordering,
                    f: impl FnMut(Self::Primitive) -> Self::Primitive,
                ) -> Self::Primitive {
                    self.update(set_order, fetch_order, f)
                }
            }
        )*
    };
}

types! {
    #[cfg(target_has_atomic = "8")]
    u8: AtomicU8;
    #[cfg(target_has_atomic = "16")]
    u16: AtomicU16;
    #[cfg(target_has_atomic = "32")]
    u32: AtomicU32;
    #[cfg(target_has_atomic = "64")]
    u64: AtomicU64;
    #[cfg(target_has_atomic = "ptr")]
    usize: AtomicUsize;

    #[cfg(target_has_atomic = "8")]
    i8: AtomicI8;
    #[cfg(target_has_atomic = "16")]
    i16: AtomicI16;
    #[cfg(target_has_atomic = "32")]
    i32: AtomicI32;
    #[cfg(target_has_atomic = "64")]
    i64: AtomicI64;
    #[cfg(target_has_atomic = "ptr")]
    isize: AtomicIsize;

    #[cfg(target_has_atomic = "8")]
    bool: AtomicBool;
}