tipc_std/

sync.rs

// SPDX-License-Identifier: Apache-2.0
// Copyright 2026 KylinSoft Co., Ltd. <https://www.kylinos.cn/>
// See LICENSES for license details.

//! Synchronization primitives for tipc-std.
//!
//! Provides multi-thread-safe `Mutex`, `OnceLock`, and `LazyLock`
//! using atomics, suitable for both x-kernel and LK environments.

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

/// A mutual exclusion primitive based on `AtomicBool` spinlock.
///
/// Safe for multi-threaded environments.
pub struct Mutex<T> {
    locked: AtomicBool,
    data: UnsafeCell<T>,
}

/// A guard that provides access to the data protected by a [`Mutex`].
pub struct MutexGuard<'a, T> {
    mutex: &'a Mutex<T>,
}

impl<T> Mutex<T> {
    /// Creates a new mutex in an unlocked state.
    pub const fn new(data: T) -> Self {
        Self { locked: AtomicBool::new(false), data: UnsafeCell::new(data) }
    }

    /// Acquires the mutex, blocking until it is available.
    ///
    /// Returns `Ok` if the lock was acquired. In this implementation
    /// the mutex is never poisoned, so `Ok` is always returned.
    pub fn lock(&self) -> Result<MutexGuard<'_, T>, PoisonError<MutexGuard<'_, T>>> {
        while self.locked.swap(true, Ordering::Acquire) {
            core::hint::spin_loop();
        }
        Ok(MutexGuard { mutex: self })
    }
}

// SAFETY: Mutex provides exclusive access via spinlock.
unsafe impl<T: Send> Sync for Mutex<T> {}
unsafe impl<T: Send> Send for Mutex<T> {}

impl<T> Deref for MutexGuard<'_, T> {
    type Target = T;

    fn deref(&self) -> &T {
        // SAFETY: We hold the lock, so we have exclusive access.
        unsafe { &*self.mutex.data.get() }
    }
}

impl<T> DerefMut for MutexGuard<'_, T> {
    fn deref_mut(&mut self) -> &mut T {
        // SAFETY: We hold the lock, so we have exclusive access.
        unsafe { &mut *self.mutex.data.get() }
    }
}

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

/// Error type for source-level compatibility with `std::sync::PoisonError`.
///
/// This implementation never actually poisons.
pub struct PoisonError<T> {
    _guard: T,
}

impl<T> PoisonError<T> {
    #[allow(dead_code)]
    pub fn new(guard: T) -> Self {
        PoisonError { _guard: guard }
    }
}

impl<T> core::fmt::Debug for PoisonError<T> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.write_str("PoisonError { .. }")
    }
}

impl<T> core::fmt::Display for PoisonError<T> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.write_str("poisoned lock: another task failed while holding the lock")
    }
}

const UNINIT: u8 = 0;
const INITIALIZING: u8 = 1;
const INITIALIZED: u8 = 2;

/// A synchronization primitive for one-time initialization.
///
/// Uses a three-state atomic with `compare_exchange` to ensure exactly
/// one thread performs initialization.
pub struct OnceLock<T> {
    inner: UnsafeCell<Option<T>>,
    state: AtomicU8,
}

// SAFETY: OnceLock ensures exclusive access during initialization via
// compare_exchange, and only provides shared (&T) access after init.
unsafe impl<T: Send> Sync for OnceLock<T> {}
unsafe impl<T: Send> Send for OnceLock<T> {}

impl<T> OnceLock<T> {
    /// Creates a new empty `OnceLock`.
    pub const fn new() -> Self {
        OnceLock { inner: UnsafeCell::new(None), state: AtomicU8::new(UNINIT) }
    }

    /// Returns `Some(&value)` if initialized, `None` otherwise.
    pub fn get(&self) -> Option<&T> {
        if self.state.load(Ordering::Acquire) == INITIALIZED {
            // SAFETY: After INITIALIZED is visible, inner is immutable.
            unsafe { (*self.inner.get()).as_ref() }
        } else {
            None
        }
    }

    /// Initializes the cell with `value` if not yet initialized.
    pub fn set(&self, value: T) -> Result<(), T> {
        match self.state.compare_exchange(
            UNINIT,
            INITIALIZING,
            Ordering::Acquire,
            Ordering::Acquire,
        ) {
            Ok(_) => {
                // SAFETY: We hold the initialization lock.
                unsafe {
                    *self.inner.get() = Some(value);
                }
                self.state.store(INITIALIZED, Ordering::Release);
                Ok(())
            }
            Err(_) => Err(value),
        }
    }

    /// Initializes the cell with the result of `f` if not yet initialized,
    /// then returns a reference to the value.
    pub fn get_or_init<F>(&self, f: F) -> &T
    where
        F: FnOnce() -> T,
    {
        match self.state.compare_exchange(
            UNINIT,
            INITIALIZING,
            Ordering::Acquire,
            Ordering::Acquire,
        ) {
            Ok(_) => {
                let value = f();
                // SAFETY: We hold the initialization lock.
                unsafe {
                    *self.inner.get() = Some(value);
                }
                self.state.store(INITIALIZED, Ordering::Release);
            }
            Err(INITIALIZED) => {}
            Err(_) => {
                while self.state.load(Ordering::Acquire) != INITIALIZED {
                    core::hint::spin_loop();
                }
            }
        }
        // SAFETY: state is INITIALIZED, so inner is Some and immutable.
        unsafe { (*self.inner.get()).as_ref().unwrap() }
    }
}

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

/// A synchronization primitive for lazy initialization.
///
/// Wraps `OnceLock<T>` and stores the initialization function in a
/// `Mutex<Option<F>>`.
pub struct LazyLock<T, F = fn() -> T> {
    once: OnceLock<T>,
    init: Mutex<Option<F>>,
}

// SAFETY: OnceLock<T> is Sync when T: Send+Sync.
// Mutex<Option<F>> is Sync when Option<F>: Send, i.e. F: Send.
unsafe impl<T: Send + Sync, F: Send> Sync for LazyLock<T, F> {}
unsafe impl<T: Send, F: Send> Send for LazyLock<T, F> {}

impl<T, F: FnOnce() -> T> LazyLock<T, F> {
    /// Creates a new `LazyLock` with the given initialization function.
    pub const fn new(f: F) -> Self {
        LazyLock { once: OnceLock::new(), init: Mutex::new(Some(f)) }
    }

    /// Forces initialization and returns a reference to the value.
    pub fn force(&self) -> &T {
        self.once.get_or_init(|| {
            let f = self
                .init
                .lock()
                .expect("LazyLock init mutex poisoned")
                .take()
                .expect("LazyLock init function already taken");
            f()
        })
    }
}

impl<T, F: FnOnce() -> T> core::ops::Deref for LazyLock<T, F> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        self.force()
    }
}