embassy_sync/
lazy_lock.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
//! Synchronization primitive for initializing a value once, allowing others to get a reference to the value.

use core::cell::UnsafeCell;
use core::mem::ManuallyDrop;
use core::sync::atomic::{AtomicBool, Ordering};

/// The `LazyLock` is a synchronization primitive that allows for
/// initializing a value once, and allowing others to obtain a
/// reference to the value. This is useful for lazy initialization of
/// a static value.
///
/// # Example
/// ```
/// use futures_executor::block_on;
/// use embassy_sync::lazy_lock::LazyLock;
///
/// // Define a static value that will be lazily initialized
/// // at runtime at the first access.
/// static VALUE: LazyLock<u32> = LazyLock::new(|| 20);
///
/// let reference = VALUE.get();
/// assert_eq!(reference, &20);
/// ```
pub struct LazyLock<T, F = fn() -> T> {
    init: AtomicBool,
    data: UnsafeCell<Data<T, F>>,
}

union Data<T, F> {
    value: ManuallyDrop<T>,
    f: ManuallyDrop<F>,
}

unsafe impl<T, F> Sync for LazyLock<T, F> {}

impl<T, F: FnOnce() -> T> LazyLock<T, F> {
    /// Create a new uninitialized `StaticLock`.
    pub const fn new(init_fn: F) -> Self {
        Self {
            init: AtomicBool::new(false),
            data: UnsafeCell::new(Data {
                f: ManuallyDrop::new(init_fn),
            }),
        }
    }

    /// Get a reference to the underlying value, initializing it if it
    /// has not been done already.
    #[inline]
    pub fn get(&self) -> &T {
        self.ensure_init_fast();
        unsafe { &(*self.data.get()).value }
    }

    /// Consume the `LazyLock`, returning the underlying value. The
    /// initialization function will be called if it has not been
    /// already.
    #[inline]
    pub fn into_inner(self) -> T {
        self.ensure_init_fast();
        let this = ManuallyDrop::new(self);
        let data = unsafe { core::ptr::read(&this.data) }.into_inner();

        ManuallyDrop::into_inner(unsafe { data.value })
    }

    /// Initialize the `LazyLock` if it has not been initialized yet.
    /// This function is a fast track to [`Self::ensure_init`]
    /// which does not require a critical section in most cases when
    /// the value has been initialized already.
    /// When this function returns, `self.data` is guaranteed to be
    /// initialized and visible on the current core.
    #[inline]
    fn ensure_init_fast(&self) {
        if !self.init.load(Ordering::Acquire) {
            self.ensure_init();
        }
    }

    /// Initialize the `LazyLock` if it has not been initialized yet.
    /// When this function returns, `self.data` is guaranteed to be
    /// initialized and visible on the current core.
    fn ensure_init(&self) {
        critical_section::with(|_| {
            if !self.init.load(Ordering::Acquire) {
                let data = unsafe { &mut *self.data.get() };
                let f = unsafe { ManuallyDrop::take(&mut data.f) };
                let value = f();
                data.value = ManuallyDrop::new(value);

                self.init.store(true, Ordering::Release);
            }
        });
    }
}

impl<T, F> Drop for LazyLock<T, F> {
    fn drop(&mut self) {
        if self.init.load(Ordering::Acquire) {
            unsafe { ManuallyDrop::drop(&mut self.data.get_mut().value) };
        } else {
            unsafe { ManuallyDrop::drop(&mut self.data.get_mut().f) };
        }
    }
}

#[cfg(test)]
mod tests {
    use core::sync::atomic::{AtomicU32, Ordering};

    use super::*;

    #[test]
    fn test_lazy_lock() {
        static VALUE: LazyLock<u32> = LazyLock::new(|| 20);
        let reference = VALUE.get();
        assert_eq!(reference, &20);
    }
    #[test]
    fn test_lazy_lock_into_inner() {
        let lazy: LazyLock<u32> = LazyLock::new(|| 20);
        let value = lazy.into_inner();
        assert_eq!(value, 20);
    }

    static DROP_CHECKER: AtomicU32 = AtomicU32::new(0);
    struct DropCheck;

    impl Drop for DropCheck {
        fn drop(&mut self) {
            DROP_CHECKER.fetch_add(1, Ordering::Acquire);
        }
    }

    #[test]
    fn test_lazy_drop() {
        let lazy: LazyLock<DropCheck> = LazyLock::new(|| DropCheck);
        assert_eq!(DROP_CHECKER.load(Ordering::Acquire), 0);
        lazy.get();
        drop(lazy);
        assert_eq!(DROP_CHECKER.load(Ordering::Acquire), 1);

        let dropper = DropCheck;
        let lazy_fn: LazyLock<u32, _> = LazyLock::new(move || {
            let _a = dropper;
            20
        });
        assert_eq!(DROP_CHECKER.load(Ordering::Acquire), 1);
        drop(lazy_fn);
        assert_eq!(DROP_CHECKER.load(Ordering::Acquire), 2);
    }
}