Skip to content

Commit

Permalink
Rollup merge of rust-lang#116540 - daxpedda:once-cell-lock-try-insert…
Browse files Browse the repository at this point in the history
…, 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: rust-lang#116693
  • Loading branch information
GuillaumeGomez authored Oct 14, 2023
2 parents ee8c9d3 + dd34d90 commit fcd75cc
Show file tree
Hide file tree
Showing 2 changed files with 77 additions and 12 deletions.
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

0 comments on commit fcd75cc

Please sign in to comment.