1use core::{
2 cell::UnsafeCell,
3 marker::PhantomData,
4 ops::Deref,
5 sync::atomic::{AtomicBool, Ordering},
6};
7
8use crate::syscalls;
9
10enum LazyData<T, F: FnOnce() -> T> {
11 Uninitialized(F),
12 Initialized(T),
13 Initializing,
14}
15
16impl<T, F: FnOnce() -> T> LazyData<T, F> {
17 fn get_value(&self) -> Option<&T> {
18 match self {
19 LazyData::Initialized(ref value) => Some(value),
20 _ => None,
21 }
22 }
23
24 fn start_initialize(&mut self) -> F {
25 match self {
26 LazyData::Uninitialized(_) => {
27 let LazyData::Uninitialized(res) = core::mem::replace(self, LazyData::Initializing)
28 else {
29 unreachable!()
30 };
31 res
32 }
33 _ => panic!("LazyData::start_initialize called on initialized data"),
34 }
35 }
36}
37
38pub struct LazyCell<T> {
40 running_init: AtomicBool,
41 value: UnsafeCell<LazyData<T, fn() -> T>>,
42 _marker: PhantomData<T>,
43}
44
45impl<T> LazyCell<T> {
46 pub const fn new(call: fn() -> T) -> Self {
47 Self {
48 running_init: AtomicBool::new(false),
49 value: UnsafeCell::new(LazyData::Uninitialized(call)),
50 _marker: PhantomData,
51 }
52 }
53
54 pub fn get(&self) -> &T {
56 let wait_for_init = || {
57 while self.running_init.load(Ordering::Acquire) {
58 syscalls::thread::yield_now();
59 }
60
61 unsafe {
62 (&*self.value.get())
63 .get_value()
64 .expect("Lazy awaited initialization but the value was never initialized")
65 }
66 };
67
68 match unsafe { &*self.value.get() } {
69 LazyData::Initialized(ref value) => value,
70 LazyData::Initializing => wait_for_init(),
71 LazyData::Uninitialized(_) => {
72 if self
73 .running_init
74 .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
75 .is_err()
76 {
77 wait_for_init()
78 } else {
79 let f = unsafe { (*self.value.get()).start_initialize() };
80 let value = (f)();
81 unsafe {
82 *self.value.get() = LazyData::Initialized(value);
83 }
84
85 self.running_init.store(false, Ordering::Release);
86 unsafe { (&*self.value.get()).get_value().unwrap() }
87 }
88 }
89 }
90 }
91}
92
93impl<T> Deref for LazyCell<T> {
94 type Target = T;
95
96 fn deref(&self) -> &Self::Target {
97 self.get()
98 }
99}
100
101unsafe impl<T: Send> Send for LazyCell<T> {}
102unsafe impl<T: Sync> Sync for LazyCell<T> {}