embassy_sync/
lazy_lock.rs

1//! Synchronization primitive for initializing a value once, allowing others to get a reference to the value.
2
3use core::cell::UnsafeCell;
4use core::mem::ManuallyDrop;
5use core::sync::atomic::{AtomicBool, Ordering};
6
7/// The `LazyLock` is a synchronization primitive that allows for
8/// initializing a value once, and allowing others to obtain a
9/// reference to the value. This is useful for lazy initialization of
10/// a static value.
11///
12/// # Example
13/// ```
14/// use futures_executor::block_on;
15/// use embassy_sync::lazy_lock::LazyLock;
16///
17/// // Define a static value that will be lazily initialized
18/// // at runtime at the first access.
19/// static VALUE: LazyLock<u32> = LazyLock::new(|| 20);
20///
21/// let reference = VALUE.get();
22/// assert_eq!(reference, &20);
23/// ```
24pub struct LazyLock<T, F = fn() -> T> {
25    init: AtomicBool,
26    data: UnsafeCell<Data<T, F>>,
27}
28
29union Data<T, F> {
30    value: ManuallyDrop<T>,
31    f: ManuallyDrop<F>,
32}
33
34unsafe impl<T, F> Sync for LazyLock<T, F>
35where
36    T: Sync,
37    F: Sync,
38{
39}
40
41impl<T, F: FnOnce() -> T> LazyLock<T, F> {
42    /// Create a new uninitialized `StaticLock`.
43    pub const fn new(init_fn: F) -> Self {
44        Self {
45            init: AtomicBool::new(false),
46            data: UnsafeCell::new(Data {
47                f: ManuallyDrop::new(init_fn),
48            }),
49        }
50    }
51
52    /// Get a reference to the underlying value, initializing it if it
53    /// has not been done already.
54    #[inline]
55    pub fn get(&self) -> &T {
56        self.ensure_init_fast();
57        unsafe { &(*self.data.get()).value }
58    }
59
60    /// Get a mutable reference to the underlying value, initializing it if it
61    /// has not been done already.
62    #[inline]
63    pub fn get_mut(&mut self) -> &mut T {
64        self.ensure_init_fast();
65        unsafe { &mut (*self.data.get()).value }
66    }
67
68    /// Consume the `LazyLock`, returning the underlying value. The
69    /// initialization function will be called if it has not been
70    /// already.
71    #[inline]
72    pub fn into_inner(self) -> T {
73        self.ensure_init_fast();
74        let this = ManuallyDrop::new(self);
75        let data = unsafe { core::ptr::read(&this.data) }.into_inner();
76
77        ManuallyDrop::into_inner(unsafe { data.value })
78    }
79
80    /// Initialize the `LazyLock` if it has not been initialized yet.
81    /// This function is a fast track to [`Self::ensure_init`]
82    /// which does not require a critical section in most cases when
83    /// the value has been initialized already.
84    /// When this function returns, `self.data` is guaranteed to be
85    /// initialized and visible on the current core.
86    #[inline]
87    fn ensure_init_fast(&self) {
88        if !self.init.load(Ordering::Acquire) {
89            self.ensure_init();
90        }
91    }
92
93    /// Initialize the `LazyLock` if it has not been initialized yet.
94    /// When this function returns, `self.data` is guaranteed to be
95    /// initialized and visible on the current core.
96    fn ensure_init(&self) {
97        critical_section::with(|_| {
98            if !self.init.load(Ordering::Acquire) {
99                let data = unsafe { &mut *self.data.get() };
100                let f = unsafe { ManuallyDrop::take(&mut data.f) };
101                let value = f();
102                data.value = ManuallyDrop::new(value);
103
104                self.init.store(true, Ordering::Release);
105            }
106        });
107    }
108}
109
110impl<T, F> Drop for LazyLock<T, F> {
111    fn drop(&mut self) {
112        if self.init.load(Ordering::Acquire) {
113            unsafe { ManuallyDrop::drop(&mut self.data.get_mut().value) };
114        } else {
115            unsafe { ManuallyDrop::drop(&mut self.data.get_mut().f) };
116        }
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    use core::sync::atomic::{AtomicU32, Ordering};
123
124    use super::*;
125
126    #[test]
127    fn test_lazy_lock() {
128        static VALUE: LazyLock<u32> = LazyLock::new(|| 20);
129        let reference = VALUE.get();
130        assert_eq!(reference, &20);
131    }
132    #[test]
133    fn test_lazy_lock_mutation() {
134        let mut value: LazyLock<u32> = LazyLock::new(|| 20);
135        *value.get_mut() = 21;
136        let reference = value.get();
137        assert_eq!(reference, &21);
138    }
139    #[test]
140    fn test_lazy_lock_into_inner() {
141        let lazy: LazyLock<u32> = LazyLock::new(|| 20);
142        let value = lazy.into_inner();
143        assert_eq!(value, 20);
144    }
145
146    static DROP_CHECKER: AtomicU32 = AtomicU32::new(0);
147    struct DropCheck;
148
149    impl Drop for DropCheck {
150        fn drop(&mut self) {
151            DROP_CHECKER.fetch_add(1, Ordering::Acquire);
152        }
153    }
154
155    #[test]
156    fn test_lazy_drop() {
157        let lazy: LazyLock<DropCheck> = LazyLock::new(|| DropCheck);
158        assert_eq!(DROP_CHECKER.load(Ordering::Acquire), 0);
159        lazy.get();
160        drop(lazy);
161        assert_eq!(DROP_CHECKER.load(Ordering::Acquire), 1);
162
163        let dropper = DropCheck;
164        let lazy_fn: LazyLock<u32, _> = LazyLock::new(move || {
165            let _a = dropper;
166            20
167        });
168        assert_eq!(DROP_CHECKER.load(Ordering::Acquire), 1);
169        drop(lazy_fn);
170        assert_eq!(DROP_CHECKER.load(Ordering::Acquire), 2);
171    }
172}
embassy_sync/lazy_lock.rs

embassy_sync/
lazy_lock.rs
