Skip to main content

ostd/sync/
once.rs

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