xref: /DragonOS/kernel/src/libs/cpumask.rs (revision d2b28acb4d1f160779b25d76afca49ed60ad5d48)
1 use bitmap::{traits::BitMapOps, AllocBitmap};
2 
3 use crate::mm::percpu::PerCpu;
4 
5 pub struct CpuMask {
6     bmp: AllocBitmap,
7 }
8 
9 #[allow(dead_code)]
10 impl CpuMask {
11     pub fn new() -> Self {
12         let bmp = AllocBitmap::new(PerCpu::MAX_CPU_NUM);
13         Self { bmp }
14     }
15 
16     /// 获取CpuMask中的第一个cpu
17     pub fn first(&self) -> Option<usize> {
18         self.bmp.first_index()
19     }
20 
21     /// 获取CpuMask中第一个未被置位的cpu
22     pub fn first_zero(&self) -> Option<usize> {
23         self.bmp.first_false_index()
24     }
25 
26     /// 获取CpuMask中的最后一个被置位的cpu
27     pub fn last(&self) -> Option<usize> {
28         self.bmp.last_index()
29     }
30 
31     /// 获取指定cpu之后第一个为1的位的cpu
32     pub fn next_index(&self, cpu: usize) -> Option<usize> {
33         self.bmp.next_index(cpu)
34     }
35 
36     /// 获取指定cpu之后第一个为未被置位的cpu
37     pub fn next_zero_index(&self, cpu: usize) -> Option<usize> {
38         self.bmp.next_false_index(cpu)
39     }
40 
41     /// 迭代所有被置位的cpu
42     pub fn iter_cpu(&self) -> CpuMaskIter {
43         CpuMaskIter {
44             mask: self,
45             index: 0,
46             set: true,
47         }
48     }
49 
50     /// 迭代所有未被置位的cpu
51     pub fn iter_zero_cpu(&self) -> CpuMaskIter {
52         CpuMaskIter {
53             mask: self,
54             index: 0,
55             set: false,
56         }
57     }
58 }
59 
60 pub struct CpuMaskIter<'a> {
61     mask: &'a CpuMask,
62     index: usize,
63     set: bool,
64 }
65 
66 impl<'a> Iterator for CpuMaskIter<'a> {
67     type Item = usize;
68 
69     fn next(&mut self) -> Option<Self::Item> {
70         if self.index == 0 {
71             if self.set {
72                 self.index = self.mask.first()?;
73             } else {
74                 self.index = self.mask.first_zero()?;
75             }
76         }
77 
78         if self.set {
79             self.index = self.mask.next_index(self.index)?;
80         } else {
81             self.index = self.mask.next_zero_index(self.index)?;
82         }
83         Some(self.index)
84     }
85 }
86