Skip to content

Implement OnceCell/Lock::try_insert() #116540

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Oct 15, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 37 additions & 9 deletions library/core/src/cell/once.rs
Original file line number Diff line number Diff line change
Expand Up @@ -87,19 +87,48 @@ impl<T> OnceCell<T> {
#[inline]
#[stable(feature = "once_cell", since = "1.70.0")]
pub fn set(&self, value: T) -> Result<(), T> {
// SAFETY: Safe because we cannot have overlapping mutable borrows
let slot = unsafe { &*self.inner.get() };
if slot.is_some() {
return Err(value);
match self.try_insert(value) {
Ok(_) => Ok(()),
Err((_, value)) => Err(value),
}
}

/// Sets the contents of the cell to `value` if the cell was empty, then
/// returns a reference to it.
///
/// # Errors
///
/// This method returns `Ok(&value)` if the cell was empty and
/// `Err(&current_value, value)` if it was full.
///
/// # Examples
///
/// ```
/// #![feature(once_cell_try_insert)]
///
/// use std::cell::OnceCell;
///
/// let cell = OnceCell::new();
/// assert!(cell.get().is_none());
///
/// assert_eq!(cell.try_insert(92), Ok(&92));
/// assert_eq!(cell.try_insert(62), Err((&92, 62)));
///
/// assert!(cell.get().is_some());
/// ```
#[inline]
#[unstable(feature = "once_cell_try_insert", issue = "116693")]
pub fn try_insert(&self, value: T) -> Result<&T, (&T, T)> {
if let Some(old) = self.get() {
return Err((old, value));
}

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

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

/// Consumes the cell, returning the wrapped value.
Expand Down
43 changes: 40 additions & 3 deletions library/std/src/sync/once_lock.rs
Original file line number Diff line number Diff line change
Expand Up @@ -126,11 +126,48 @@ impl<T> OnceLock<T> {
#[inline]
#[stable(feature = "once_cell", since = "1.70.0")]
pub fn set(&self, value: T) -> Result<(), T> {
match self.try_insert(value) {
Ok(_) => Ok(()),
Err((_, value)) => Err(value),
}
}

/// Sets the contents of this cell to `value` if the cell was empty, then
/// returns a reference to it.
///
/// May block if another thread is currently attempting to initialize the cell. The cell is
/// guaranteed to contain a value when set returns, though not necessarily the one provided.
///
/// Returns `Ok(&value)` if the cell was empty and `Err(&current_value, value)` if it was full.
///
/// # Examples
///
/// ```
/// #![feature(once_cell_try_insert)]
///
/// use std::sync::OnceLock;
///
/// static CELL: OnceLock<i32> = OnceLock::new();
///
/// fn main() {
/// assert!(CELL.get().is_none());
///
/// std::thread::spawn(|| {
/// assert_eq!(CELL.try_insert(92), Ok(&92));
/// }).join().unwrap();
///
/// assert_eq!(CELL.try_insert(62), Err((&92, 62)));
/// assert_eq!(CELL.get(), Some(&92));
/// }
/// ```
#[inline]
#[unstable(feature = "once_cell_try_insert", issue = "116693")]
pub fn try_insert(&self, value: T) -> Result<&T, (&T, T)> {
let mut value = Some(value);
self.get_or_init(|| value.take().unwrap());
let res = self.get_or_init(|| value.take().unwrap());
match value {
None => Ok(()),
Some(value) => Err(value),
None => Ok(res),
Some(value) => Err((res, value)),
}
}

Expand Down