Skip to main content

ostd/task/scheduler/
info.rs

1// SPDX-License-Identifier: MPL-2.0
2//! Scheduling related information in a task.
3use vstd::{
4    atomic::{PAtomicU32, PermissionU32},
5    prelude::*,
6};
7
8use crate::cpu::{axiom_cpu_count_bounds, cpu_count, lemma_cpu_id_model, lemma_type_inv_range};
9
10// use core::sync::atomic::{AtomicU32, Ordering};
11
12use crate::{cpu::CpuId /*, task::Task */};
13
14/// Fields of a task that OSTD will never touch.
15///
16/// The type ought to be defined by the OSTD user and injected into the task.
17/// They are not part of the dynamic task data because it's slower there for
18/// the user-defined scheduler to access. The better ways to let the user
19/// define them, such as
20/// [existential types](https://github.com/rust-lang/rfcs/pull/2492) do not
21/// exist yet. So we decide to define them in OSTD.
22/* Origin Rust: #[derive(Debug)] (dropped: PAtomicU32 has no Debug impl) */
23#[verus_verify]
24pub struct TaskScheduleInfo {
25    /// The CPU that the task would like to be running on.
26    pub cpu: AtomicCpuId,
27}
28
29/// An atomic CPUID container.
30/* Origin Rust: #[derive(Debug)] pub struct AtomicCpuId(AtomicU32); */
31#[verus_verify]
32pub struct AtomicCpuId(PAtomicU32);
33
34verus! {
35
36impl AtomicCpuId {
37    /// The null value of CPUID.
38    ///
39    /// An `AtomicCpuId` with `AtomicCpuId::NONE` as its inner value is empty.
40    /* `pub` is required by Verus since public contracts refer to it. */
41    pub const NONE: u32 = u32::MAX;
42
43    /// Sets the inner value of an `AtomicCpuId` if it's empty.
44    ///
45    /// The return value is a result indicating whether the new value was written
46    /// and containing the previous value. If the previous value is empty, it returns
47    /// `Ok(())`. Otherwise, it returns `Err(previous_value)` which the previous
48    /// value is a valid CPU ID.
49    #[verus_spec(
50        ret =>
51        with
52            Tracked(perm): Tracked<&mut PermissionU32>,
53        requires
54            old(perm).view().patomic == self.patomic_id(),
55            Self::inv_value(old(perm).view().value),
56        ensures
57            self.patomic_id() == final(perm).view().patomic,
58            ret is Ok ==> final(perm).view().value == cpu_id.as_usize() as u32,
59            ret matches Err(prev) ==> prev@ == old(perm).view().value
60                && final(perm).view().value == old(perm).view().value,
61            Self::inv_value(final(perm).view().value),
62    )]
63    pub fn set_if_is_none(&self, cpu_id: CpuId) -> core::result::Result<(), CpuId> {
64        proof! {
65            broadcast use {
66                axiom_cpu_count_bounds,
67                lemma_cpu_id_model,
68                lemma_type_inv_range,
69            };
70            use_type_invariant(&cpu_id);
71            assert(Self::inv_value(cpu_id.as_usize() as u32));
72        }
73        /* PAtomicU32::compare_exchange takes the tracked permission and hardcodes
74         * `SeqCst`, so the call drops the two `Ordering::Relaxed` arguments; the
75         * `map_err` closure needs its own spec to carry the `unwrap` precondition.
76         * Origin Rust: self.0.compare_exchange(Self::NONE, cpu_id.as_usize() as u32, Ordering::Relaxed, Ordering::Relaxed).map(|_| ()).map_err(|prev| (prev as usize).try_into().unwrap())
77         */
78        self.0.compare_exchange(Tracked(perm), Self::NONE, cpu_id.as_usize() as u32).map(
79            |_| (),
80        ).map_err(
81            |prev: u32| -> (payload: CpuId)
82                requires
83                    cpu_count() > prev,
84                ensures
85                    payload@ == prev,
86                { (prev as usize).try_into().unwrap() },
87        )
88    }
89
90    /// Sets the inner value of an `AtomicCpuId` anyway.
91    #[verus_spec(
92        with
93            Tracked(perm): Tracked<&mut PermissionU32>,
94        requires
95            old(perm).view().patomic == self.patomic_id(),
96        ensures
97            self.patomic_id() == final(perm).view().patomic,
98            final(perm).view().value == cpu_id.as_usize() as u32,
99            Self::inv_value(final(perm).view().value),
100    )]
101    pub fn set_anyway(&self, cpu_id: CpuId) {
102        proof! {
103            broadcast use {
104                axiom_cpu_count_bounds,
105                lemma_cpu_id_model,
106                lemma_type_inv_range,
107            };
108            use_type_invariant(&cpu_id);
109            assert(Self::inv_value(cpu_id.as_usize() as u32));
110        }
111        /* Origin Rust: self.0.store(cpu_id.as_usize() as u32, Ordering::Relaxed); */
112        self.0.store(Tracked(perm), cpu_id.as_usize() as u32);
113    }
114
115    /// Sets the inner value of an `AtomicCpuId` to `AtomicCpuId::NONE`, i.e. makes
116    /// an `AtomicCpuId` empty.
117    #[verus_spec(
118        with
119            Tracked(perm): Tracked<&mut PermissionU32>,
120        requires
121            old(perm).view().patomic == self.patomic_id(),
122        ensures
123            self.patomic_id() == final(perm).view().patomic,
124            final(perm).view().value == Self::NONE,
125            Self::inv_value(final(perm).view().value),
126    )]
127    pub fn set_to_none(&self) {
128        /* Origin Rust: self.0.store(Self::NONE, Ordering::Relaxed); */
129        self.0.store(Tracked(perm), Self::NONE);
130    }
131
132    /// Gets the inner value of an `AtomicCpuId`.
133    #[verus_spec(
134        ret =>
135        with
136            Tracked(perm): Tracked<&PermissionU32>,
137        requires
138            perm.view().patomic == self.patomic_id(),
139            Self::inv_value(perm.view().value),
140        ensures
141            ret is None == (perm.view().value == Self::NONE),
142            ret matches Some(id) ==> id@ == perm.view().value,
143    )]
144    pub fn get(&self) -> Option<CpuId> {
145        /* Origin Rust: let val = self.0.load(Ordering::Relaxed); */
146        let val = self.0.load(Tracked(perm));
147        if val == Self::NONE {
148            None
149        } else {
150            Some((val as usize).try_into().ok()?)
151        }
152    }
153
154    /// The identity of the inner atomic cell, for tying permissions to `self`.
155    /// A closed helper, since a `pub` contract cannot use the private field.
156    pub closed spec fn patomic_id(&self) -> int {
157        self.0.id()
158    }
159
160    /// The trace invariant of the CPUID cell: the stored word is the null
161    /// constant `NONE` or encodes a valid CPU ID (below the CPU count).
162    pub open spec fn inv_value(v: u32) -> bool {
163        v == Self::NONE || 0 <= v < cpu_count()
164    }
165}
166
167} // verus!
168impl Default for AtomicCpuId {
169    fn default() -> Self {
170        Self(PAtomicU32::new(Self::NONE).0)
171    }
172}
173
174/* impl CommonSchedInfo for Task {
175    fn cpu(&self) -> &AtomicCpuId {
176        &self.schedule_info().cpu
177    }
178} */
179
180/// Trait for fetching common scheduling information.
181pub trait CommonSchedInfo {
182    /// Gets the CPU that the task is running on or lately ran on.
183    fn cpu(&self) -> &AtomicCpuId;
184}