ax-task 0.8.2

//! Public spin-lock types whose acquisition methods express context policy.

use core::{fmt, ptr};

use self::{
    base::{BaseSpinLock, BaseSpinLockGuard},
    rwlock::{BaseSpinRwLock, BaseSpinRwLockReadGuard, BaseSpinRwLockWriteGuard},
};
use super::context::{GuardState, PreemptIrqSaveState, PreemptState, RawState};

pub(crate) mod atomic;
mod base;
#[cfg(feature = "lockdep")]
pub(crate) mod lockdep;
mod raw;
mod rwlock;

pub use self::raw::RawIrqSaveMutex;

/// A non-sleeping mutual-exclusion lock.
///
/// The lock object does not bake in an execution-context policy. Callers
/// choose the policy at the acquisition site with [`Self::lock`],
/// [`Self::lock_irqsave`], or [`Self::lock_raw`].
#[repr(transparent)]
pub struct RawSpinLock<T: ?Sized>(BaseSpinLock<RawState, T>);

/// A guard returned by [`RawSpinLock::lock`].
pub type RawSpinLockGuard<'a, T> = BaseSpinLockGuard<'a, PreemptState, T>;

/// A guard returned by [`RawSpinLock::lock_irqsave`].
pub type RawSpinLockIrqSaveGuard<'a, T> = BaseSpinLockGuard<'a, PreemptIrqSaveState, T>;

/// A guard returned by [`RawSpinLock::lock_raw`].
pub type RawSpinLockUnpinnedGuard<'a, T> = BaseSpinLockGuard<'a, RawState, T>;

impl<T> RawSpinLock<T> {
    /// Creates an unlocked spin lock.
    #[inline(always)]
    #[track_caller]
    pub const fn new(data: T) -> Self {
        Self(BaseSpinLock::new(data))
    }

    /// Consumes the lock and returns the protected value.
    #[inline(always)]
    pub fn into_inner(self) -> T {
        self.0.into_inner()
    }
}

impl<T: ?Sized> RawSpinLock<T> {
    #[inline(always)]
    fn with_state<G: GuardState>(&self) -> &BaseSpinLock<G, T> {
        // SAFETY: `BaseSpinLock` has a stable C layout, and its guard-state
        // parameter is represented only by `PhantomData`. The atomic state,
        // lockdep map, and protected value therefore have identical addresses
        // for every `G`.
        unsafe { &*(ptr::from_ref(&self.0) as *const BaseSpinLock<G, T>) }
    }

    #[inline(always)]
    fn with_state_mut<G: GuardState>(&mut self) -> &mut BaseSpinLock<G, T> {
        // SAFETY: see `with_state`; the exclusive borrow prevents aliases.
        unsafe { &mut *(ptr::from_mut(&mut self.0) as *mut BaseSpinLock<G, T>) }
    }

    /// Acquires the lock after disabling kernel preemption.
    #[inline(always)]
    #[track_caller]
    pub fn lock(&self) -> RawSpinLockGuard<'_, T> {
        self.with_state::<PreemptState>().lock()
    }

    /// Acquires the lock after disabling preemption, using a lockdep subclass.
    ///
    /// This is intended for structurally nested acquisitions of different
    /// locks with the same class. Without `lockdep`, `subclass` has no effect.
    #[inline(always)]
    #[track_caller]
    pub fn lock_nested(&self, subclass: u32) -> RawSpinLockGuard<'_, T> {
        self.with_state::<PreemptState>().lock_nested(subclass)
    }

    /// Attempts to acquire the lock after disabling kernel preemption.
    #[inline(always)]
    #[track_caller]
    pub fn try_lock(&self) -> Option<RawSpinLockGuard<'_, T>> {
        self.with_state::<PreemptState>().try_lock()
    }

    /// Acquires the lock after disabling preemption and saving/disabling IRQs.
    #[inline(always)]
    #[track_caller]
    pub fn lock_irqsave(&self) -> RawSpinLockIrqSaveGuard<'_, T> {
        self.with_state::<PreemptIrqSaveState>().lock()
    }

    /// Acquires the lock after disabling preemption and saving/disabling IRQs,
    /// using a lockdep subclass.
    ///
    /// This is intended for structurally nested acquisitions of different
    /// locks with the same class. Without `lockdep`, `subclass` has no effect.
    #[inline(always)]
    #[track_caller]
    pub fn lock_irqsave_nested(&self, subclass: u32) -> RawSpinLockIrqSaveGuard<'_, T> {
        self.with_state::<PreemptIrqSaveState>()
            .lock_nested(subclass)
    }

    /// Attempts to acquire the lock after disabling preemption and IRQs.
    #[inline(always)]
    #[track_caller]
    pub fn try_lock_irqsave(&self) -> Option<RawSpinLockIrqSaveGuard<'_, T>> {
        self.with_state::<PreemptIrqSaveState>().try_lock()
    }

    /// Acquires the lock without changing preemption or interrupt state.
    ///
    /// # Safety
    ///
    /// The caller must prevent same-CPU re-entry and all concurrent access
    /// which could violate exclusive access, including on single-core builds
    /// where the atomic lock word is compiled out.
    #[inline(always)]
    #[track_caller]
    pub unsafe fn lock_raw(&self) -> RawSpinLockUnpinnedGuard<'_, T> {
        self.with_state::<RawState>().lock()
    }

    /// Attempts a raw acquisition without changing execution context.
    ///
    /// # Safety
    ///
    /// The caller must uphold the same exclusion contract as
    /// [`Self::lock_raw`], even when this function returns `None`.
    #[inline(always)]
    #[track_caller]
    pub unsafe fn try_lock_raw(&self) -> Option<RawSpinLockUnpinnedGuard<'_, T>> {
        self.with_state::<RawState>().try_lock()
    }

    /// Returns whether the lock appears held.
    ///
    /// This is only a diagnostic snapshot and provides no synchronization.
    #[inline(always)]
    pub fn is_locked(&self) -> bool {
        self.0.is_locked()
    }

    /// Returns mutable access without locking.
    #[inline(always)]
    pub fn get_mut(&mut self) -> &mut T {
        self.with_state_mut::<RawState>().get_mut()
    }

    /// Releases a preemption-mode lock without consuming its guard.
    ///
    /// # Safety
    ///
    /// The caller must own exactly one guard returned by [`Self::lock`] and
    /// must ensure that guard will never subsequently be dropped.
    #[doc(hidden)]
    #[inline(always)]
    pub unsafe fn force_unlock(&self) {
        unsafe { self.with_state::<PreemptState>().force_unlock() };
    }
}

