diff --git a/library/core/src/cell/once.rs b/library/core/src/cell/once.rs index 14c587e0c482e..6d2110ec0b273 100644 --- a/library/core/src/cell/once.rs +++ b/library/core/src/cell/once.rs @@ -1,4 +1,5 @@ use crate::cell::UnsafeCell; +use crate::ops::{ControlFlow, FromResidual, Residual, Try}; use crate::{fmt, mem}; /// A cell which can nominally be written to only once. @@ -228,14 +229,18 @@ impl OnceCell { /// assert_eq!(cell.get(), Some(&92)) /// ``` #[unstable(feature = "once_cell_try", issue = "109737")] - pub fn get_or_try_init(&self, f: F) -> Result<&T, E> + pub fn get_or_try_init<'a, F, R>(&'a self, f: F) -> >::TryType where - F: FnOnce() -> Result, + F: FnOnce() -> R, + R: Try>, { - if let Some(val) = self.get() { - return Ok(val); + if self.get().is_none() { + if let ControlFlow::Break(residual) = self.try_init(f) { + return FromResidual::from_residual(residual); + } } - self.try_init(f) + + try { self.get().unwrap() } } /// Gets the mutable reference of the contents of the cell, initializing @@ -266,29 +271,40 @@ impl OnceCell { /// assert_eq!(cell.get(), Some(&1236)) /// ``` #[unstable(feature = "once_cell_get_mut", issue = "121641")] - pub fn get_mut_or_try_init(&mut self, f: F) -> Result<&mut T, E> + pub fn get_mut_or_try_init<'a, F, R>( + &'a mut self, + f: F, + ) -> >::TryType where - F: FnOnce() -> Result, + F: FnOnce() -> R, + R: Try>, { if self.get().is_none() { - self.try_init(f)?; + if let ControlFlow::Break(residual) = self.try_init(f) { + return FromResidual::from_residual(residual); + } } - Ok(self.get_mut().unwrap()) + + try { self.get_mut().unwrap() } } // Avoid inlining the initialization closure into the common path that fetches // the already initialized value #[cold] - fn try_init(&self, f: F) -> Result<&T, E> + fn try_init(&self, f: F) -> ControlFlow where - F: FnOnce() -> Result, + F: FnOnce() -> R, + R: Try, { - let val = f()?; + let val = f().branch()?; // Note that *some* forms of reentrant initialization might lead to // UB (see `reentrant_init` test). I believe that just removing this // `panic`, while keeping `try_insert` would be sound, but it seems // better to panic, rather than to silently use an old value. - if let Ok(val) = self.try_insert(val) { Ok(val) } else { panic!("reentrant init") } + match self.try_insert(val) { + Ok(_) => ControlFlow::Continue(()), + Err(_) => panic!("reentrant init"), + } } /// Consumes the cell, returning the wrapped value. diff --git a/library/std/src/lib.rs b/library/std/src/lib.rs index b81e7c18abbb0..ef118defd0b50 100644 --- a/library/std/src/lib.rs +++ b/library/std/src/lib.rs @@ -361,6 +361,8 @@ #![feature(str_internals)] #![feature(strict_provenance)] #![feature(strict_provenance_atomic_ptr)] +#![feature(try_trait_v2)] +#![feature(try_trait_v2_residual)] #![feature(ub_checks)] // tidy-alphabetical-end // diff --git a/library/std/src/sync/once_lock.rs b/library/std/src/sync/once_lock.rs index be615a5a8ef37..3bbf9a4247638 100644 --- a/library/std/src/sync/once_lock.rs +++ b/library/std/src/sync/once_lock.rs @@ -2,6 +2,7 @@ use crate::cell::UnsafeCell; use crate::fmt; use crate::marker::PhantomData; use crate::mem::MaybeUninit; +use crate::ops::{ControlFlow, FromResidual, Residual, Try}; use crate::panic::{RefUnwindSafe, UnwindSafe}; use crate::sync::Once; @@ -375,24 +376,23 @@ impl OnceLock { /// ``` #[inline] #[unstable(feature = "once_cell_try", issue = "109737")] - pub fn get_or_try_init(&self, f: F) -> Result<&T, E> + pub fn get_or_try_init<'a, F, R>(&'a self, f: F) -> >::TryType where - F: FnOnce() -> Result, + F: FnOnce() -> R, + R: Try>, { // Fast path check // NOTE: We need to perform an acquire on the state in this method // in order to correctly synchronize `LazyLock::force`. This is - // currently done by calling `self.get()`, which in turn calls - // `self.is_initialized()`, which in turn performs the acquire. - if let Some(value) = self.get() { - return Ok(value); + // currently done by calling `self.is_initialized()`. + if !self.is_initialized() { + if let ControlFlow::Break(residual) = self.initialize(f) { + return FromResidual::from_residual(residual); + } } - self.initialize(f)?; - - debug_assert!(self.is_initialized()); // SAFETY: The inner value has been initialized - Ok(unsafe { self.get_unchecked() }) + try { unsafe { self.get_unchecked() } } } /// Gets the mutable reference of the contents of the cell, initializing @@ -426,16 +426,22 @@ impl OnceLock { /// ``` #[inline] #[unstable(feature = "once_cell_get_mut", issue = "121641")] - pub fn get_mut_or_try_init(&mut self, f: F) -> Result<&mut T, E> + pub fn get_mut_or_try_init<'a, F, R>( + &'a mut self, + f: F, + ) -> >::TryType where - F: FnOnce() -> Result, + F: FnOnce() -> R, + R: Try>, { - if self.get().is_none() { - self.initialize(f)?; + if !self.is_initialized() { + if let ControlFlow::Break(residual) = self.initialize(f) { + return FromResidual::from_residual(residual); + } } - debug_assert!(self.is_initialized()); + // SAFETY: The inner value has been initialized - Ok(unsafe { self.get_unchecked_mut() }) + try { unsafe { self.get_unchecked_mut() } } } /// Consumes the `OnceLock`, returning the wrapped value. Returns @@ -499,22 +505,22 @@ impl OnceLock { #[cold] #[optimize(size)] - fn initialize(&self, f: F) -> Result<(), E> + fn initialize(&self, f: F) -> ControlFlow where - F: FnOnce() -> Result, + F: FnOnce() -> R, + R: Try, { - let mut res: Result<(), E> = Ok(()); - let slot = &self.value; + let mut res = ControlFlow::Continue(()); // Ignore poisoning from other threads // If another thread panics, then we'll be able to run our closure self.once.call_once_force(|p| { - match f() { - Ok(value) => { - unsafe { (&mut *slot.get()).write(value) }; + match f().branch() { + ControlFlow::Continue(value) => { + unsafe { (&mut *self.value.get()).write(value) }; } - Err(e) => { - res = Err(e); + ControlFlow::Break(residual) => { + res = ControlFlow::Break(residual); // Treat the underlying `Once` as poisoned since we // failed to initialize our value. diff --git a/library/std/src/sync/once_lock/tests.rs b/library/std/src/sync/once_lock/tests.rs index 5113d436c3c99..e43ba8b1591f8 100644 --- a/library/std/src/sync/once_lock/tests.rs +++ b/library/std/src/sync/once_lock/tests.rs @@ -83,7 +83,7 @@ fn clone() { #[test] #[cfg_attr(not(panic = "unwind"), ignore = "test requires unwinding support")] fn get_or_try_init() { - let cell: OnceLock = OnceLock::new(); + let mut cell: OnceLock = OnceLock::new(); assert!(cell.get().is_none()); let res = panic::catch_unwind(|| cell.get_or_try_init(|| -> Result<_, ()> { panic!() })); @@ -91,10 +91,14 @@ fn get_or_try_init() { assert!(!cell.is_initialized()); assert!(cell.get().is_none()); + assert_eq!(cell.get_or_try_init(|| None), None); assert_eq!(cell.get_or_try_init(|| Err(())), Err(())); assert_eq!(cell.get_or_try_init(|| Ok::<_, ()>("hello".to_string())), Ok(&"hello".to_string())); - assert_eq!(cell.get(), Some(&"hello".to_string())); + assert_eq!(cell.take(), Some("hello".to_string())); + + assert_eq!(cell.get_or_try_init(|| Some("42".to_string())), Some(&"42".to_string())); + assert_eq!(cell.get(), Some(&"42".to_string())); } #[test]