meowalloc 0.1.1

use core::cell::UnsafeCell;
use core::sync::atomic::{AtomicBool, Ordering};
use core::hint;
use core::ops::{Deref, DerefMut};


pub struct SpinLock<T> {
    inner: UnsafeCell<T>,
    locked: AtomicBool
}

impl<T> SpinLock<T> {
    pub fn lock(&self) -> SpinLockGuard<'_, T> {
        while let Err(_) = self.locked.compare_exchange_weak(
            false, true, Ordering::SeqCst, Ordering::SeqCst
        ) { hint::spin_loop() }
        SpinLockGuard { inner: &self }
    }
    pub fn try_lock(&self) -> Option<SpinLockGuard<'_, T>> {
        if let Err(_) = self.locked.compare_exchange_weak(
            false, true, Ordering::SeqCst, Ordering::SeqCst
        ) { return None }
        Some(SpinLockGuard { inner: &self })
    }
}

pub struct SpinLockGuard<'a, T> {
    inner: &'a SpinLock<T>
}

impl<'a, T> Deref for SpinLockGuard<'a, T> {
    type Target = T;
    fn deref(&self) -> &Self::Target {
        unsafe { self.inner.inner.get().as_ref_unchecked() }
    }
}

impl<'a, T> DerefMut for SpinLockGuard<'a, T> {
    fn deref_mut(&mut self) -> &mut T {
        unsafe { self.inner.inner.get().as_mut_unchecked() }
    }
}

impl<'a, T> Drop for SpinLockGuard<'a, T> {
    fn drop(&mut self) {
        self.inner.locked.store(false, Ordering::SeqCst);
    }
}

impl<T> SpinLock<T> {
    pub const fn new(inner: T) -> Self {
        Self { inner: UnsafeCell::new(inner), locked: const { AtomicBool::new(false) } }
    }
}

#[test]
fn try_lock() {
    let res = SpinLock::new(1);
    {
        let mut h1 = res.try_lock().unwrap();
        *h1 = 2;
        assert!(res.try_lock().is_none())
    }
    *res.try_lock().unwrap() = 3;
    assert_eq!(*res.try_lock().unwrap(), 3);
}