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

cubecl_common/
stub.rs
