set_once.rs

  1// SPDX-License-Identifier: GPL-2.0
  2
  3//! A container that can be initialized at most once.
  4
  5use super::atomic::{
  6    ordering::{Acquire, Relaxed, Release},
  7    Atomic,
  8};
  9use core::{cell::UnsafeCell, mem::MaybeUninit};
 10
 11/// A container that can be populated at most once. Thread safe.
 12///
 13/// Once the a [`SetOnce`] is populated, it remains populated by the same object for the
 14/// lifetime `Self`.
 15///
 16/// # Invariants
 17///
 18/// - `init` may only increase in value.
 19/// - `init` may only assume values in the range `0..=2`.
 20/// - `init == 0` if and only if `value` is uninitialized.
 21/// - `init == 1` if and only if there is exactly one thread with exclusive
 22///   access to `self.value`.
 23/// - `init == 2` if and only if `value` is initialized and valid for shared
 24///   access.
 25///
 26/// # Example
 27///
 28/// ```
 29/// # use kernel::sync::SetOnce;
 30/// let value = SetOnce::new();
 31/// assert_eq!(None, value.as_ref());
 32///
 33/// let status = value.populate(42u8);
 34/// assert_eq!(true, status);
 35/// assert_eq!(Some(&42u8), value.as_ref());
 36/// assert_eq!(Some(42u8), value.copy());
 37///
 38/// let status = value.populate(101u8);
 39/// assert_eq!(false, status);
 40/// assert_eq!(Some(&42u8), value.as_ref());
 41/// assert_eq!(Some(42u8), value.copy());
 42/// ```
 43pub struct SetOnce<T> {
 44    init: Atomic<u32>,
 45    value: UnsafeCell<MaybeUninit<T>>,
 46}
 47
 48impl<T> Default for SetOnce<T> {
 49    fn default() -> Self {
 50        Self::new()
 51    }
 52}
 53
 54impl<T> SetOnce<T> {
 55    /// Create a new [`SetOnce`].
 56    ///
 57    /// The returned instance will be empty.
 58    pub const fn new() -> Self {
 59        // INVARIANT: The container is empty and we initialize `init` to `0`.
 60        Self {
 61            value: UnsafeCell::new(MaybeUninit::uninit()),
 62            init: Atomic::new(0),
 63        }
 64    }
 65
 66    /// Get a reference to the contained object.
 67    ///
 68    /// Returns [`None`] if this [`SetOnce`] is empty.
 69    pub fn as_ref(&self) -> Option<&T> {
 70        if self.init.load(Acquire) == 2 {
 71            // SAFETY: By the type invariants of `Self`, `self.init == 2` means that `self.value`
 72            // is initialized and valid for shared access.
 73            Some(unsafe { &*self.value.get().cast() })
 74        } else {
 75            None
 76        }
 77    }
 78
 79    /// Populate the [`SetOnce`].
 80    ///
 81    /// Returns `true` if the [`SetOnce`] was successfully populated.
 82    pub fn populate(&self, value: T) -> bool {
 83        // INVARIANT: If the swap succeeds:
 84        //  - We increase `init`.
 85        //  - We write the valid value `1` to `init`.
 86        //  - Only one thread can succeed in this write, so we have exclusive access after the
 87        //    write.
 88        if let Ok(0) = self.init.cmpxchg(0, 1, Relaxed) {
 89            // SAFETY: By the type invariants of `Self`, the fact that we succeeded in writing `1`
 90            // to `self.init` means we obtained exclusive access to `self.value`.
 91            unsafe { core::ptr::write(self.value.get().cast(), value) };
 92            // INVARIANT:
 93            //  - We increase `init`.
 94            //  - We write the valid value `2` to `init`.
 95            //  - We release our exclusive access to `self.value` and it is now valid for shared
 96            //    access.
 97            self.init.store(2, Release);
 98            true
 99        } else {
100            false
101        }
102    }
103
104    /// Get a copy of the contained object.
105    ///
106    /// Returns [`None`] if the [`SetOnce`] is empty.
107    pub fn copy(&self) -> Option<T>
108    where
109        T: Copy,
110    {
111        self.as_ref().copied()
112    }
113}
114
115impl<T> Drop for SetOnce<T> {
116    fn drop(&mut self) {
117        if *self.init.get_mut() == 2 {
118            let value = self.value.get_mut();
119            // SAFETY: By the type invariants of `Self`, `self.init == 2` means that `self.value`
120            // contains a valid value. We have exclusive access, as we hold a `mut` reference to
121            // `self`.
122            unsafe { value.assume_init_drop() };
123        }
124    }
125}