projects/linux·blob master · refs

133 lines · 4.2 KiB

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}
126
127// SAFETY: `SetOnce` can be transferred across thread boundaries iff the data it contains can.
128unsafe impl<T: Send> Send for SetOnce<T> {}
129
130// SAFETY: `SetOnce` synchronises access to the inner value via atomic operations,
131// so shared references are safe when `T: Sync`. Since the inner `T` may be dropped
132// on any thread, we also require `T: Send`.
133unsafe impl<T: Send + Sync> Sync for SetOnce<T> {}