cpumask: add functions similar to for_each_*_cpu

This commit adds three iterators which returns the CPU index. They are
PossibleCpusIndexIter, OnlineCpusIndexIter, and PresentCpusIndexIter
corresponding to C macros for_each_possible_cpu, for_each_online_cpu, and
for_each_present_cpu exclusively.

Signed-off-by: Li Hongyu <[email protected]>

Author: Richardhongyu

rust/kernel/cpumask.rs

@@ -0,0 +1,140 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! Cpumask variables and related functions.
+//!
+//! C header: [`include/linux/cpumask.h`](../../../../include/linux/cpumask.h).
+
+use crate::bindings;
+use core::iter::Iterator;
+
+/// An possible CPU index iterator.
+///
+/// This iterator has a similar abilitiy to the kernel's macro `for_each_possible_cpu`.
+pub struct PossibleCpusIndexIter {
+    index: i32,
+}
+
+/// An online CPU index iterator.
+///
+/// This iterator has a similar abilitiy to the kernel's macro `for_each_online_cpu`.
+pub struct OnlineCpusIndexIter {
+    index: i32,
+}
+
+/// An present CPU index iterator.
+///
+/// This iterator has a similar abilitiy to the kernel's macro `for_each_present_cpu`.
+pub struct PresentCpusIndexIter {
+    index: i32,
+}
+
+impl Iterator for PossibleCpusIndexIter {
+    type Item = u32;
+
+    fn next(&mut self) -> Option<u32> {
+        let next_cpu_id =
+            unsafe { bindings::cpumask_next(self.index, &bindings::__cpu_possible_mask) };
+        // When [`bindings::cpumask_next`] can not find further CPUs set in the
+        // [`bindings::__cpu_possible_mask`], it returns a value >= [`bindings::nr_cpu_ids`].
+        if next_cpu_id >= unsafe { bindings::nr_cpu_ids } {
+            return None;
+        }
+        self.index = next_cpu_id as i32;
+        Some(next_cpu_id)
+    }
+}
+
+impl Iterator for OnlineCpusIndexIter {
+    type Item = u32;
+
+    fn next(&mut self) -> Option<u32> {
+        let next_cpu_id =
+            unsafe { bindings::cpumask_next(self.index, &bindings::__cpu_online_mask) };
+        // When [`bindings::cpumask_next`] can not find further CPUs set in the
+        // [`bindings::__cpu_online_mask`], it returns a value >= [`bindings::nr_cpu_ids`].
+        if next_cpu_id >= unsafe { bindings::nr_cpu_ids } {
+            return None;
+        }
+        self.index = next_cpu_id as i32;
+        Some(next_cpu_id)
+    }
+}
+
+impl Iterator for PresentCpusIndexIter {
+    type Item = u32;
+
+    fn next(&mut self) -> Option<u32> {
+        let next_cpu_id =
+            unsafe { bindings::cpumask_next(self.index, &bindings::__cpu_present_mask) };
+        // When [`bindings::cpumask_next`] can not find further CPUs set in the
+        // [`bindings::__cpu_present_mask`], it returns a value >= [`bindings::nr_cpu_ids`].
+        if next_cpu_id >= unsafe { bindings::nr_cpu_ids } {
+            return None;
+        }
+        self.index = next_cpu_id as i32;
+        Some(next_cpu_id)
+    }
+}
+
+/// Returns a [`PossibleCpusIndexIter`] that gives the possible CPU indexes.
+///
+/// # Examples
+///
+/// ```
+/// # use kernel::prelude::*;
+/// # use kernel::cpumask::possible_cpus;
+///
+/// fn example() {
+///     // This prints all the possible cpu indexes.
+///     for cpu in possible_cpus(){
+///         pr_info!("{}\n", cpu);
+///     }
+/// }
+/// ```
+pub fn possible_cpus() -> PossibleCpusIndexIter {
+    // Initial index is set to -1. Since [`bindings::cpumask_next`] return the next set bit in a
+    // [`bindings::__cpu_possible_mask`], the CPU index should begins from 0.
+    PossibleCpusIndexIter { index: -1 }
+}
+
+/// Returns a [`OnlineCpusIndexIter`] that gives the online CPU indexes.
+///
+/// # Examples
+///
+/// ```
+/// # use kernel::prelude::*;
+/// # use kernel::cpumask::online_cpus;
+///
+/// fn example() {
+///     // This prints all the online cpu indexes.
+///     for cpu in online_cpus(){
+///         pr_info!("{}\n", cpu);
+///     }
+/// }
+/// ```
+pub fn online_cpus() -> OnlineCpusIndexIter {
+    // Initial index is set to -1. Since [`bindings::cpumask_next`] return the next set bit in a
+    // [`bindings::__cpu_online_mask`], the CPU index should begins from 0.
+    OnlineCpusIndexIter { index: -1 }
+}
+
+/// Returns a [`PresentCpusIndexIter`] that gives the present CPU indexes.
+///
+/// # Examples
+///
+/// ```
+/// # use kernel::prelude::*;
+/// # use kernel::cpumask::present_cpus;
+///
+/// fn example() {
+///     // This prints all the present cpu indexes.
+///     for cpu in present_cpus(){
+///         pr_info!("{}\n", cpu);
+///     }
+/// }
+/// ```
+pub fn present_cpus() -> PresentCpusIndexIter {
+    // Initial index is set to -1. Since [`bindings::cpumask_next`] return the next set bit in a
+    // [`bindings::__cpu_present_mask`], the CPU index should begins from 0.
+    PresentCpusIndexIter { index: -1 }
+}

rust/kernel/lib.rs

@@ -94,6 +94,7 @@ pub mod platform;
 mod types;
 pub mod user_ptr;
 
+pub mod cpumask;
 #[doc(hidden)]
 pub use build_error::build_error;