impl<T: Default> Default for RawSpinLock<T> {
    fn default() -> Self {
        Self::new(T::default())
    }
}

impl<T: fmt::Debug> fmt::Debug for RawSpinLock<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self.try_lock() {
            Some(guard) => f
                .debug_struct("RawSpinLock")
                .field("data", &&*guard)
                .finish(),
            None => f
                .debug_struct("RawSpinLock")
                .field("data", &"<locked>")
                .finish(),
        }
    }
}

/// A non-sleeping read-write lock with acquisition-site context policy.
#[repr(transparent)]
pub struct RawSpinRwLock<T: ?Sized>(BaseSpinRwLock<RawState, T>);

/// A read guard returned by [`RawSpinRwLock::read`].
pub type RawSpinRwLockReadGuard<'a, T> = BaseSpinRwLockReadGuard<'a, PreemptState, T>;

/// A write guard returned by [`RawSpinRwLock::write`].
pub type RawSpinRwLockWriteGuard<'a, T> = BaseSpinRwLockWriteGuard<'a, PreemptState, T>;

/// An IRQ-save read guard.
pub type RawSpinRwLockIrqSaveReadGuard<'a, T> = BaseSpinRwLockReadGuard<'a, PreemptIrqSaveState, T>;

/// An IRQ-save write guard.
pub type RawSpinRwLockIrqSaveWriteGuard<'a, T> =
    BaseSpinRwLockWriteGuard<'a, PreemptIrqSaveState, T>;

/// A raw read guard.
pub type RawSpinRwLockUnpinnedReadGuard<'a, T> = BaseSpinRwLockReadGuard<'a, RawState, T>;

/// A raw write guard.
pub type RawSpinRwLockUnpinnedWriteGuard<'a, T> = BaseSpinRwLockWriteGuard<'a, RawState, T>;

impl<T> RawSpinRwLock<T> {
    /// Creates an unlocked spin read-write lock.
    #[inline(always)]
    #[track_caller]
    pub const fn new(data: T) -> Self {
        Self(BaseSpinRwLock::new(data))
    }

    /// Consumes the lock and returns the protected value.
    #[inline(always)]
    pub fn into_inner(self) -> T {
        self.0.into_inner()
    }
}

impl<T: ?Sized> RawSpinRwLock<T> {
    #[inline(always)]
    fn with_state<G: GuardState>(&self) -> &BaseSpinRwLock<G, T> {
        // SAFETY: identical to `RawSpinLock::with_state`; the generic parameter
        // is represented only by `PhantomData` in a stable C layout.
        unsafe { &*(ptr::from_ref(&self.0) as *const BaseSpinRwLock<G, T>) }
    }

    #[inline(always)]
    fn with_state_mut<G: GuardState>(&mut self) -> &mut BaseSpinRwLock<G, T> {
        // SAFETY: see `with_state`; the exclusive borrow prevents aliases.
        unsafe { &mut *(ptr::from_mut(&mut self.0) as *mut BaseSpinRwLock<G, T>) }
    }

