Skip to content

Commit fcd75cc

Browse files
Rollup merge of #116540 - daxpedda:once-cell-lock-try-insert, r=Mark-Simulacrum
Implement `OnceCell/Lock::try_insert()` I took inspiration from [`once_cell`](https://crates.io/crates/once_cell): - [`once_cell::unsync::OnceCell::try_insert()`](https://github.com/matklad/once_cell/blob/874f9373abd7feaf923a3b3c34bfb3383529c671/src/lib.rs#L551-L563) - [`once_cell::sync::OnceCell::try_insert()`](https://github.com/matklad/once_cell/blob/874f9373abd7feaf923a3b3c34bfb3383529c671/src/lib.rs#L1080-L1087) I tried to change as little code as possible in the first commit and applied some obvious optimizations in the second one. ACP: rust-lang/libs-team#276 Tracking issue: #116693
2 parents ee8c9d3 + dd34d90 commit fcd75cc

File tree

2 files changed

+77
-12
lines changed

2 files changed

+77
-12
lines changed

library/core/src/cell/once.rs

+37-9
Original file line numberDiff line numberDiff line change
@@ -87,19 +87,48 @@ impl<T> OnceCell<T> {
8787
#[inline]
8888
#[stable(feature = "once_cell", since = "1.70.0")]
8989
pub fn set(&self, value: T) -> Result<(), T> {
90-
// SAFETY: Safe because we cannot have overlapping mutable borrows
91-
let slot = unsafe { &*self.inner.get() };
92-
if slot.is_some() {
93-
return Err(value);
90+
match self.try_insert(value) {
91+
Ok(_) => Ok(()),
92+
Err((_, value)) => Err(value),
93+
}
94+
}
95+
96+
/// Sets the contents of the cell to `value` if the cell was empty, then
97+
/// returns a reference to it.
98+
///
99+
/// # Errors
100+
///
101+
/// This method returns `Ok(&value)` if the cell was empty and
102+
/// `Err(&current_value, value)` if it was full.
103+
///
104+
/// # Examples
105+
///
106+
/// ```
107+
/// #![feature(once_cell_try_insert)]
108+
///
109+
/// use std::cell::OnceCell;
110+
///
111+
/// let cell = OnceCell::new();
112+
/// assert!(cell.get().is_none());
113+
///
114+
/// assert_eq!(cell.try_insert(92), Ok(&92));
115+
/// assert_eq!(cell.try_insert(62), Err((&92, 62)));
116+
///
117+
/// assert!(cell.get().is_some());
118+
/// ```
119+
#[inline]
120+
#[unstable(feature = "once_cell_try_insert", issue = "116693")]
121+
pub fn try_insert(&self, value: T) -> Result<&T, (&T, T)> {
122+
if let Some(old) = self.get() {
123+
return Err((old, value));
94124
}
95125

96126
// SAFETY: This is the only place where we set the slot, no races
97127
// due to reentrancy/concurrency are possible, and we've
98128
// checked that slot is currently `None`, so this write
99129
// maintains the `inner`'s invariant.
100130
let slot = unsafe { &mut *self.inner.get() };
101-
*slot = Some(value);
102-
Ok(())
131+
Ok(slot.insert(value))
103132
}
104133

105134
/// Gets the contents of the cell, initializing it with `f`
@@ -183,10 +212,9 @@ impl<T> OnceCell<T> {
183212
let val = outlined_call(f)?;
184213
// Note that *some* forms of reentrant initialization might lead to
185214
// UB (see `reentrant_init` test). I believe that just removing this
186-
// `assert`, while keeping `set/get` would be sound, but it seems
215+
// `panic`, while keeping `try_insert` would be sound, but it seems
187216
// better to panic, rather than to silently use an old value.
188-
assert!(self.set(val).is_ok(), "reentrant init");
189-
Ok(self.get().unwrap())
217+
if let Ok(val) = self.try_insert(val) { Ok(val) } else { panic!("reentrant init") }
190218
}
191219

192220
/// Consumes the cell, returning the wrapped value.

library/std/src/sync/once_lock.rs

+40-3
Original file line numberDiff line numberDiff line change
@@ -126,11 +126,48 @@ impl<T> OnceLock<T> {
126126
#[inline]
127127
#[stable(feature = "once_cell", since = "1.70.0")]
128128
pub fn set(&self, value: T) -> Result<(), T> {
129+
match self.try_insert(value) {
130+
Ok(_) => Ok(()),
131+
Err((_, value)) => Err(value),
132+
}
133+
}
134+
135+
/// Sets the contents of this cell to `value` if the cell was empty, then
136+
/// returns a reference to it.
137+
///
138+
/// May block if another thread is currently attempting to initialize the cell. The cell is
139+
/// guaranteed to contain a value when set returns, though not necessarily the one provided.
140+
///
141+
/// Returns `Ok(&value)` if the cell was empty and `Err(&current_value, value)` if it was full.
142+
///
143+
/// # Examples
144+
///
145+
/// ```
146+
/// #![feature(once_cell_try_insert)]
147+
///
148+
/// use std::sync::OnceLock;
149+
///
150+
/// static CELL: OnceLock<i32> = OnceLock::new();
151+
///
152+
/// fn main() {
153+
/// assert!(CELL.get().is_none());
154+
///
155+
/// std::thread::spawn(|| {
156+
/// assert_eq!(CELL.try_insert(92), Ok(&92));
157+
/// }).join().unwrap();
158+
///
159+
/// assert_eq!(CELL.try_insert(62), Err((&92, 62)));
160+
/// assert_eq!(CELL.get(), Some(&92));
161+
/// }
162+
/// ```
163+
#[inline]
164+
#[unstable(feature = "once_cell_try_insert", issue = "116693")]
165+
pub fn try_insert(&self, value: T) -> Result<&T, (&T, T)> {
129166
let mut value = Some(value);
130-
self.get_or_init(|| value.take().unwrap());
167+
let res = self.get_or_init(|| value.take().unwrap());
131168
match value {
132-
None => Ok(()),
133-
Some(value) => Err(value),
169+
None => Ok(res),
170+
Some(value) => Err((res, value)),
134171
}
135172
}
136173

0 commit comments

Comments
 (0)