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

vstd_extra/
once.rs
