zyx 0.15.0

Zyx machine learning library
Documentation
// Copyright (C) 2025 zk4x
// SPDX-License-Identifier: LGPL-3.0-only

//! Simple implementation of mutex based on spinlock.

/*use std::sync::{LockResult, MutexGuard};
pub struct Mutex<T>(std::sync::Mutex<T>);

impl<T> Mutex<T> {
    pub const fn new(data: T) -> Self {
        Self(std::sync::Mutex::new(data))
    }

    pub fn lock(&self) -> MutexGuard<'_, T> {
        match self.0.lock() {
            Ok(guard) => guard,
            Err(err) => {
                println!("Poisoned mutex, forcefully recovering");
                err.into_inner()
            }
        }
    }

    pub fn try_lock(&self) -> LockResult<MutexGuard<'_, T>> {
        self.0.lock()
    }
}*/

// Good for debugging too
/*pub struct Mutex<T>(parking_lot::Mutex<T>);

impl<T> Mutex<T> {
    pub const fn new(data: T) -> Self {
        Self(parking_lot::Mutex::new(data))
    }

    pub fn lock(&self) -> parking_lot::MutexGuard<'_, T> {
        self.0.lock()
    }

    pub fn try_lock(&self) -> Result<parking_lot::MutexGuard<'_, T>, ()> {
        Ok(self.0.lock())
    }
}*/

// Spinlock is better for debugging, but std::mutex::Mutex is better for release
use std::{
    cell::UnsafeCell,
    sync::atomic::{AtomicBool, Ordering},
};

// Standard spinlock.
#[derive(Debug)]
pub struct Mutex<T> {
    data: UnsafeCell<T>,
    lock: AtomicBool,
}

#[derive(Debug)]
pub struct MutexGuard<'a, T: 'a> {
    mutex: &'a Mutex<T>,
}

unsafe impl<T: Send> Sync for Mutex<T> {}
unsafe impl<T: Send> Send for Mutex<T> {}

unsafe impl<T: Sync> Sync for MutexGuard<'_, T> {}
unsafe impl<T: Send> Send for MutexGuard<'_, T> {}

impl<T> Mutex<T> {
    pub(super) const fn new(data: T) -> Self {
        Self { data: UnsafeCell::new(data), lock: AtomicBool::new(false) }
    }

    pub(super) fn lock(&self) -> MutexGuard<'_, T> {
        loop {
            if self
                .lock
                .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
                .is_ok()
            {
                return MutexGuard { mutex: self };
            }

            while self.lock.load(Ordering::Relaxed) {
                //core::sync::atomic::spin_loop_hint();
                //std::thread::sleep(std::time::Duration::from_secs(1));
                core::hint::spin_loop();
                /*i += 1;
                debug_assert!(
                    i > N,
                    "Failed to unlock mutex after {N} tries. Panicking in order to avoid deadlock."
                );*/
            }
            /*debug_assert!(
                i > N,
                "Failed to unlock mutex after million tries. Panicking in order to avoid deadlock."
            );*/
        }
    }

    pub(super) fn try_lock(&self) -> Result<MutexGuard<'_, T>, ()> {
        loop {
            if self
                .lock
                .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
                .is_ok()
            {
                return Ok(MutexGuard { mutex: self });
            }

            while self.lock.load(Ordering::Relaxed) {
                //core::sync::atomic::spin_loop_hint();
                //std::thread::sleep(std::time::Duration::from_secs(1));
                core::hint::spin_loop();
                //i += 1;
                /*if i > N {
                    return Err(());
                }*/
            }
        }
    }
}

impl<T> core::ops::Deref for MutexGuard<'_, T> {
    type Target = T;
    fn deref(&self) -> &T {
        unsafe { &*self.mutex.data.get() }
    }
}

impl<T> core::ops::DerefMut for MutexGuard<'_, T> {
    fn deref_mut(&mut self) -> &mut T {
        unsafe { &mut *self.mutex.data.get() }
    }
}

impl<T> Drop for MutexGuard<'_, T> {
    fn drop(&mut self) {
        self.mutex.lock.store(false, Ordering::Release);
    }
}