openhcl_boot/
single_threaded.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! Support for working with global variables in a single-threaded environment.
5//! In such an environment, it is safe to access globals even if they don't
6//! implement [`Sync`], since there is only one thread that can access them. But
7//! code still needs to be careful to avoid creating multiple _mutable_
8//! references to the same global. These types provide abstractions for doing
9//! this safely.
10
11use core::cell::Cell;
12use core::cell::UnsafeCell;
13use core::ops::Deref;
14use core::ops::DerefMut;
15
16/// A wrapper around a value that implements `Sync` even if `T` does not
17/// implement `Sync`.
18///
19/// This is only safe to use in a single-threaded environment. Do not compile
20/// this type into a multi-threaded environment.
21pub struct SingleThreaded<T>(pub T);
22
23// SAFETY: we must mark this as Sync so that it can be `static`. It is
24// not actually necessarily Sync, so this can only be used in a
25// single-threaded environment.
26unsafe impl<T> Sync for SingleThreaded<T> {}
27
28impl<T> Deref for SingleThreaded<T> {
29    type Target = T;
30
31    fn deref(&self) -> &T {
32        &self.0
33    }
34}
35
36/// A reference returned by [`off_stack`].
37pub struct OffStackRef<'a, T>(&'a mut T, BorrowRef<'a>);
38
39impl<'a, T> OffStackRef<'a, T> {
40    #[track_caller]
41    #[doc(hidden)]
42    pub unsafe fn new_internal(value: &'a UnsafeCell<T>, used: &'a Cell<bool>) -> Self {
43        let r = BorrowRef::try_new(used).expect("function recursed");
44        // SAFETY: we just set `used` to true, so we know that we are the only
45        // one accessing `value`.
46        let value = unsafe { &mut *value.get() };
47        OffStackRef(value, r)
48    }
49
50    /// Leaks the borrow, returning the reference.
51    ///
52    /// This will lead to a panic if there is an attempt to borrow the value
53    /// again (e.g., if the function invoking the `off_stack` macro is called
54    /// again).
55    pub fn leak(this: Self) -> &'a mut T {
56        core::mem::forget(this.1);
57        this.0
58    }
59}
60
61struct BorrowRef<'a>(&'a Cell<bool>);
62
63impl<'a> BorrowRef<'a> {
64    // It is UB to use `off_stack` in a multi-threaded environment, which can
65    // happen depending on how unittests are run. To help catch this, we check
66    // that only the first thread to call `off_stack` matches all subsequent
67    // usage.
68    #[cfg(test)]
69    #[track_caller]
70    fn assert_single_threaded() {
71        static OFF_STACK_THREAD_ID: std::sync::OnceLock<std::thread::ThreadId> =
72            std::sync::OnceLock::new();
73
74        let current_thread = std::thread::current().id();
75        let off_stack_thread = OFF_STACK_THREAD_ID.get_or_init(|| current_thread);
76        if *off_stack_thread != current_thread {
77            panic!(
78                "off_stack! used from multiple test threads; run openhcl_boot \
79                tests with nextest or specify only a single test"
80            );
81        }
82    }
83
84    fn try_new(used: &'a Cell<bool>) -> Option<Self> {
85        #[cfg(test)]
86        Self::assert_single_threaded();
87
88        if used.replace(true) {
89            None
90        } else {
91            Some(Self(used))
92        }
93    }
94}
95
96impl Drop for BorrowRef<'_> {
97    fn drop(&mut self) {
98        self.0.set(false);
99    }
100}
101
102impl<T> Deref for OffStackRef<'_, T> {
103    type Target = T;
104    fn deref(&self) -> &T {
105        self.0
106    }
107}
108
109impl<T> DerefMut for OffStackRef<'_, T> {
110    fn deref_mut(&mut self) -> &mut T {
111        self.0
112    }
113}
114
115/// Returns a mutable reference to a value that is stored as a global `static`
116/// variable rather than exist on the stack.
117///
118/// This is useful for working with large objects that don't fit on the stack.
119/// It is an alternative to using [`SingleThreaded`] with
120/// [`RefCell`](core::cell::RefCell); `RefCell` has the disadvantage of putting
121/// an extra `bool` next to the value in memory, which can waste a lot of space
122/// for heavily-aligned objects.
123///
124/// Panics if this function is called recursively, since this would attempt to
125/// create multiple mutable references to the same global variable.
126///
127/// This only works in a single-threaded environment.
128///
129/// Note that when `off_stack` is used in a function that can be called multiple
130/// times, the caller must not assume the value is initialized with the value
131/// specified in the macro. For example, if `off_stack!(ArrayVec<u8>,
132/// ArrayVec::new_const())` is used in a function, the caller must not assume
133/// that the value is an empty vector, since the function could have been called
134/// before and left stale values, due to this being a wrapper around a global
135/// variable.
136macro_rules! off_stack {
137    ($ty:ty, $val:expr) => {{
138        use core::cell::Cell;
139        use core::cell::UnsafeCell;
140        use $crate::single_threaded::OffStackRef;
141        use $crate::single_threaded::SingleThreaded;
142
143        static VALUE: SingleThreaded<UnsafeCell<$ty>> = SingleThreaded(UnsafeCell::new($val));
144        static USED: SingleThreaded<Cell<bool>> = SingleThreaded(Cell::new(false));
145
146        // SAFETY: `USED` is always used to track the usage of `VALUE`.
147        unsafe { OffStackRef::new_internal(&VALUE.0, &USED.0) }
148    }};
149}
150pub(crate) use off_stack;
151
152#[cfg(test)]
153mod tests {
154    fn borrow_a() -> super::OffStackRef<'static, u32> {
155        off_stack!(u32, 0)
156    }
157
158    #[test]
159    fn off_stack_panics_when_used_by_two_threads() {
160        let _value = borrow_a();
161
162        let result = std::thread::Builder::new()
163            .spawn(|| {
164                let _other = borrow_a();
165            })
166            .unwrap()
167            .join();
168
169        assert!(result.is_err());
170    }
171}
openhcl_boot/single_threaded.rs

openhcl_boot/
single_threaded.rs
