Skip to main content

o1_utils/
lazy_lock.rs

1//! A `no_std`-compatible replacement for [`std::sync::LazyLock`].
2//!
3//! When the `std` feature is enabled, this delegates directly to
4//! [`std::sync::LazyLock`]. In `no_std` mode, this delegates to
5//! [`spin::Lazy`] which uses a spin-based `Once` internally.
6
7use core::ops::Deref;
8
9/// A thread-safe, lazily-initialized value.
10///
11/// Constructed with a function (typically `fn() -> T`) so that it can be used
12/// in `static` items via [`LazyLock::new`], which is `const`.
13///
14/// # Examples
15///
16/// ```
17/// use o1_utils::lazy_lock::LazyLock;
18///
19/// static VALUE: LazyLock<u64> = LazyLock::new(|| 1 + 2);
20///
21/// assert_eq!(*VALUE, 3);
22/// ```
23#[cfg(feature = "std")]
24pub struct LazyLock<T, F = fn() -> T> {
25    inner: std::sync::LazyLock<T, F>,
26}
27
28#[cfg(feature = "std")]
29impl<T, F: FnOnce() -> T> LazyLock<T, F> {
30    pub const fn new(init: F) -> Self {
31        Self {
32            inner: std::sync::LazyLock::new(init),
33        }
34    }
35}
36
37#[cfg(feature = "std")]
38impl<T, F: FnOnce() -> T> Deref for LazyLock<T, F> {
39    type Target = T;
40
41    fn deref(&self) -> &T {
42        &self.inner
43    }
44}
45
46/// A thread-safe, lazily-initialized value (no-std fallback via [`spin::Lazy`]).
47#[cfg(not(feature = "std"))]
48pub struct LazyLock<T, F = fn() -> T> {
49    inner: spin::Lazy<T, F>,
50}
51
52#[cfg(not(feature = "std"))]
53impl<T, F: FnOnce() -> T> LazyLock<T, F> {
54    pub const fn new(init: F) -> Self {
55        Self {
56            inner: spin::Lazy::new(init),
57        }
58    }
59}
60
61#[cfg(not(feature = "std"))]
62impl<T, F: FnOnce() -> T> Deref for LazyLock<T, F> {
63    type Target = T;
64
65    fn deref(&self) -> &T {
66        &self.inner
67    }
68}
69
70#[cfg(test)]
71mod tests {
72    use super::LazyLock;
73    extern crate std;
74
75    #[test]
76    fn lazy_lock_panic() {
77        static VALUE: LazyLock<u64> = LazyLock::new(|| {
78            panic!("test_lazy_lock_panic");
79        });
80
81        std::thread::scope(|s| {
82            let error_counts = (0..4)
83                .map(|_| {
84                    s.spawn(|| {
85                        assert_eq!(*VALUE, 3);
86                    })
87                })
88                .map(|thread| thread.join().unwrap_err())
89                .map(|err| *err.downcast_ref::<&'static str>().unwrap())
90                .fold(std::collections::HashMap::new(), |mut acc, err| {
91                    *acc.entry(err).or_insert(0) += 1;
92                    acc
93                });
94
95            #[cfg(feature = "std")]
96            let poisoned_msg = "LazyLock instance has previously been poisoned";
97            #[cfg(not(feature = "std"))]
98            let poisoned_msg = "Once panicked";
99
100            assert_eq!(error_counts.get("test_lazy_lock_panic").copied(), Some(1));
101            assert_eq!(
102                error_counts.get(poisoned_msg).copied(),
103                Some(3),
104                "missing poisoned errors, use `std::dbg!` to see all errors"
105            );
106        });
107    }
108
109    #[test]
110    fn lazy_lock_success() {
111        static VALUE: LazyLock<u64> = LazyLock::new(|| 3);
112
113        std::thread::scope(|s| {
114            let threads = (0..4)
115                .map(|_| {
116                    s.spawn(|| {
117                        assert_eq!(*VALUE, 3);
118                    })
119                })
120                .collect::<alloc::vec::Vec<_>>();
121
122            for thread in threads {
123                thread.join().unwrap();
124            }
125
126            assert_eq!(*VALUE, 3);
127        });
128    }
129}