cpumask: add a function like for_each_online_cpu

This commit add a function return a struct implementing Iterator trait. Thus, this struct can be used as an iterator that returns an online cpu index each time.

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

Author: Richardhongyu

rust/kernel/cpumask.rs

@@ -0,0 +1,52 @@
+// 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 online CPU index iterator.
+///
+/// This iterator has a similar abilitiy to the kernel's macro `for_each_online_cpu`.
+pub struct OnlineCpusIndexIter {
+    index: i32,
+}
+
+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)
+    }
+}
+
+/// 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 {
+    // Inital 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 }
+}

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;