diff --git a/kernel/src/libs/cpumask.rs b/kernel/src/libs/cpumask.rs new file mode 100644 index 000000000..2b9cdacb1 --- /dev/null +++ b/kernel/src/libs/cpumask.rs @@ -0,0 +1,85 @@ +use bitmap::{traits::BitMapOps, AllocBitmap}; + +use crate::mm::percpu::PerCpu; + +pub struct CpuMask { + bmp: AllocBitmap, +} + +#[allow(dead_code)] +impl CpuMask { + pub fn new() -> Self { + let bmp = AllocBitmap::new(PerCpu::MAX_CPU_NUM); + Self { bmp } + } + + /// 获取CpuMask中的第一个cpu + pub fn first(&self) -> Option { + self.bmp.first_index() + } + + /// 获取CpuMask中第一个未被置位的cpu + pub fn first_zero(&self) -> Option { + self.bmp.first_false_index() + } + + /// 获取CpuMask中的最后一个被置位的cpu + pub fn last(&self) -> Option { + self.bmp.last_index() + } + + /// 获取指定cpu之后第一个为1的位的cpu + pub fn next_index(&self, cpu: usize) -> Option { + self.bmp.next_index(cpu) + } + + /// 获取指定cpu之后第一个为未被置位的cpu + pub fn next_zero_index(&self, cpu: usize) -> Option { + self.bmp.next_false_index(cpu) + } + + /// 迭代所有被置位的cpu + pub fn iter_cpu(&self) -> CpuMaskIter { + CpuMaskIter { + mask: self, + index: 0, + set: true, + } + } + + /// 迭代所有未被置位的cpu + pub fn iter_zero_cpu(&self) -> CpuMaskIter { + CpuMaskIter { + mask: self, + index: 0, + set: false, + } + } +} + +pub struct CpuMaskIter<'a> { + mask: &'a CpuMask, + index: usize, + set: bool, +} + +impl<'a> Iterator for CpuMaskIter<'a> { + type Item = usize; + + fn next(&mut self) -> Option { + if self.index == 0 { + if self.set { + self.index = self.mask.first()?; + } else { + self.index = self.mask.first_zero()?; + } + } + + if self.set { + self.index = self.mask.next_index(self.index)?; + } else { + self.index = self.mask.next_zero_index(self.index)?; + } + Some(self.index) + } +} diff --git a/kernel/src/libs/mod.rs b/kernel/src/libs/mod.rs index 8c5c86e0d..7a23aa53c 100644 --- a/kernel/src/libs/mod.rs +++ b/kernel/src/libs/mod.rs @@ -1,5 +1,6 @@ pub mod align; pub mod casting; +pub mod cpumask; pub mod elf; pub mod ffi_convert; #[macro_use]