xref: /linux/rust/kernel/sync/set_once.rs (revision b0319c4642638bad4b36974055b1c0894b2c7aa9)
1 // SPDX-License-Identifier: GPL-2.0
2 
3 //! A container that can be initialized at most once.
4 
5 use super::atomic::{
6     ordering::{Acquire, Relaxed, Release},
7     Atomic,
8 };
9 use 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 /// ```
43 pub struct SetOnce<T> {
44     init: Atomic<u32>,
45     value: UnsafeCell<MaybeUninit<T>>,
46 }
47 
48 impl<T> Default for SetOnce<T> {
49     fn default() -> Self {
50         Self::new()
51     }
52 }
53 
54 impl<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 
115 impl<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 }
126