[PATCH V2 1/2] rust: cpu: Introduce CpuId abstraction

[Viresh Kumar <viresh.kumar@linaro.org>]

This adds abstraction for representing a CPU identifier.

Suggested-by: Boqun Feng <boqun.feng@gmail.com>
Signed-off-by: Viresh Kumar <viresh.kumar@linaro.org>
---
 rust/kernel/cpu.rs | 102 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 102 insertions(+)

diff --git a/rust/kernel/cpu.rs b/rust/kernel/cpu.rs
index 10c5c3b25873..0b545dbf5c83 100644
--- a/rust/kernel/cpu.rs
+++ b/rust/kernel/cpu.rs
@@ -6,6 +6,108 @@
 
 use crate::{bindings, device::Device, error::Result, prelude::ENODEV};
 
+/// Returns the maximum number of possible CPUs in the current system configuration.
+#[inline]
+pub fn nr_cpu_ids() -> u32 {
+    #[cfg(any(NR_CPUS_1, CONFIG_FORCE_NR_CPUS))]
+    {
+        bindings::NR_CPUS
+    }
+
+    #[cfg(not(any(NR_CPUS_1, CONFIG_FORCE_NR_CPUS)))]
+    // SAFETY: `nr_cpu_ids` is a valid global provided by the kernel.
+    unsafe {
+        bindings::nr_cpu_ids
+    }
+}
+
+/// The CPU ID.
+///
+/// Represents a CPU identifier as a wrapper around an `u32`.
+///
+/// # Invariants
+///
+/// The CPU ID must always lie within the range `[0, nr_cpu_ids())`.
+///
+/// ## Examples
+///
+/// ```
+/// use kernel::cpu::CpuId;
+///
+/// let cpu = 0;
+///
+/// // SAFETY: 0 is always a valid CPU number.
+/// let id = unsafe { CpuId::from_u32_unchecked(cpu) };
+///
+/// assert_eq!(id.as_u32(), cpu);
+/// assert!(CpuId::from_i32(0).is_some());
+/// assert!(CpuId::from_i32(-1).is_none());
+/// ```
+#[derive(Copy, Clone, PartialEq, Eq, Debug)]
+pub struct CpuId(u32);
+
+impl CpuId {
+    /// Creates a new [`CpuId`] from the given `id` without checking bounds.
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure that `id` is a valid CPU ID (i.e., `0 <= id < nr_cpu_ids()`).
+    #[inline]
+    pub unsafe fn from_i32_unchecked(id: i32) -> Self {
+        // INVARIANT: The function safety guarantees `id` is a valid CPU id.
+        Self(id as u32)
+    }
+
+    /// Creates a new [`CpuId`] from the given `id`, checking that it is valid.
+    pub fn from_i32(id: i32) -> Option<Self> {
+        if id < 0 || id as u32 >= nr_cpu_ids() {
+            None
+        } else {
+            // SAFETY: `id` has just been checked as a valid CPU ID.
+            Some(unsafe { Self::from_i32_unchecked(id) })
+        }
+    }
+
+    /// Creates a new [`CpuId`] from the given `id` without checking bounds.
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure that `id` is a valid CPU ID (i.e., `0 <= id < nr_cpu_ids()`).
+    #[inline]
+    pub unsafe fn from_u32_unchecked(id: u32) -> Self {
+        // INVARIANT: The function safety guarantees `id` is a valid CPU id.
+        Self(id)
+    }
+
+    /// Creates a new [`CpuId`] from the given `id`, checking that it is valid.
+    pub fn from_u32(id: u32) -> Option<Self> {
+        if id >= nr_cpu_ids() {
+            None
+        } else {
+            // SAFETY: `id` has just been checked as a valid CPU ID.
+            Some(unsafe { Self::from_u32_unchecked(id) })
+        }
+    }
+
+    /// Returns CPU number.
+    #[inline]
+    pub fn as_u32(&self) -> u32 {
+        self.0
+    }
+}
+
+impl From<CpuId> for u32 {
+    fn from(id: CpuId) -> Self {
+        id.as_u32()
+    }
+}
+
+impl From<CpuId> for i32 {
+    fn from(id: CpuId) -> Self {
+        id.as_u32() as i32
+    }
+}
+
 /// Creates a new instance of CPU's device.
 ///
 /// # Safety
-- 
2.31.1.272.g89b43f80a514