option_lock/
once.rs

1use core::{
2    cell::Cell,
3    fmt::{self, Debug, Display, Formatter},
4    hint::spin_loop,
5    ops::Deref,
6};
7
8use super::{error::OptionLockError, lock::OptionLock};
9
10/// An `Option` value which can be safely written once.
11#[repr(transparent)]
12pub struct OnceCell<T>(OptionLock<T>);
13
14impl<T> OnceCell<T> {
15    /// Create a new, empty `OnceCell`.
16    pub const fn empty() -> Self {
17        Self(OptionLock::empty())
18    }
19
20    /// Create a `OnceCell` from an owned value.
21    pub const fn new(value: T) -> Self {
22        Self(OptionLock::new(value))
23    }
24
25    /// Get a shared reference to the contained value, if any.
26    pub fn get(&self) -> Option<&T> {
27        if self.0.is_some() {
28            // safe because the value is never reassigned
29            Some(unsafe { &*self.0.as_ptr() })
30        } else {
31            None
32        }
33    }
34
35    /// Get a mutable reference to the contained value, if any.
36    pub fn get_mut(&mut self) -> Option<&mut T> {
37        self.0.get_mut()
38    }
39
40    /// Get a reference to the contained value, initializing it if necessary.
41    /// The initializer will only be run by one thread if multiple are in competition.
42    pub fn get_or_init(&self, init: impl FnOnce() -> T) -> &T {
43        if let Some(value) = self.get() {
44            return value;
45        }
46        match self.0.try_lock_none() {
47            Ok(mut guard) => {
48                let prev = guard.replace(init());
49                assert!(prev.is_none());
50            }
51            Err(OptionLockError::FillState) => {
52                // filled
53            }
54            Err(OptionLockError::Unavailable) => loop {
55                while !self.0.is_some_unlocked() {
56                    spin_loop();
57                }
58            },
59        }
60        unsafe { &*self.0.as_ptr() }
61    }
62
63    /// Get a reference to the contained value, initializing it if necessary.
64    /// The initializer will only be run by one thread if multiple are in competition.
65    pub fn get_or_try_init<E>(&self, init: impl FnOnce() -> Result<T, E>) -> Result<&T, E> {
66        if let Some(value) = self.get() {
67            return Ok(value);
68        }
69        match self.0.try_lock_none() {
70            Ok(mut guard) => {
71                let prev = guard.replace(init()?);
72                assert!(prev.is_none());
73            }
74            Err(OptionLockError::FillState) => {
75                // filled
76            }
77            Err(OptionLockError::Unavailable) => loop {
78                while !self.0.is_some_unlocked() {
79                    spin_loop();
80                }
81            },
82        }
83        Ok(unsafe { &*self.0.as_ptr() })
84    }
85
86    /// Assign the value of the OnceCell, returning `Some(value)` if
87    /// the cell is already locked or populated.
88    pub fn set(&self, value: T) -> Result<(), T> {
89        self.0.try_fill(value)
90    }
91
92    /// Extract the inner value.
93    pub fn into_inner(self) -> Option<T> {
94        self.0.into_inner()
95    }
96
97    /// Check if the lock is currently acquired.
98    pub fn is_locked(&self) -> bool {
99        self.0.is_locked()
100    }
101}
102
103impl<T: Clone> Clone for OnceCell<T> {
104    fn clone(&self) -> Self {
105        Self::from(self.get().cloned())
106    }
107}
108
109impl<T> Default for OnceCell<T> {
110    fn default() -> Self {
111        Self(None.into())
112    }
113}
114
115impl<T: Debug> Debug for OnceCell<T> {
116    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
117        f.debug_tuple("OnceCell").field(&self.get()).finish()
118    }
119}
120
121impl<T: Display> Display for OnceCell<T> {
122    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
123        if let Some(val) = self.get() {
124            Display::fmt(val, f)
125        } else {
126            write!(f, "None")
127        }
128    }
129}
130
131impl<T> From<T> for OnceCell<T> {
132    fn from(data: T) -> Self {
133        Self(data.into())
134    }
135}
136
137impl<T> From<Option<T>> for OnceCell<T> {
138    fn from(data: Option<T>) -> Self {
139        Self(data.into())
140    }
141}
142
143impl<T> From<OptionLock<T>> for OnceCell<T> {
144    fn from(lock: OptionLock<T>) -> Self {
145        Self(lock)
146    }
147}
148
149/// A convenient wrapper around a `OnceCell<T>` with an initializer.
150pub struct Lazy<T, F = fn() -> T> {
151    cell: OnceCell<T>,
152    init: Cell<Option<F>>,
153}
154
155unsafe impl<T, F: Send> Sync for Lazy<T, F> where OnceCell<T>: Sync {}
156
157impl<T, F> Lazy<T, F> {
158    /// Create a new Lazy instance
159    pub const fn new(init: F) -> Self {
160        Self {
161            cell: OnceCell::empty(),
162            init: Cell::new(Some(init)),
163        }
164    }
165}
166
167impl<T: Debug, F> Debug for Lazy<T, F> {
168    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
169        f.debug_struct("Lazy")
170            .field("cell", &self.cell)
171            .field("init", &"..")
172            .finish()
173    }
174}
175
176impl<T: Display, F: FnOnce() -> T> Display for Lazy<T, F> {
177    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
178        Display::fmt(&**self, f)
179    }
180}
181
182impl<T, F: FnOnce() -> T> Lazy<T, F> {
183    /// Ensure that the initializer has run
184    pub fn force(this: &Self) -> &T {
185        this.cell.get_or_init(|| (this.init.take().unwrap())())
186    }
187}
188
189impl<T: Default> Default for Lazy<T> {
190    fn default() -> Lazy<T> {
191        Lazy::new(T::default)
192    }
193}
194
195impl<T, F: FnOnce() -> T> Deref for Lazy<T, F> {
196    type Target = T;
197
198    fn deref(&self) -> &T {
199        Lazy::force(self)
200    }
201}