cubecl_common/
stub.rs

1#[cfg(not(feature = "std"))]
2use spin::{
3    Mutex as MutexImported, MutexGuard, Once as OnceImported, RwLock as RwLockImported,
4    RwLockReadGuard, RwLockWriteGuard,
5};
6#[cfg(feature = "std")]
7use std::sync::{
8    Mutex as MutexImported, MutexGuard, OnceLock as OnceImported, RwLock as RwLockImported,
9    RwLockReadGuard, RwLockWriteGuard,
10};
11
12/// A mutual exclusion primitive useful for protecting shared data
13///
14/// This mutex will block threads waiting for the lock to become available. The
15/// mutex can also be statically initialized or created via a [Mutex::new]
16///
17/// [Mutex] wrapper to make `spin::Mutex` API compatible with `std::sync::Mutex` to swap
18#[derive(Debug)]
19pub struct Mutex<T> {
20    inner: MutexImported<T>,
21}
22
23impl<T> Mutex<T> {
24    /// Creates a new mutex in an unlocked state ready for use.
25    #[inline(always)]
26    pub const fn new(value: T) -> Self {
27        Self {
28            inner: MutexImported::new(value),
29        }
30    }
31
32    /// Locks the mutex blocking the current thread until it is able to do so.
33    #[inline(always)]
34    pub fn lock(&self) -> Result<MutexGuard<T>, alloc::string::String> {
35        #[cfg(not(feature = "std"))]
36        {
37            Ok(self.inner.lock())
38        }
39
40        #[cfg(feature = "std")]
41        {
42            self.inner.lock().map_err(|err| err.to_string())
43        }
44    }
45}
46
47/// A reader-writer lock which is exclusively locked for writing or shared for reading.
48/// This reader-writer lock will block threads waiting for the lock to become available.
49/// The lock can also be statically initialized or created via a [RwLock::new]
50/// [RwLock] wrapper to make `spin::RwLock` API compatible with `std::sync::RwLock` to swap
51#[derive(Debug)]
52pub struct RwLock<T> {
53    inner: RwLockImported<T>,
54}
55
56impl<T> RwLock<T> {
57    /// Creates a new reader-writer lock in an unlocked state ready for use.
58    #[inline(always)]
59    pub const fn new(value: T) -> Self {
60        Self {
61            inner: RwLockImported::new(value),
62        }
63    }
64
65    /// Locks this rwlock with shared read access, blocking the current thread
66    /// until it can be acquired.
67    #[inline(always)]
68    pub fn read(&self) -> Result<RwLockReadGuard<T>, alloc::string::String> {
69        #[cfg(not(feature = "std"))]
70        {
71            Ok(self.inner.read())
72        }
73        #[cfg(feature = "std")]
74        {
75            self.inner.read().map_err(|err| err.to_string())
76        }
77    }
78
79    /// Locks this rwlock with exclusive write access, blocking the current thread
80    /// until it can be acquired.
81    #[inline(always)]
82    pub fn write(&self) -> Result<RwLockWriteGuard<T>, alloc::string::String> {
83        #[cfg(not(feature = "std"))]
84        {
85            Ok(self.inner.write())
86        }
87
88        #[cfg(feature = "std")]
89        {
90            self.inner.write().map_err(|err| err.to_string())
91        }
92    }
93}
94
95/// A unique identifier for a running thread.
96///
97/// This module is a stub when no std is available to swap with std::thread::ThreadId.
98#[allow(dead_code)]
99#[derive(Eq, PartialEq, Clone, Copy, Hash, Debug)]
100pub struct ThreadId(core::num::NonZeroU64);
101
102/// A cell that provides lazy one-time initialization that implements [Sync] and [Send].
103///
104/// This module is a stub when no std is available to swap with [std::sync::OnceLock].
105pub struct SyncOnceCell<T>(OnceImported<T>);
106
107impl<T: core::fmt::Debug> Default for SyncOnceCell<T> {
108    fn default() -> Self {
109        Self::new()
110    }
111}
112
113impl<T: core::fmt::Debug> SyncOnceCell<T> {
114    /// Create a new once.
115    #[inline(always)]
116    pub fn new() -> Self {
117        Self(OnceImported::new())
118    }
119
120    /// Initialize the cell with a value.
121    #[inline(always)]
122    pub fn initialized(value: T) -> Self {
123        #[cfg(not(feature = "std"))]
124        {
125            let cell = OnceImported::initialized(value);
126            Self(cell)
127        }
128
129        #[cfg(feature = "std")]
130        {
131            let cell = OnceImported::new();
132            cell.set(value).unwrap();
133
134            Self(cell)
135        }
136    }
137
138    /// Gets the contents of the cell, initializing it with `f` if the cell
139    /// was empty.
140    #[inline(always)]
141    pub fn get_or_init<F>(&self, f: F) -> &T
142    where
143        F: FnOnce() -> T,
144    {
145        #[cfg(not(feature = "std"))]
146        {
147            self.0.call_once(f)
148        }
149
150        #[cfg(feature = "std")]
151        {
152            self.0.get_or_init(f)
153        }
154    }
155}
cubecl_common/stub.rs

cubecl_common/
stub.rs
