ostd/sync/once.rs
1use vstd::{
2 atomic_with_ghost,
3 cell::{MemContents, PCell, PointsTo},
4 modes::tracked_static_ref,
5 prelude::*,
6};
7
8verus! {
9
10pub const UNINIT: u64 = 0;
11
12pub const OCCUPIED: u64 = 1;
13
14pub const INITED: u64 = 2;
15
16/// A tracked state of a¸ [`Once`] that can be used to ensure that the cell is
17/// initialized before accessing its value.
18pub tracked enum OnceState<V: 'static> {
19 /// The cell is uninitialized.
20 Uninit(PointsTo<Option<V>>),
21 /// The cell is occupied meaning it is being *written*.
22 Occupied,
23 /// The cell is initialized with a value and extended with
24 /// static lifetime.
25 Init(&'static PointsTo<Option<V>>),
26}
27
28/// A structure that combines some data with a permission to access it.
29///
30/// For example, in `aster_common` we can see a lot of structs with
31/// its `owner` associated. E.g., `MetaSlotOwner` is the owner of
32/// `MetaSlot`. This struct can be used to represent such a combination
33/// because now the permission is no longer exclusively owner by some
34/// specific CPU and is "shared" among multiple threads via atomic
35/// operations.
36///
37/// This struct is especially useful when used in conjunction with
38/// synchronization primitives like [`Once`], where we want to ensure that
39/// the data is initialized only once and the permission is preserved
40/// throughout the lifetime of the data.
41#[repr(transparent)]
42#[allow(repr_transparent_non_zst_fields)]
43pub struct AtomicDataWithOwner<V, Own> {
44 /// The underlying data.
45 pub data: V,
46 /// The permission to access the data.
47 pub permission: Tracked<Own>,
48}
49
50impl<U, Own> View for AtomicDataWithOwner<U, Own> {
51 type V = U;
52
53 open spec fn view(&self) -> Self::V {
54 self.data
55 }
56}
57
58impl<V, Own> AtomicDataWithOwner<V, Own> {
59 /// Creates a new `AtomicDataWithOwner` with the given data and permission.
60 pub const fn new(data: V, permission: Tracked<Own>) -> (r: Self)
61 ensures
62 r.data == data,
63 r.permission == permission,
64 {
65 AtomicDataWithOwner { data, permission }
66 }
67}
68
69/// A [`Predicate`] is something you're gonna preserve during the lifetime
70/// of any synchronization primitives like [`Once`].
71pub trait Predicate<V> {
72 spec fn inv(self, v: V) -> bool;
73}
74
75/// A trivial predicate that holds for any value.
76/// Use with [`OnceImpl`] when no invariant is needed.
77pub struct TrivialPred;
78
79impl<V> Predicate<V> for TrivialPred {
80 open spec fn inv(self, v: V) -> bool {
81 true
82 }
83}
84
85struct_with_invariants! {
86/// A synchronization primitive which can nominally be written to only once.
87///
88/// This type is a thread-safe [`Once`], and can be used in statics.
89/// In many simple cases, you can use [`LazyLock<T, F>`] instead to get the benefits of this type
90/// with less effort: `LazyLock<T, F>` "looks like" `&T` because it initializes with `F` on deref!
91/// Where OnceLock shines is when LazyLock is too simple to support a given case, as LazyLock
92/// doesn't allow additional inputs to its function after you call [`LazyLock::new(|| ...)`].
93///
94/// A `OnceLock` can be thought of as a safe abstraction over uninitialized data that becomes
95/// initialized once written.
96///
97/// # Examples
98///
99/// ```rust
100/// static MY_ONCE: Once<i32> = Once::new();
101///
102/// let value = MY_ONCE.get();
103/// assert(value.is_some()); // unsatisfied precondition, as MY_ONCE is uninitialized.
104/// ```
105#[verifier::reject_recursive_types(V)]
106pub struct OnceImpl<V: 'static, F: Predicate<V>> {
107 cell: PCell<Option<V>>,
108 state: vstd::atomic_ghost::AtomicU64<_, OnceState<V>, _>,
109 f: Ghost<F>,
110}
111
112#[verifier::type_invariant]
113pub closed spec fn wf(&self) -> bool {
114 invariant on state with (cell, f) is (v: u64, g: OnceState<V>) {
115 match g {
116 OnceState::Uninit(points_to) => {
117 &&& v == UNINIT
118 &&& points_to.id() == cell.id()
119 &&& points_to.mem_contents() matches MemContents::Init(None)
120 }
121 OnceState::Occupied => {
122 &&& v == OCCUPIED
123 }
124 OnceState::Init(points_to) => {
125 &&& v == INITED
126 &&& points_to.id() == cell.id()
127 &&& points_to.mem_contents() matches MemContents::Init(Some(value))
128 &&& f@.inv(value)
129 }
130 }
131 }
132}
133
134}
135
136#[verifier::external]
137unsafe impl<V, F: Predicate<V>> Send for OnceImpl<V, F> {
138
139}
140
141#[verifier::external]
142unsafe impl<V, F: Predicate<V>> Sync for OnceImpl<V, F> {
143
144}
145
146impl<V, F: Predicate<V>> OnceImpl<V, F> {
147 pub closed spec fn inv(&self) -> F {
148 self.f@
149 }
150
151 /// Creates a new uninitialized [`Once`].
152 pub const fn new(Ghost(f): Ghost<F>) -> (r: Self)
153 ensures
154 r.wf(),
155 r.inv() == f,
156 {
157 let (cell, Tracked(points_to)) = PCell::new(None);
158 let tracked state = OnceState::Uninit(points_to);
159 let state = vstd::atomic_ghost::AtomicU64::new(
160 Ghost((cell, Ghost(f))),
161 UNINIT,
162 Tracked(state),
163 );
164
165 Self { cell, state, f: Ghost(f) }
166 }
167
168 /// Initializes the [`Once`] with the given value `v`.
169 pub fn init(&self, v: V)
170 requires
171 self.inv().inv(v),
172 self.wf(),
173 {
174 let cur_state =
175 atomic_with_ghost! {
176 &self.state => load(); ghost g => {}
177 };
178
179 if cur_state != UNINIT {
180 return ;
181 } else {
182 let tracked mut points_to = None;
183 let res =
184 atomic_with_ghost! {
185 &self.state => compare_exchange(UNINIT, OCCUPIED);
186 returning res; ghost g => {
187 g = match g {
188 OnceState::Uninit(points_to_inner) => {
189 points_to = Some(points_to_inner);
190 OnceState::Occupied
191 }
192 _ => {
193 // If we are not in Uninit state, we cannot do anything.
194 g
195 }
196 }
197 }
198 };
199
200 if !res.is_err() {
201 let tracked mut points_to = points_to.tracked_unwrap();
202 self.cell.replace(Tracked(&mut points_to), Some(v));
203 // Extending the permission to static because `OnceLock` is
204 // often shared among threads and we want to ensure that
205 // the value is accessible globally.
206 let tracked static_points_to = tracked_static_ref(points_to);
207 // let tracked _ = self.inst.borrow().do_deposit(points_to, points_to, &mut token);
208 atomic_with_ghost! {
209 &self.state => store(INITED); ghost g => {
210 g = OnceState::Init(static_points_to);
211 }
212 }
213 return ;
214 } else {
215 // wait or abort.
216 return ;
217 }
218 }
219 }
220
221 /// Try to get the value stored in the [`Once`]. A [`Option::Some`]
222 /// is returned if the cell is initialized; otherwise, [`Option::None`]
223 /// is returned.
224 pub fn get<'a>(&'a self) -> (r: Option<&'a V>)
225 requires
226 self.wf(),
227 ensures
228 self.wf(),
229 r matches Some(res) ==> self.inv().inv(*res),
230 {
231 let tracked mut points_to = None;
232 let res =
233 atomic_with_ghost! {
234 &self.state => load(); ghost g => {
235 match g {
236 OnceState::Init(points_to_opt) => {
237 points_to = Some(points_to_opt);
238 }
239 _ => {}
240 }
241 }
242 };
243
244 if res == INITED {
245 let tracked points_to = points_to.tracked_unwrap();
246 let tracked static_points_to = tracked_static_ref(points_to);
247
248 self.cell.borrow(Tracked(static_points_to)).as_ref()
249 } else {
250 None
251 }
252 }
253}
254
255/// A `Once` that combines some data with a permission to access it.
256///
257/// This type alias automatically lifts the target value `V` into
258/// a wrapper [`AtomicDataWithOwner<V, Own>`] where `Own` is the
259/// permission type so that we can reason about non-trivial runtime
260/// properties in verification.
261pub type Once<V, Own, F> = OnceImpl<AtomicDataWithOwner<V, Own>, F>;
262
263} // verus!