impl get_mut_or_init and get_mut_or_try_init for OnceCell

See also rust-lang#74465 (comment)

Signed-off-by: tison <wander4096@gmail.com>

Author: tisonkun

library/core/src/cell/once.rs

@@ -135,6 +135,16 @@ impl<T> OnceCell<T> {
         }
     }
 
+    #[inline]
+    pub fn get_mut_or_init<F>(&mut self, f: F) -> &mut T
+    where
+        F: FnOnce() -> T,
+    {
+        match self.get_mut_or_try_init(|| Ok::<T, !>(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.
@@ -171,22 +181,35 @@ impl<T> OnceCell<T> {
         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, T, E>(f: F) -> Result<T, E>
-        where
-            F: FnOnce() -> Result<T, E>,
-        {
-            f()
+        self.try_init(f)?;
+        Ok(self.get().unwrap())
+    }
+
+    pub fn get_mut_or_try_init<F, E>(&mut self, f: F) -> Result<&mut T, E>
+    where
+        F: FnOnce() -> Result<T, E>,
+    {
+        if let Some(val) = self.get() {
+            return Ok(val);
         }
-        let val = outlined_call(f)?;
+        self.try_init(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<F, E>(&self, f: F) -> Result<(), E>
+    where
+        F: FnOnce() -> Result<T, E>,
+    {
+        let val = 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
         // better to panic, rather than to silently use an old value.
         assert!(self.set(val).is_ok(), "reentrant init");
-        Ok(self.get().unwrap())
+        Ok(())
     }
 
     /// Consumes the cell, returning the wrapped value.