add Singleton<T>

Author: Freax13

src/lib.rs

@@ -13,6 +13,9 @@
 #![warn(missing_docs)]
 #![deny(missing_debug_implementations)]
 
+use core::cell::UnsafeCell;
+use core::sync::atomic::{AtomicBool, Ordering};
+
 pub use crate::addr::{align_down, align_up, PhysAddr, VirtAddr};
 
 /// Makes a function const only when `feature = "const_fn"` is enabled.
@@ -112,3 +115,60 @@ impl PrivilegeLevel {
         }
     }
 }
+
+/// A wrapper that can be used to safely create one mutable reference `&'static mut T` from a static variable.
+///
+/// `Singleton` is safe because it ensures that it only ever gives out one reference.
+///
+/// ``Singleton<T>` is a safe alternative to `static mut` or a static `UnsafeCell<T>`.
+#[derive(Debug)]
+pub struct Singleton<T> {
+    used: AtomicBool,
+    value: UnsafeCell<T>,
+}
+
+impl<T> Singleton<T> {
+    /// Construct a new singleton.
+    pub const fn new(value: T) -> Self {
+        Self {
+            used: AtomicBool::new(false),
+            value: UnsafeCell::new(value),
+        }
+    }
+
+    /// Try to acquire a mutable reference to the wrapped value.
+    /// This will only succeed the first time the function is
+    /// called and fail on all following calls.
+    ///
+    /// ```
+    /// use x86_64::Singleton;
+    ///
+    /// static FOO: Singleton<i32> = Singleton::new(0);
+    ///
+    /// // Call `try_get_mut` for the first time and get a reference.
+    /// let first: &'static mut i32 = FOO.try_get_mut().unwrap();
+    /// assert_eq!(first, &0);
+    ///
+    /// // Calling `try_get_mut` again will return `None`.
+    /// assert_eq!(FOO.try_get_mut(), None);
+    /// ```
+    pub fn try_get_mut(&self) -> Option<&mut T> {
+        match self
+            .used
+            .compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed)
+        {
+            Ok(_) => Some(unsafe {
+                // SAFETY: no reference has been given out yet and we won't give out another.
+                &mut *self.value.get()
+            }),
+            Err(_) => None,
+        }
+    }
+}
+
+// SAFETY: Sharing a `Singleton<T>` between threads is safe regardless of whether `T` is `Sync`
+// because we only expose the inner value once to one thread.
+unsafe impl<T> Sync for Singleton<T> {}
+
+// SAFETY: It's safe to send a `Singleton<T>` to another thread if it's safe to send `T`.
+unsafe impl<T: Send> Send for Singleton<T> {}