kernel_api/sync/
once.rs

1use core::cell::UnsafeCell;
2use core::mem::MaybeUninit;
3use core::ops::Deref;
4use core::sync::atomic::{AtomicU8, fence, Ordering};
5
6pub struct Once(AtomicU8);
7
8#[derive(Debug, Copy, Clone, Eq, PartialEq)]
9#[repr(u8)]
10enum State {
11    Uncalled = 0,
12    Running = 1,
13    Called = 2,
14    Poison = 3
15}
16
17impl State {
18    const fn const_into_u8(self) -> u8 {
19        match self {
20            State::Uncalled => 0,
21            State::Running => 1,
22            State::Called => 2,
23            State::Poison => 3
24        }
25    }
26
27    const fn const_from_u8(value: u8) -> Result<Self, ()> {
28        match value {
29            0 => Ok(State::Uncalled),
30            1 => Ok(State::Running),
31            2 => Ok(State::Called),
32            3 => Ok(State::Poison),
33            _ => Err(())
34        }
35    }
36}
37
38impl From<State> for u8 {
39    fn from(value: State) -> Self {
40        value.const_into_u8()
41    }
42}
43
44impl TryFrom<u8> for State {
45    type Error = ();
46
47    fn try_from(value: u8) -> Result<Self, Self::Error> {
48        Self::const_from_u8(value)
49    }
50}
51
52impl Once {
53    pub const fn new() -> Self {
54        Self(AtomicU8::new(State::Uncalled.const_into_u8()))
55    }
56
57    pub fn call_once<F: FnOnce()>(&self, f: F) {
58        loop {
59            let current = self.0.compare_exchange_weak(State::Uncalled.into(), State::Running.into(), Ordering::Relaxed, Ordering::Acquire);
60            match current {
61                Ok(_) => break, // Switched from Uncalled to Running, call the function
62                Err(s) if s == State::Poison.into() => panic!("poisoned `Once`"),
63                Err(s) if s == State::Running.into() => {}, // Currently running, spin until state changes
64                Err(s) if s == State::Called.into() => return, // Already called, return immediately
65                Err(s) if s == State::Uncalled.into() => {}, // weak cas fail, try again
66                _ => unreachable!()
67            }
68            core::hint::spin_loop();
69        }
70
71        struct DropGuard<'a>(&'a Once);
72        impl Drop for DropGuard<'_> {
73            fn drop(&mut self) {
74                self.0.0.store(State::Poison.into(), Ordering::Relaxed);
75            }
76        }
77        let drop_guard = DropGuard(self);
78
79        f();
80
81        core::mem::forget(drop_guard);
82
83        self.0.store(State::Called.into(), Ordering::Release);
84    }
85
86    pub fn is_complete(&self) -> bool {
87        let state = self.0.load(Ordering::Relaxed).try_into().unwrap();
88        match state {
89            State::Called => true,
90            _ => false
91        }
92    }
93}
94
95pub struct OnceLock<T> {
96    data: UnsafeCell<MaybeUninit<T>>,
97    once: Once
98}
99
100unsafe impl<T: Send> Send for OnceLock<T> {}
101unsafe impl<T: Send + Sync> Sync for OnceLock<T> {}
102
103impl<T> OnceLock<T> {
104    pub const fn new() -> Self {
105        Self {
106            data: UnsafeCell::new(MaybeUninit::uninit()),
107            once: Once::new()
108        }
109    }
110
111    #[inline]
112    pub fn get(&self) -> Option<&T> {
113        if !self.once.is_complete() { return None; }
114        fence(Ordering::Acquire);
115
116        unsafe {
117            Some((*self.data.get()).assume_init_ref())
118        }
119    }
120
121    #[inline]
122    pub fn get_mut(&mut self) -> Option<&mut T> {
123        if !self.once.is_complete() { return None; }
124        fence(Ordering::Acquire);
125
126        unsafe {
127            Some((*self.data.get()).assume_init_mut())
128        }
129    }
130
131    #[inline]
132    pub fn get_or_init(&self, f: impl FnOnce() -> T) -> &T {
133        self.once.call_once(|| unsafe { (*self.data.get()).write(f()); });
134        unsafe { (*self.data.get()).assume_init_ref() }
135    }
136}
137
138pub struct LazyLock<T, F = fn() -> T> {
139    once: OnceLock<T>,
140    // FIXME: actually drop this when needed
141    f: MaybeUninit<F>
142}
143
144//unsafe impl<T, F: FnOnce() -> T> Send for LazyLock<T, F> {}
145unsafe impl<T: Sync + Send, F: Send> Sync for LazyLock<T, F> {}
146
147impl<T, F: FnOnce() -> T> LazyLock<T, F> {
148    pub const fn new(f: F) -> Self {
149        Self {
150            once: OnceLock::new(),
151            f: MaybeUninit::new(f)
152        }
153    }
154
155    pub fn force(this: &Self) -> &T {
156        this.once.get_or_init(unsafe {
157            core::ptr::read(this.f.as_ptr())
158        })
159    }
160}
161
162impl<T, F: FnOnce() -> T> Deref for LazyLock<T, F> {
163    type Target = T;
164
165    fn deref(&self) -> &Self::Target {
166        Self::force(self)
167    }
168}
169
170pub use bootstrap::BootstrapOnceLock;
171
172mod bootstrap {
173    use core::cell::UnsafeCell;
174    use core::mem::MaybeUninit;
175    use core::sync::atomic::{AtomicU8, fence, Ordering};
176
177    #[derive(Debug, Copy, Clone, Eq, PartialEq)]
178    #[repr(u8)]
179    enum State {
180        Uncalled = 0,
181        Saving = 1,
182        Running = 2,
183        Init = 3,
184        Poison = 4
185    }
186
187    impl State {
188        const fn const_into_u8(self) -> u8 {
189            match self {
190                State::Uncalled => 0,
191                State::Saving => 1,
192                State::Running => 2,
193                State::Init => 3,
194                State::Poison => 4
195            }
196        }
197
198        const fn const_from_u8(value: u8) -> Result<Self, ()> {
199            match value {
200                0 => Ok(State::Uncalled),
201                1 => Ok(State::Saving),
202                2 => Ok(State::Running),
203                3 => Ok(State::Init),
204                4 => Ok(State::Poison),
205                _ => Err(())
206            }
207        }
208    }
209
210    impl From<State> for u8 {
211        fn from(value: State) -> Self {
212            value.const_into_u8()
213        }
214    }
215
216    impl TryFrom<u8> for State {
217        type Error = ();
218
219        fn try_from(value: u8) -> Result<Self, Self::Error> {
220            Self::const_from_u8(value)
221        }
222    }
223
224    pub struct BootstrapOnceLock<T> {
225        data: UnsafeCell<MaybeUninit<T>>,
226        state: AtomicU8
227    }
228
229    unsafe impl<T> Send for BootstrapOnceLock<T> {}
230    unsafe impl<T> Sync for BootstrapOnceLock<T> {}
231
232    impl<T> BootstrapOnceLock<T> {
233        pub const fn new() -> Self {
234            Self {
235                data: UnsafeCell::new(MaybeUninit::uninit()),
236                state: AtomicU8::new(State::Uncalled.const_into_u8())
237            }
238        }
239
240        pub fn get(&self) -> Option<&T> {
241            let state: State = self.state.load(Ordering::Relaxed).try_into().unwrap();
242
243            if state == State::Poison { panic!("poisoned `BootstrapOnceLock`") }
244            if state == State::Uncalled || state == State::Saving { return None; }
245            fence(Ordering::Acquire);
246
247            unsafe {
248                Some((*self.data.get()).assume_init_ref())
249            }
250        }
251
252        /*
253        - Starts as `Uncalled`
254        - Move to `Saving`
255        - Store bootstrap value
256        - Move to `Running` - value is now legal to access
257        - Call function
258        - Move to `Saving` - value is now illegal to access
259        - Store new value
260        - Move to `Init` - value is now illegal to access
261         */
262        pub fn bootstrap(&self, bootstrap_value: T, f: impl FnOnce() -> T) -> &T {
263            loop {
264                let current = self.state.compare_exchange_weak(State::Uncalled.into(), State::Saving.into(), Ordering::Relaxed, Ordering::Acquire);
265                match current {
266                    Ok(_) => break, // Switched from Uncalled to Saving, bootstrap then call the function
267                    Err(s) if s == State::Poison.into() => panic!("poisoned `BootstrapOnceLock`"),
268                    Err(s) if s == State::Running.into() || s == State::Saving.into() => {}, // Currently running, spin until state changes
269                    Err(s) if s == State::Init.into() => {
270                        // Already called, return immediately
271                        return unsafe {
272                            (*self.data.get()).assume_init_ref()
273                        };
274                    },
275                    Err(s) if s == State::Uncalled.into() => {}, // Weak CAS failure so retry
276                    _ => unreachable!()
277                }
278                core::hint::spin_loop();
279            }
280
281            // We now need to bootstrap and init
282            unsafe { (*self.data.get()).write(bootstrap_value); }
283            // Release ordering so bootstrapped value syncs with Acquire ordering in Self::get
284            self.state.store(State::Running.into(), Ordering::Release);
285
286            let true_value = f();
287
288            // Relaxed ordering since no memory stuff to sync with (???)
289            self.state.store(State::Saving.into(), Ordering::Relaxed);
290            let ret = unsafe { (*self.data.get()).write(true_value) };
291            // Release ordering so bootstrapped value syncs with Acquire ordering in Self::get
292            self.state.store(State::Init.into(), Ordering::Release);
293            ret
294        }
295    }
296}
kernel_api/sync/once.rs

kernel_api/sync/
once.rs
