Add get_mut_or_init and get_mut_or_try_init for OnceCell

[tisonkun]

Proposal

Problem statement

OnceCell has been a stable API since 1.70.0 with the following methods:

• pub fn get(&self) -> Option<&T> - stale since 1.70.0;
• pub fn get_mut(&mut self) -> Option<&mut T> - stable since 1.70.0;
• pub fn set(&self, value: T) -> Result<(), T> - stable since 1.70.0;
• pub fn get_or_init<F>(&self, f: F) -> &T - stable since 1.70.0;
• pub fn get_or_try_init<F, E>(&self, f: F) -> Result<&T, E> - unstable https://github.com/rust-lang/rust/issues/109737

The problem I encountered at https://github.com/rust-lang/rust/issues/74465#issuecomment-1668766041 is that I need to obtain a mut of OnceCell like get_mut reference and am unsure if it's initialized. So, a get_mut_or_init or get_mut_or_try_init counterpart to the ones without mut is needed.

Motivating examples or use cases

This proposal is straightforward so the problem is the motivation I want to add new APIs.

Current workaround:

    pub fn mut_batches(&mut self) -> IterMut<'_, RecordBatch> {
        self.batches.get_or_init(|| load_batches(&self.buf));
        // SAFETY - init above
        unsafe { self.batches.get_mut().unwrap_unchecked() }.iter_mut()
    }

Expected code with get_mut_or_init:

    pub fn mut_batches(&mut self) -> IterMut<'_, RecordBatch> {
        self.batches.get_mut_or_init(|| load_batches(&self.buf));
    }

Solution sketch

Add these two methods. It should be intuitive and I try it in https://github.com/rust-lang/rust/pull/114788. Since it's an API change, I was guided to create this ACP.

It looks like:

    #[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,
        }
    }

    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);
        }
        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(())
    }

Alternatives

Not yet. The solution proposed above should be the simplest one.

Links and related work

• https://github.com/rust-lang/rust/issues/74465
• https://github.com/rust-lang/rust/issues/109737
• https://github.com/rust-lang/rust/pull/114788

What happens now?

This issue contains an API change proposal (or ACP) and is part of the libs-api team feature lifecycle. Once this issue is filed, the libs-api team will review open proposals as capability becomes available. Current response times do not have a clear estimate, but may be up to several months.

Possible responses

The libs team may respond in various different ways. First, the team will consider the problem (this doesn't require any concrete solution or alternatives to have been proposed):

• We think this problem seems worth solving, and the standard library might be the right place to solve it.
• We think that this probably doesn't belong in the standard library.

Second, if there's a concrete solution:

• We think this specific solution looks roughly right, approved, you or someone else should implement this. (Further review will still happen on the subsequent implementation PR.)
• We're not sure this is the right solution, and the alternatives or other materials don't give us enough information to be sure about that. Here are some questions we have that aren't answered, or rough ideas about alternatives we'd want to see discussed.

[tgross35]

This proposal is straightforward so the problem is the motivation I want to add new APIs.

even if straightforward, we need an example of “this is what the workaround looks like now” and “this is how the same thing would look with this proposal”. This can just be a simplified version of your linked comment.

And we want to see code, not links, in the solution sketch.

[tisonkun]

@tgross35 Updated.

[tisonkun]

cc @tgross35 @dtolnay anything blocks here or how can I get some feedback to move forward?

[dtolnay]

Thank you! This looks great to me.

These methods are the equivalent of the lazycell crate's borrow_mut_with and try_borrow_mut_with, which are used by Cargo, blocking https://github.com/rust-lang/cargo/issues/9310. So I am eager for this to be supported better by OnceCell.