Use CpuId with validity ensurance rather than u32

ostd/src/cpu/mod.rs

@@ -21,6 +21,37 @@ use crate::{
     trap::DisabledLocalIrqGuard,
 };
 
+/// The ID of a CPU in the system.
+///
+/// If converting from/to an integer, the integer must start from 0 and be less
+/// than the number of CPUs.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct CpuId(u32);
+
+impl CpuId {
+    /// Returns the CPU ID of the bootstrap processor (BSP).
+    pub const fn bsp() -> Self {
+        CpuId(0)
+    }
+
+    /// Converts the CPU ID to an `usize`.
+    pub const fn as_usize(self) -> usize {
+        self.0 as usize
+    }
+}
+
+impl TryFrom<usize> for CpuId {
+    type Error = &'static str;
+
+    fn try_from(value: usize) -> Result<Self, Self::Error> {
+        if value < num_cpus() {
+            Ok(CpuId(value as u32))
+        } else {
+            Err("The given CPU ID is out of range")
+        }
+    }
+}
+
 /// The number of CPUs.
 static NUM_CPUS: Once<u32> = Once::new();
 
@@ -46,14 +77,20 @@ pub(crate) unsafe fn set_this_cpu_id(id: u32) {
 }
 
 /// Returns the number of CPUs.
-pub fn num_cpus() -> u32 {
+pub fn num_cpus() -> usize {
     debug_assert!(
         NUM_CPUS.get().is_some(),
         "The number of CPUs is not initialized"
     );
     // SAFETY: The number of CPUs is initialized. The unsafe version is used
     // to avoid the overhead of the check.
-    unsafe { *NUM_CPUS.get_unchecked() }
+    let num = unsafe { *NUM_CPUS.get_unchecked() };
+    num as usize
+}
+
+/// Returns an iterator over all CPUs.
+pub fn all_cpus() -> impl Iterator<Item = CpuId> {
+    (0..num_cpus()).map(|id| CpuId(id as u32))
 }
 
 /// A marker trait for guard types that can "pin" the current task to the
@@ -70,10 +107,10 @@ pub fn num_cpus() -> u32 {
 /// CPU while any one of the instances of the implemented structure exists.
 pub unsafe trait PinCurrentCpu {
     /// Returns the number of the current CPU.
-    fn current_cpu(&self) -> u32 {
+    fn current_cpu(&self) -> CpuId {
         let id = CURRENT_CPU.load();
         debug_assert_ne!(id, u32::MAX, "This CPU is not initialized");
-        id
+        CpuId(id)
     }
 }