    /// Acquires a read guard after disabling preemption.
    #[inline(always)]
    #[track_caller]
    pub fn read(&self) -> RawSpinRwLockReadGuard<'_, T> {
        self.with_state::<PreemptState>().read()
    }

    /// Attempts a read acquisition after disabling preemption.
    #[inline(always)]
    #[track_caller]
    pub fn try_read(&self) -> Option<RawSpinRwLockReadGuard<'_, T>> {
        self.with_state::<PreemptState>().try_read()
    }

    /// Acquires a write guard after disabling preemption.
    #[inline(always)]
    #[track_caller]
    pub fn write(&self) -> RawSpinRwLockWriteGuard<'_, T> {
        self.with_state::<PreemptState>().write()
    }

    /// Attempts a write acquisition after disabling preemption.
    #[inline(always)]
    #[track_caller]
    pub fn try_write(&self) -> Option<RawSpinRwLockWriteGuard<'_, T>> {
        self.with_state::<PreemptState>().try_write()
    }

    /// Acquires an IRQ-save read guard.
    #[inline(always)]
    #[track_caller]
    pub fn read_irqsave(&self) -> RawSpinRwLockIrqSaveReadGuard<'_, T> {
        self.with_state::<PreemptIrqSaveState>().read()
    }

    /// Attempts an IRQ-save read acquisition.
    #[inline(always)]
    #[track_caller]
    pub fn try_read_irqsave(&self) -> Option<RawSpinRwLockIrqSaveReadGuard<'_, T>> {
        self.with_state::<PreemptIrqSaveState>().try_read()
    }

    /// Acquires an IRQ-save write guard.
    #[inline(always)]
    #[track_caller]
    pub fn write_irqsave(&self) -> RawSpinRwLockIrqSaveWriteGuard<'_, T> {
        self.with_state::<PreemptIrqSaveState>().write()
    }

    /// Attempts an IRQ-save write acquisition.
    #[inline(always)]
    #[track_caller]
    pub fn try_write_irqsave(&self) -> Option<RawSpinRwLockIrqSaveWriteGuard<'_, T>> {
        self.with_state::<PreemptIrqSaveState>().try_write()
    }

    /// Acquires a raw read guard without changing execution context.
    ///
    /// # Safety
    ///
    /// The caller must prevent re-entry and uphold the read-side exclusion
    /// contract, including on single-core builds.
    #[inline(always)]
    #[track_caller]
    pub unsafe fn read_raw(&self) -> RawSpinRwLockUnpinnedReadGuard<'_, T> {
        self.with_state::<RawState>().read()
    }

    /// Attempts a raw read acquisition.
    ///
    /// # Safety
    ///
    /// The caller must uphold the contract of [`Self::read_raw`].
    #[inline(always)]
    #[track_caller]
    pub unsafe fn try_read_raw(&self) -> Option<RawSpinRwLockUnpinnedReadGuard<'_, T>> {
        self.with_state::<RawState>().try_read()
    }

    /// Acquires a raw write guard without changing execution context.
    ///
    /// # Safety
    ///
    /// The caller must prevent re-entry and concurrent readers or writers,
    /// including on single-core builds.
    #[inline(always)]
    #[track_caller]
    pub unsafe fn write_raw(&self) -> RawSpinRwLockUnpinnedWriteGuard<'_, T> {
        self.with_state::<RawState>().write()
    }

    /// Attempts a raw write acquisition.
    ///
    /// # Safety
    ///
    /// The caller must uphold the contract of [`Self::write_raw`].
    #[inline(always)]
    #[track_caller]
    pub unsafe fn try_write_raw(&self) -> Option<RawSpinRwLockUnpinnedWriteGuard<'_, T>> {
        self.with_state::<RawState>().try_write()
    }

    /// Returns mutable access without locking.
    #[inline(always)]
    pub fn get_mut(&mut self) -> &mut T {
        self.with_state_mut::<RawState>().get_mut()
    }

    /// Removes one deliberately leaked raw read guard from the reader count.
    ///
    /// # Safety
    ///
    /// The caller must own a deliberately forgotten guard returned by
    /// [`Self::read_raw`] and must prove that no live reference from it remains.
    #[doc(hidden)]
    #[inline(always)]
    pub unsafe fn force_read_decrement_raw(&self) {
        unsafe {
            self.with_state::<RawState>().force_read_decrement();
        }
    }
}

impl<T: Default> Default for RawSpinRwLock<T> {
    fn default() -> Self {
        Self::new(T::default())
    }
}

impl<T> From<T> for RawSpinRwLock<T> {
    fn from(value: T) -> Self {
        Self::new(value)
    }
}

impl<T: fmt::Debug> fmt::Debug for RawSpinRwLock<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self.try_read() {
            Some(guard) => f
                .debug_struct("RawSpinRwLock")
                .field("data", &&*guard)
                .finish(),
            None => f
                .debug_struct("RawSpinRwLock")
                .field("data", &"<write locked>")
                .finish(),
        }
    }
}