Re: [PATCH v14 1/7] rust: sync: add `OnceLock`

[Benno Lossin]

On Thu Jul 3, 2025 at 11:03 AM CEST, Andreas Hindborg wrote:
> "Benno Lossin" <lossin@xxxxxxxxxx> writes:
>> On Wed Jul 2, 2025 at 3:18 PM CEST, Andreas Hindborg wrote:
>>> +///
>>> +/// # Example
>>> +///
>>> +/// ```
>>> +/// # use kernel::sync::once_lock::OnceLock;
>>> +/// let value = OnceLock::new();
>>> +/// assert_eq!(None, value.as_ref());
>>> +///
>>> +/// let status = value.populate(42u8);
>>> +/// assert_eq!(true, status);
>>> +/// assert_eq!(Some(&42u8), value.as_ref());
>>> +/// assert_eq!(Some(42u8), value.copy());
>>> +///
>>> +/// let status = value.populate(101u8);
>>> +/// assert_eq!(false, status);
>>> +/// assert_eq!(Some(&42u8), value.as_ref());
>>> +/// assert_eq!(Some(42u8), value.copy());
>>> +/// ```
>>> +pub struct OnceLock<T> {
>>> +    init: Atomic<u32>,
>>> +    value: Opaque<T>,
>>> +}
>>> +
>>> +impl<T> Default for OnceLock<T> {
>>> +    fn default() -> Self {
>>> +        Self::new()
>>> +    }
>>> +}
>>> +
>>> +impl<T> OnceLock<T> {
>>> +    /// Create a new [`OnceLock`].
>>> +    ///
>>> +    /// The returned instance will be empty.
>>> +    pub const fn new() -> Self {
>>> +        // INVARIANT: The container is empty and we set `init` to `0`.
>>> +        Self {
>>> +            value: Opaque::uninit(),
>>> +            init: Atomic::new(0),
>>> +        }
>>> +    }
>>> +
>>> +    /// Get a reference to the contained object.
>>> +    ///
>>> +    /// Returns [`None`] if this [`OnceLock`] is empty.
>>> +    pub fn as_ref(&self) -> Option<&T> {
>>> +        if self.init.load(Acquire) == 2 {
>>> +            // SAFETY: As determined by the load above, the object is ready for shared access.
>>
>>     // SAFETY: By the safety requirements of `Self`, `self.init == 2` means that `self.value` contains
>>     // a valid value.
>
> By the *type invariants* I guess?

Oh yeah.

>>> +            Some(unsafe { &*self.value.get() })
>>> +        } else {
>>> +            None
>>> +        }
>>> +    }
>>> +
>>> +    /// Populate the [`OnceLock`].
>>> +    ///
>>> +    /// Returns `true` if the [`OnceLock`] was successfully populated.
>>> +    pub fn populate(&self, value: T) -> bool {
>>> +        // INVARIANT: We obtain exclusive access to the contained allocation and write 1 to
>>> +        // `init`.
>>> +        if let Ok(0) = self.init.cmpxchg(0, 1, Acquire) {
>>> +            // SAFETY: We obtained exclusive access to the contained object.
>>> +            unsafe { core::ptr::write(self.value.get(), value) };
>>> +            // INVARIANT: We release our exclusive access and transition the object to shared
>>> +            // access.
>>> +            self.init.store(2, Release);
>>> +            true
>>> +        } else {
>>> +            false
>>> +        }
>>> +    }
>>> +}
>>> +
>>> +impl<T: Copy> OnceLock<T> {
>>> +    /// Get a copy of the contained object.
>>> +    ///
>>> +    /// Returns [`None`] if the [`OnceLock`] is empty.
>>> +    pub fn copy(&self) -> Option<T> {
>>> +        if self.init.load(Acquire) == 2 {
>>> +            // SAFETY: As determined by the load above, the object is ready for shared access.
>>> +            Some(unsafe { *self.value.get() })
>>> +        } else {
>>> +            None
>>> +        }
>>
>> The impl can just be:
>>
>>     self.as_ref().copied()
>
> Nice. I was thinking of dropping this method and just have callers do
>
>  my_once_lock.as_ref().map(|v| v.copied())
>
> What do you think?

There is `Option::copied`, so no need for the `.map` call. I don't
really have a preference, if users always want to access it by-value,
then we should have `copy`.

---
Cheers,
Benno