Enum API

Signed-off-by: Klim Tsoutsman <klim@tsoutsman.com>

Author: tsoutsman

kernel/cpu_local/src/lib.rs

@@ -36,20 +36,64 @@ pub struct FixedCpuLocal {
     pub size:   usize,
     pub align:  usize,
 }
+
 // NOTE: These fields must be kept in sync with `cpu_local::FixedCpuLocal`.
-impl FixedCpuLocal {
-    pub const CPU_ID:                       Self = Self { offset: 8,  size: 4, align: 4 };
-    pub const PREEMPTION_COUNT:             Self = Self { offset: 12, size: 1, align: 1 };
-    pub const TASK_SWITCH_PREEMPTION_GUARD: Self = Self { offset: 16, size: 8, align: 4 };
-    pub const DROP_AFTER_TASK_SWITCH:       Self = Self { offset: 24, size: 8, align: 8 };
+pub enum CpuLocalField {
+    CpuId,
+    PreemptionCount,
+    TaskSwitchPreemptionGuard,
+    DropAfterTaskSwitch,
+}
+
+impl CpuLocalField {
+    pub const fn offset(&self) -> usize {
+        match self {
+            Self::CpuId => 8,
+            Self::PreemptionCount => 12,
+            Self::TaskSwitchPreemptionGuard => 16,
+            Self::DropAfterTaskSwitch => 24,
+        }
+    }
+
+    // TODO: Do we even need to know the size and alignment?
+
+    // pub const fn size(&self) -> usize {
+    //     match self {
+    //         Self::CpuId => 4,
+    //         Self::PreemptionCount => 1,
+    //         Self::TaskSwitchPreemptionGuard => 8,
+    //         Self::DropAfterTaskSwitch => 8,
+    //     }
+    // }
+
+    // pub const fn align(&self) -> usize {
+    //     match self {
+    //         Self::CpuId => 4,
+    //         Self::PreemptionCount => 1,
+    //         Self::TaskSwitchPreemptionGuard => 4,
+    //         Self::DropAfterTaskSwitch => 8,
+    //     }
+    // }
+}
+
+// TODO: Could come up with a more imaginative name.
+pub unsafe trait Field: Sized {
+    const FIELD: CpuLocalField;
+
+    // const SIZE_CHECK: () = assert!(Self::FIELD.size() == size_of::<Self>());
+    // const ALIGN_CHECK: () = assert!(Self::FIELD.align() == align_of::<Self>());
 }
 
 /// A reference to a CPU-local variable.
 ///
 /// Note that this struct doesn't contain an instance of the type `T`,
 /// and dropping it has no effect.
-pub struct CpuLocal<const OFFSET: usize, T>(PhantomData<*mut T>);
-impl<const OFFSET: usize, T> CpuLocal<OFFSET, T> {
+pub struct CpuLocal<T: Field>(PhantomData<*mut T>);
+
+impl<T> CpuLocal<T>
+where
+    T: Field,
+{
     /// Creates a new reference to a predefined CPU-local variable.
     ///
     /// # Arguments
@@ -63,12 +107,8 @@ impl<const OFFSET: usize, T> CpuLocal<OFFSET, T> {
     ///
     /// Thus, the caller must guarantee that the type `T` is correct for the
     /// given `FixedCpuLocal`.
-    pub const unsafe fn new_fixed(
-        fixed: FixedCpuLocal,
-    ) -> Self {
-        assert!(OFFSET == fixed.offset);
-        assert!(size_of::<T>()  == fixed.size);
-        assert!(align_of::<T>() == fixed.align);
+    pub const unsafe fn new_fixed() -> Self {
+        // FIXME: Should this still be unsafe?
         Self(PhantomData)
     }
 
@@ -97,7 +137,7 @@ impl<const OFFSET: usize, T> CpuLocal<OFFSET, T> {
     where
         F: FnOnce(&T) -> R,
     {
-        let ptr_to_cpu_local = self.self_ptr() + OFFSET;
+        let ptr_to_cpu_local = self.self_ptr() + T::FIELD.offset();
         let local_ref = unsafe { &*(ptr_to_cpu_local as *const T) };
         func(local_ref)
     }
@@ -111,7 +151,7 @@ impl<const OFFSET: usize, T> CpuLocal<OFFSET, T> {
         F: FnOnce(&mut T) -> R,
     {
         let _held_interrupts = irq_safety::hold_interrupts();
-        let ptr_to_cpu_local = self.self_ptr() + OFFSET;
+        let ptr_to_cpu_local = self.self_ptr() + T::FIELD.offset();
         let local_ref_mut = unsafe { &mut *(ptr_to_cpu_local as *mut T) };
         func(local_ref_mut)
     }
@@ -134,7 +174,10 @@ impl<const OFFSET: usize, T> CpuLocal<OFFSET, T> {
     }
 }
 
-impl<const OFFSET: usize, T: Copy> CpuLocal<OFFSET, T> {
+impl<T> CpuLocal<T>
+where
+    T: Copy + Field,
+{
     /// Returns a copy of this `CpuLocal`'s inner value of type `T`.
     ///
     /// This is a convenience function only available for types where `T: Copy`.
@@ -209,18 +252,23 @@ pub fn init<P>(
 
     // TODO Remove, temp tests
     if true {
-        let test_value = CpuLocal::<32, u64>(PhantomData);
-        test_value.with(|tv| log::warn!("Got test_value: {:#X}", *tv));
-        log::warn!("Setting test_value to 0x12345678...");
-        test_value.with_mut(|tv| { *tv = 0x12345678; });
-        test_value.with(|tv| log::warn!("Got test_value: {:#X}", *tv));
-
-        let test_string = CpuLocal::<40, alloc::string::String>(PhantomData);
-        test_string.with(|s| log::warn!("Got test_string: {:?}", s));
-        let new_str = ", world!";
-        log::warn!("Appending {:?} to test_string...", new_str);
-        test_string.with_mut(|s| s.push_str(new_str));
-        test_string.with(|s| log::warn!("Got test_string: {:?}", s));
+        // NOTE: The fact that this previously compiled in `cpu_local` is
+        // indicative of an unsafe API, as the offsets (and types) are
+        // completely arbitrary. There's nothing stopping us from trying to
+        // access CpuLocal::<16, String> which would be undefined behaviour.
+
+        // let test_value = CpuLocal::<32, u64>(PhantomData);
+        // test_value.with(|tv| log::warn!("Got test_value: {:#X}", *tv));
+        // log::warn!("Setting test_value to 0x12345678...");
+        // test_value.with_mut(|tv| { *tv = 0x12345678; });
+        // test_value.with(|tv| log::warn!("Got test_value: {:#X}", *tv));
+
+        // let test_string = CpuLocal::<40, alloc::string::String>(PhantomData);
+        // test_string.with(|s| log::warn!("Got test_string: {:?}", s));
+        // let new_str = ", world!";
+        // log::warn!("Appending {:?} to test_string...", new_str);
+        // test_string.with_mut(|s| s.push_str(new_str));
+        // test_string.with(|s| log::warn!("Got test_string: {:?}", s));
     }
     Ok(())
 }

kernel/per_cpu/src/lib.rs

@@ -35,8 +35,8 @@
 
 extern crate alloc; // TODO temp remove this 
 
-use cpu::CpuId;
-use preemption::{PreemptionCount, PreemptionGuard};
+use cpu_local::{CpuLocalField, Field};
+use preemption::PreemptionGuard;
 use task::TaskRef;
 
 /// The data stored on a per-CPU basis in Theseus.
@@ -71,35 +71,64 @@ pub struct PerCpuData {
     /// A preemption guard used during task switching to ensure that one task switch
     /// cannot interrupt (preempt) another task switch already in progress.
     // task_switch_preemption_guard: Option<TestU32>, // TODO temp remove this
-    task_switch_preemption_guard: Option<PreemptionGuard>,
+    task_switch_preemption_guard: TaskSwitchPreemptionGuard,
     /// Data that should be dropped after switching away from a task that has exited.
     /// Currently, this contains the previous task's `TaskRef` that was removed
     /// from its TLS area during the last task switch away from it.
-    drop_after_task_switch: Option<TaskRef>,
+    drop_after_task_switch: DropAfterTaskSwitch,
     test_value: u64,
     test_string: alloc::string::String,
 }
+
 impl PerCpuData {
     /// Defines the initial values of each per-CPU state.
-    fn new(self_ptr: usize, cpu_id: CpuId) -> Self {
+    fn new(self_ptr: usize, cpu_id: cpu::CpuId) -> Self {
         Self {
             self_ptr,
-            cpu_id,
-            preemption_count: PreemptionCount::new(),
-            task_switch_preemption_guard: None,
-            drop_after_task_switch: None,
+            cpu_id: CpuId(cpu_id),
+            preemption_count: PreemptionCount(preemption::PreemptionCount::new()),
+            task_switch_preemption_guard: TaskSwitchPreemptionGuard(None),
+            drop_after_task_switch: DropAfterTaskSwitch(None),
             test_value: 0xDEADBEEF,
             test_string: alloc::string::String::from("test_string hello"),
 
         }
     }
 }
 
+#[repr(transparent)]
+pub struct CpuId(pub cpu::CpuId);
+
+unsafe impl Field for CpuId {
+    const FIELD: CpuLocalField = CpuLocalField::CpuId;
+}
+
+#[repr(transparent)]
+pub struct PreemptionCount(pub preemption::PreemptionCount);
+
+unsafe impl Field for PreemptionCount {
+    const FIELD: CpuLocalField = CpuLocalField::PreemptionCount;
+}
+
+#[repr(transparent)]
+pub struct TaskSwitchPreemptionGuard(pub Option<PreemptionGuard>);
+
+unsafe impl Field for TaskSwitchPreemptionGuard {
+    const FIELD: CpuLocalField = CpuLocalField::TaskSwitchPreemptionGuard;
+}
+
+#[repr(transparent)]
+pub struct DropAfterTaskSwitch(pub Option<TaskRef>);
+
+unsafe impl Field for DropAfterTaskSwitch {
+    const FIELD: CpuLocalField = CpuLocalField::DropAfterTaskSwitch;
+}
+
 /// Initializes the current CPU's `PerCpuData`.
 ///
 /// This must be invoked from (run on) the actual CPU with the given `cpu_id`;
 /// the main bootstrap CPU cannot run this for all CPUs itself.
-pub fn init(cpu_id: CpuId) -> Result<(), &'static str> {
+pub fn init(cpu_id: cpu::CpuId) -> Result<(), &'static str> {
     cpu_local::init(
         cpu_id.value(),
         core::mem::size_of::<PerCpuData>(),
@@ -109,27 +138,16 @@ pub fn init(cpu_id: CpuId) -> Result<(), &'static str> {
 
 mod const_assertions {
     use core::mem::{align_of, size_of};
-    use cpu_local::FixedCpuLocal;
+    use cpu_local::CpuLocalField;
     use memoffset::offset_of;
     use super::*;
 
-    const _: () = assert!(0 == offset_of!(PerCpuData, self_ptr));
     const _: () = assert!(8 == size_of::<usize>());
     const _: () = assert!(8 == align_of::<usize>());
 
-    const _: () = assert!(FixedCpuLocal::CPU_ID.offset == offset_of!(PerCpuData, cpu_id));
-    const _: () = assert!(FixedCpuLocal::CPU_ID.size   == size_of::<CpuId>());
-    const _: () = assert!(FixedCpuLocal::CPU_ID.align  == align_of::<CpuId>());
-
-    const _: () = assert!(FixedCpuLocal::PREEMPTION_COUNT.offset == offset_of!(PerCpuData, preemption_count));
-    const _: () = assert!(FixedCpuLocal::PREEMPTION_COUNT.size   == size_of::<PreemptionCount>());
-    const _: () = assert!(FixedCpuLocal::PREEMPTION_COUNT.align  == align_of::<PreemptionCount>());
-
-    const _: () = assert!(FixedCpuLocal::TASK_SWITCH_PREEMPTION_GUARD.offset == offset_of!(PerCpuData, task_switch_preemption_guard));
-    const _: () = assert!(FixedCpuLocal::TASK_SWITCH_PREEMPTION_GUARD.size   == size_of::<Option<PreemptionGuard>>());
-    const _: () = assert!(FixedCpuLocal::TASK_SWITCH_PREEMPTION_GUARD.align  == align_of::<Option<PreemptionGuard>>());
-
-    const _: () = assert!(FixedCpuLocal::DROP_AFTER_TASK_SWITCH.offset == offset_of!(PerCpuData, drop_after_task_switch));
-    const _: () = assert!(FixedCpuLocal::DROP_AFTER_TASK_SWITCH.size   == size_of::<Option<TaskRef>>());
-    const _: () = assert!(FixedCpuLocal::DROP_AFTER_TASK_SWITCH.align  == align_of::<Option<TaskRef >>());
+    const _: () = assert!(0 == offset_of!(PerCpuData, self_ptr));
+    const _: () = assert!(CpuLocalField::CpuId.offset() == offset_of!(PerCpuData, cpu_id));
+    const _: () = assert!(CpuLocalField::PreemptionCount.offset() == offset_of!(PerCpuData, preemption_count));
+    const _: () = assert!(CpuLocalField::TaskSwitchPreemptionGuard.offset() == offset_of!(PerCpuData, task_switch_preemption_guard));
+    const _: () = assert!(CpuLocalField::DropAfterTaskSwitch.offset() == offset_of!(PerCpuData, drop_after_task_switch));
 }