cubecl_common/
stub.rs

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

cubecl_common/
stub.rs
