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

cubecl_common/
stub.rs
