From 95e195f41ea6cc64846b1fec4fad9bf6d8f7db82 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 14 Feb 2024 19:51:02 +0800 Subject: [PATCH] impl get_mut_or_init and get_mut_or_try_init for OnceCell and OnceLock See also https://github.com/rust-lang/rust/issues/74465#issuecomment-1676522051 Signed-off-by: tison --- library/core/src/cell/once.rs | 93 ++++++++++++++++++++++++++++--- library/std/src/sync/once_lock.rs | 81 +++++++++++++++++++++++++++ 2 files changed, 165 insertions(+), 9 deletions(-) diff --git a/library/core/src/cell/once.rs b/library/core/src/cell/once.rs index 3877a0c48cb23..a7c3dfc982d12 100644 --- a/library/core/src/cell/once.rs +++ b/library/core/src/cell/once.rs @@ -164,6 +164,42 @@ impl OnceCell { } } + /// Gets the mutable reference of the contents of the cell, + /// initializing it with `f` if the cell was empty. + /// + /// # Panics + /// + /// If `f` panics, the panic is propagated to the caller, and the cell + /// remains uninitialized. + /// + /// # Examples + /// + /// ``` + /// #![feature(once_cell_get_mut)] + /// + /// use std::cell::OnceCell; + /// + /// let mut cell = OnceCell::new(); + /// let value = cell.get_mut_or_init(|| 92); + /// assert_eq!(*value, 92); + /// + /// *value += 2; + /// assert_eq!(*value, 94); + /// + /// let value = cell.get_mut_or_init(|| unreachable!()); + /// assert_eq!(*value, 94); + /// ``` + #[inline] + #[unstable(feature = "once_cell_get_mut", issue = "121641")] + pub fn get_mut_or_init(&mut self, f: F) -> &mut T + where + F: FnOnce() -> T, + { + match self.get_mut_or_try_init(|| Ok::(f())) { + Ok(val) => val, + } + } + /// Gets the contents of the cell, initializing it with `f` if /// the cell was empty. If the cell was empty and `f` failed, an /// error is returned. @@ -200,16 +236,55 @@ impl OnceCell { if let Some(val) = self.get() { return Ok(val); } - /// Avoid inlining the initialization closure into the common path that fetches - /// the already initialized value - #[cold] - fn outlined_call(f: F) -> Result - where - F: FnOnce() -> Result, - { - f() + self.try_init(f) + } + + /// Gets the mutable reference of the contents of the cell, initializing + /// it with `f` if the cell was empty. If the cell was empty and `f` failed, + /// an error is returned. + /// + /// # Panics + /// + /// If `f` panics, the panic is propagated to the caller, and the cell + /// remains uninitialized. + /// + /// # Examples + /// + /// ``` + /// #![feature(once_cell_get_mut)] + /// + /// use std::cell::OnceCell; + /// + /// let mut cell: OnceCell = OnceCell::new(); + /// + /// // Failed initializers do not change the value + /// assert!(cell.get_mut_or_try_init(|| "not a number!".parse()).is_err()); + /// assert!(cell.get().is_none()); + /// + /// let value = cell.get_mut_or_try_init(|| "1234".parse()); + /// assert_eq!(value, Ok(&mut 1234)); + /// *value.unwrap() += 2; + /// 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> + where + F: FnOnce() -> Result, + { + if self.get().is_none() { + self.try_init(f)?; } - let val = outlined_call(f)?; + Ok(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> + where + F: FnOnce() -> Result, + { + let val = f()?; // 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 diff --git a/library/std/src/sync/once_lock.rs b/library/std/src/sync/once_lock.rs index 6d068613f8f30..fc830baccedd2 100644 --- a/library/std/src/sync/once_lock.rs +++ b/library/std/src/sync/once_lock.rs @@ -252,6 +252,46 @@ impl OnceLock { } } + /// Gets the mutable reference of the contents of the cell, initializing + /// it with `f` if the cell was empty. + /// + /// Many threads may call `get_mut_or_init` concurrently with different + /// initializing functions, but it is guaranteed that only one function + /// will be executed. + /// + /// # Panics + /// + /// If `f` panics, the panic is propagated to the caller, and the cell + /// remains uninitialized. + /// + /// # Examples + /// + /// ``` + /// #![feature(once_cell_get_mut)] + /// + /// use std::sync::OnceLock; + /// + /// let mut cell = OnceLock::new(); + /// let value = cell.get_mut_or_init(|| 92); + /// assert_eq!(*value, 92); + /// + /// *value += 2; + /// assert_eq!(*value, 94); + /// + /// let value = cell.get_mut_or_init(|| unreachable!()); + /// assert_eq!(*value, 94); + /// ``` + #[inline] + #[unstable(feature = "once_cell_get_mut", issue = "121641")] + pub fn get_mut_or_init(&mut self, f: F) -> &mut T + where + F: FnOnce() -> T, + { + match self.get_mut_or_try_init(|| Ok::(f())) { + Ok(val) => val, + } + } + /// Gets the contents of the cell, initializing it with `f` if /// the cell was empty. If the cell was empty and `f` failed, an /// error is returned. @@ -303,6 +343,47 @@ impl OnceLock { Ok(unsafe { self.get_unchecked() }) } + /// Gets the mutable reference of the contents of the cell, initializing + /// it with `f` if the cell was empty. If the cell was empty and `f` failed, + /// an error is returned. + /// + /// # Panics + /// + /// If `f` panics, the panic is propagated to the caller, and + /// the cell remains uninitialized. + /// + /// # Examples + /// + /// ``` + /// #![feature(once_cell_get_mut)] + /// + /// use std::sync::OnceLock; + /// + /// let mut cell: OnceLock = OnceLock::new(); + /// + /// // Failed initializers do not change the value + /// assert!(cell.get_mut_or_try_init(|| "not a number!".parse()).is_err()); + /// assert!(cell.get().is_none()); + /// + /// let value = cell.get_mut_or_try_init(|| "1234".parse()); + /// assert_eq!(value, Ok(&mut 1234)); + /// *value.unwrap() += 2; + /// assert_eq!(cell.get(), Some(&1236)) + /// ``` + #[inline] + #[unstable(feature = "once_cell_get_mut", issue = "121641")] + pub fn get_mut_or_try_init(&mut self, f: F) -> Result<&mut T, E> + where + F: FnOnce() -> Result, + { + if self.get().is_none() { + self.initialize(f)?; + } + debug_assert!(self.is_initialized()); + // SAFETY: The inner value has been initialized + Ok(unsafe { self.get_unchecked_mut() }) + } + /// Consumes the `OnceLock`, returning the wrapped value. Returns /// `None` if the cell was empty. ///