Start of serious allocator work.

1 files changed, 72 insertions, 0 deletions

diff --git a/crates/kernel/src/cpu_locals.rs b/crates/kernel/src/cpu_locals.rs
new file mode 100644
index 0000000..fd826aa
--- /dev/null
+++ b/crates/kernel/src/cpu_locals.rs
@@ -0,0 +1,72 @@
+use crate::{alloc::Heap, arch};
+use core::{cell::RefCell, marker::PhantomData, ops::DerefMut, ptr::addr_of_mut};
+use static_assertions::{assert_eq_size, assert_not_impl_any};
+
+/// The data that is stored in a per-CPU structure.
+#[derive(Debug)]
+pub struct CPULocals {
+    /// The index of this CPU.
+    pub cpu_number: usize,
+
+    /// The heap used by this CPU's allocator.
+    pub heap: RefCell<&'static mut Heap>,
+
+    /// A canary for the `CPULocals` being initialized.
+    canary: usize,
+
+    // This ensures that the type is not `Send`.
+    _phantom: PhantomData<*mut ()>,
+}
+
+impl CPULocals {
+    /// Creates a new instance of the `CPULocals`, using it to initialize this CPU's `CPULocals`.
+    ///
+    /// # Safety
+    ///
+    /// - The CPULocals must not have already been initialized.
+    pub unsafe fn init(cpu_number: usize, heap: &'static mut Heap) {
+        arch::get_cpu_locals().write(CPULocals {
+            cpu_number,
+            heap: RefCell::new(heap),
+            canary: CANARY,
+            _phantom: PhantomData,
+        });
+    }
+
+    /// Returns the instance of the `CPULocals` for this CPU.
+    pub fn get() -> &'static CPULocals {
+        let ptr = arch::get_cpu_locals();
+        // SAFETY: The entrypoint sets this up for hart0, and we allocate this for other harts.
+        unsafe {
+            let canary_ptr = addr_of_mut!((*ptr.as_ptr()).canary);
+            assert_eq!(
+                *canary_ptr, CANARY,
+                "CPULocals were not initialized (and we probably just did UB)"
+            );
+            ptr.as_ref()
+        }
+    }
+
+    /// Retrieves a reference to the heap.
+    ///
+    /// # Panics
+    ///
+    /// - Panics if the CPU-local heap was already borrowed.
+    /// - The returned guard will panic on deref if the heap was not initialized.
+    pub fn heap(&'static self) -> impl DerefMut<Target = &'static mut Heap> {
+        self.heap.borrow_mut()
+    }
+}
+
+assert_eq_size!(CPULocals, [usize; 4]);
+assert_not_impl_any!(CPULocals: Send);
+
+cfg_if::cfg_if! {
+    if #[cfg(target_pointer_width = "32")] {
+        const CANARY: usize = usize::from_le_bytes(*b"locl");
+    } else if #[cfg(target_pointer_width = "64")] {
+        const CANARY: usize = usize::from_le_bytes(*b"CPULocal");
+    } else {
+        compile_error!("unsupported platform");
+    }
+}