Skip to content

Commit

Permalink
Implement OnceCell/Lock::try_insert()
Browse files Browse the repository at this point in the history
  • Loading branch information
daxpedda committed Oct 8, 2023
1 parent bf9a1c8 commit 966ec87
Show file tree
Hide file tree
Showing 2 changed files with 71 additions and 8 deletions.
38 changes: 33 additions & 5 deletions library/core/src/cell/once.rs
Original file line number Diff line number Diff line change
Expand Up @@ -87,10 +87,38 @@ 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
///
/// ```
/// 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]
#[stable(feature = "once_cell", since = "CURRENT_RUSTC_VERSION")]
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
Expand All @@ -99,7 +127,7 @@ impl<T> OnceCell<T> {
// maintains the `inner`'s invariant.
let slot = unsafe { &mut *self.inner.get() };
*slot = Some(value);
Ok(())
Ok(self.get().unwrap())
}

/// Gets the contents of the cell, initializing it with `f`
Expand Down
41 changes: 38 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,46 @@ 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
///
/// ```
/// 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]
#[stable(feature = "once_cell", since = "CURRENT_RUSTC_VERSION")]
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

0 comments on commit 966ec87

Please sign in to comment.