refactor: add a safe OnceLock data structure

Taken from rust, very minor modifications on my side.

hal_aarch64/src/irq.rs

@@ -7,10 +7,9 @@ use hal_core::{Error, TimerCallbackFn};
 use crate::devices::gicv2::GicV2;
 
 use hal_core::mm::PageAlloc;
+use hal_core::once_lock::OnceLock;
 use hal_core::IrqOps;
 
-use core::cell::OnceCell;
-
 use cortex_a::registers::*;
 use tock_registers::interfaces::{ReadWriteable, Writeable};
 
@@ -100,7 +99,8 @@ enum InterruptType {
     SerrorLowerElAarch32,
 }
 
-static mut IRQS: core::cell::OnceCell<&Aarch64Irqs> = core::cell::OnceCell::new();
+/// Need a once lock because this is unset until Aarch64Irqs::init has been called by the hal.
+static IRQS: OnceLock<&Aarch64Irqs> = OnceLock::new();
 
 #[no_mangle]
 unsafe extern "C" fn aarch64_common_trap(offset: u64) {
@@ -133,17 +133,17 @@ unsafe extern "C" fn aarch64_common_trap(offset: u64) {
 
 #[derive(Debug)]
 pub struct Aarch64Irqs {
-    irq_chip: OnceCell<GicV2>,
+    irq_chip: OnceLock<GicV2>,
     timer_callback: AtomicPtr<TimerCallbackFn>,
 }
 
-/// Safety: I know what I'm doing :D
+/// Safety: Not safe because GicV2 is not Sync
 unsafe impl Sync for Aarch64Irqs {}
 
 impl Aarch64Irqs {
     pub const fn new() -> Self {
         Self {
-            irq_chip: OnceCell::new(),
+            irq_chip: OnceLock::new(),
             timer_callback: AtomicPtr::new(ptr::null_mut()),
         }
     }
@@ -185,10 +185,8 @@ impl Aarch64Irqs {
 impl IrqOps for Aarch64Irqs {
     fn init(&'static self) {
         cortex_a::registers::VBAR_EL1.set(el1_vector_table as usize as u64);
-        unsafe {
-            IRQS.set(self)
-                .expect("looks like init has already been called")
-        };
+        IRQS.set(self)
+            .expect("looks like init has already been called");
     }
 
     fn init_irq_chip(&self, _allocator: &impl PageAlloc) -> Result<(), Error> {

hal_core/src/hal.rs

@@ -1,19 +1,19 @@
-use super::fake_once_lock::FakeOnceLock;
 use super::mm::{self, Mmu, PageAlloc, PageMap};
+use super::once_lock::OnceLock;
 use super::AddressRange;
 use super::Error;
 use super::ReentrantSpinlock;
 use super::{IrqOps, TimerCallbackFn};
 
 pub struct Hal<P: PageMap + 'static, I: IrqOps> {
-    kpt: FakeOnceLock<ReentrantSpinlock<&'static mut P>>,
+    kpt: OnceLock<ReentrantSpinlock<&'static mut P>>,
     irq_ops: I,
 }
 
 impl<P: PageMap + Mmu + 'static, I: IrqOps> Hal<P, I> {
     pub const fn new(irq_ops: I) -> Hal<P, I> {
         Self {
-            kpt: FakeOnceLock::new(),
+            kpt: OnceLock::new(),
             irq_ops,
         }
     }

hal_core/src/lib.rs

@@ -1,13 +1,14 @@
 #![no_std]
 #![feature(const_mut_refs)]
+#![feature(never_type)]
 
 use core::convert::Into;
 use core::ops::Range;
 
 mod hal;
 pub use hal::Hal;
 
-mod fake_once_lock;
+pub mod once_lock;
 
 mod reentrant_spinlock;
 pub use reentrant_spinlock::ReentrantSpinlock;

hal_core/src/once_lock.rs

@@ -0,0 +1,115 @@
+// Taken from https://doc.rust-lang.org/nightly/src/std/sync/once_lock.rs.html#121-128
+// Not mine!
+use spin::Once;
+
+use core::cell::UnsafeCell;
+use core::fmt;
+use core::mem::MaybeUninit;
+
+pub struct OnceLock<T> {
+    once: Once,
+    value: UnsafeCell<MaybeUninit<T>>,
+}
+
+impl<T> OnceLock<T> {
+    pub const fn new() -> OnceLock<T> {
+        OnceLock {
+            once: Once::new(),
+            value: UnsafeCell::new(MaybeUninit::uninit()),
+        }
+    }
+
+    pub fn get(&self) -> Option<&T> {
+        if self.is_initialized() {
+            Some(unsafe { self.get_unchecked() })
+        } else {
+            None
+        }
+    }
+
+    pub fn set(&self, value: T) -> Result<(), T> {
+        match self.try_insert(value) {
+            Ok(_) => Ok(()),
+            Err((_, value)) => Err(value),
+        }
+    }
+
+    fn try_insert(&self, value: T) -> Result<&T, (&T, T)> {
+        let mut value = Some(value);
+        let res = self.get_or_init(|| value.take().unwrap());
+        match value {
+            None => Ok(res),
+            Some(value) => Err((res, value)),
+        }
+    }
+
+    pub fn get_or_init<F>(&self, f: F) -> &T
+    where
+        F: FnOnce() -> T,
+    {
+        match self.get_or_try_init(|| Ok::<T, !>(f())) {
+            Ok(val) => val,
+            Err(_) => unreachable!(), // never type above
+        }
+    }
+
+    pub fn get_or_try_init<F, E>(&self, f: F) -> Result<&T, E>
+    where
+        F: FnOnce() -> Result<T, E>,
+    {
+        // Fast path check
+        // NOTE: We need to perform an acquire on the state in this method
+        // in order to correctly synchronize `LazyLock::force`. This is
+        // currently done by calling `self.get()`, which in turn calls
+        // `self.is_initialized()`, which in turn performs the acquire.
+        if let Some(value) = self.get() {
+            return Ok(value);
+        }
+        self.initialize(f)?;
+
+        debug_assert!(self.is_initialized());
+
+        // SAFETY: The inner value has been initialized
+        Ok(unsafe { self.get_unchecked() })
+    }
+
+    fn is_initialized(&self) -> bool {
+        self.once.is_completed()
+    }
+
+    fn initialize<F, E>(&self, f: F) -> Result<(), E>
+    where
+        F: FnOnce() -> Result<T, E>,
+    {
+        let mut res: Result<(), E> = Ok(());
+        let slot = &self.value;
+
+        self.once.call_once(|| match f() {
+            Ok(value) => {
+                unsafe { (*slot.get()).write(value) };
+            }
+            Err(e) => {
+                res = Err(e);
+            }
+        });
+        res
+    }
+
+    unsafe fn get_unchecked(&self) -> &T {
+        debug_assert!(self.is_initialized());
+        (*self.value.get()).assume_init_ref()
+    }
+}
+
+impl<T: fmt::Debug> fmt::Debug for OnceLock<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let mut d = f.debug_tuple("OnceLock");
+        match self.get() {
+            Some(v) => d.field(v),
+            None => d.field(&format_args!("<uninit>")),
+        };
+        d.finish()
+    }
+}
+
+unsafe impl<T: Sync + Send> Sync for OnceLock<T> {}