frame-alloc 0.1.0

use crate::{InterruptControl, NoInterruptControl};
use core::cell::UnsafeCell;
use core::marker::PhantomData;
use core::sync::atomic::{AtomicBool, Ordering};

/// Maximum number of `spin_loop` iterations in a single backoff step.
const MAX_BACKOFF: usize = 32;

/// A test-and-test-and-set (TATAS) spinlock with exponential backoff.
///
/// The type parameter `I` selects an [`InterruptControl`] strategy. With the
/// default [`NoInterruptControl`] the lock is an ordinary spinlock and the
/// interrupt machinery compiles away (`I::State` is `()`). Supplying a real
/// strategy makes the lock disable local interrupts for the duration it is held.
pub(crate) struct Lock<I: InterruptControl = NoInterruptControl> {
    locked: AtomicBool,
    /// Interrupt state saved by `acquire`, consumed by `release`. Accessed only
    /// under `locked`.
    saved: UnsafeCell<I::State>,
    _interrupt: PhantomData<fn() -> I>,
}

impl<I: InterruptControl> Lock<I> {
    pub(crate) const fn new() -> Self {
        Self {
            locked: AtomicBool::new(false),
            // Seed only; overwritten by every `acquire` before `release` reads it.
            saved: UnsafeCell::new(I::INIT),
            _interrupt: PhantomData,
        }
    }

    #[inline(always)]
    pub(crate) fn acquire(&self) {
        let state = I::disable();
        let mut backoff = 1usize;
        loop {
            if self
                .locked
                .compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed)
                .is_ok()
            {
                // SAFETY: we hold the lock, so we have exclusive access to `saved`.
                unsafe { *self.saved.get() = state };
                return;
            }
            for _ in 0..backoff {
                core::hint::spin_loop();
            }
            backoff = (backoff * 2).min(MAX_BACKOFF);
            while self.locked.load(Ordering::Relaxed) {
                core::hint::spin_loop();
            }
        }
    }

    #[inline(always)]
    pub(crate) fn release(&self) {
        // SAFETY: we still hold the lock, so reading `saved` is exclusive.
        let state = unsafe { *self.saved.get() };
        self.locked.store(false, Ordering::Release);
        // SAFETY: paired with the `I::disable()` in the matching `acquire`.
        unsafe { I::restore(state) };
    }
}