Add an iterator for `Cpumask` making use of C's `cpumask_next`.
Signed-off-by: Mitchell Levy <levymitchell0@gmail.com>
---
rust/helpers/cpumask.c | 5 +++++
rust/kernel/cpumask.rs | 48 +++++++++++++++++++++++++++++++++++++++++++++++-
2 files changed, 52 insertions(+), 1 deletion(-)
diff --git a/rust/helpers/cpumask.c b/rust/helpers/cpumask.c
index eb10598a0242..d95bfa111191 100644
--- a/rust/helpers/cpumask.c
+++ b/rust/helpers/cpumask.c
@@ -42,6 +42,11 @@ bool rust_helper_cpumask_full(struct cpumask *srcp)
return cpumask_full(srcp);
}
+unsigned int rust_helper_cpumask_next(int n, struct cpumask *srcp)
+{
+ return cpumask_next(n, srcp);
+}
+
unsigned int rust_helper_cpumask_weight(struct cpumask *srcp)
{
return cpumask_weight(srcp);
diff --git a/rust/kernel/cpumask.rs b/rust/kernel/cpumask.rs
index 3fcbff438670..b7401848f59e 100644
--- a/rust/kernel/cpumask.rs
+++ b/rust/kernel/cpumask.rs
@@ -6,7 +6,7 @@
use crate::{
alloc::{AllocError, Flags},
- cpu::CpuId,
+ cpu::{self, CpuId},
prelude::*,
types::Opaque,
};
@@ -161,6 +161,52 @@ pub fn copy(&self, dstp: &mut Self) {
}
}
+/// Iterator for a `Cpumask`.
+pub struct CpumaskIter<'a> {
+ mask: &'a Cpumask,
+ last: Option<u32>,
+}
+
+impl<'a> CpumaskIter<'a> {
+ /// Creates a new `CpumaskIter` for the given `Cpumask`.
+ fn new(mask: &'a Cpumask) -> CpumaskIter<'a> {
+ Self { mask, last: None }
+ }
+}
+
+impl<'a> Iterator for CpumaskIter<'a> {
+ type Item = CpuId;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ // SAFETY: By the type invariant, `self.mask.as_raw` is a `struct cpumask *`.
+ let next = unsafe {
+ bindings::cpumask_next(
+ if let Some(last) = self.last {
+ last.try_into().unwrap()
+ } else {
+ -1
+ },
+ self.mask.as_raw(),
+ )
+ };
+
+ if next == cpu::nr_cpu_ids() {
+ None
+ } else {
+ self.last = Some(next);
+ // SAFETY: `cpumask_next` returns either `nr_cpu_ids` or a valid CPU ID.
+ unsafe { Some(CpuId::from_u32_unchecked(next)) }
+ }
+ }
+}
+
+impl Cpumask {
+ /// Returns an iterator over the set bits in the cpumask.
+ pub fn iter(&self) -> CpumaskIter<'_> {
+ CpumaskIter::new(self)
+ }
+}
+
/// A CPU Mask pointer.
///
/// Rust abstraction for the C `struct cpumask_var_t`.
--
2.34.1