xref: /DragonOS/kernel/src/libs/cpumask.rs (revision 338f6903262c5031abad3c8e361813355a27fcdb)
1 use bitmap::{traits::BitMapOps, AllocBitmap};
2 
3 use crate::{mm::percpu::PerCpu, smp::cpu::ProcessorId};
4 
5 #[derive(Clone)]
6 pub struct CpuMask {
7     bmp: AllocBitmap,
8 }
9 
10 #[allow(dead_code)]
11 impl CpuMask {
12     pub fn new() -> Self {
13         let bmp = AllocBitmap::new(PerCpu::MAX_CPU_NUM as usize);
14         Self { bmp }
15     }
16 
17     /// 获取CpuMask中的第一个cpu
18     pub fn first(&self) -> Option<ProcessorId> {
19         self.bmp
20             .first_index()
21             .map(|index| ProcessorId::new(index as u32))
22     }
23 
24     /// 获取CpuMask中第一个未被置位的cpu
25     pub fn first_zero(&self) -> Option<ProcessorId> {
26         self.bmp
27             .first_false_index()
28             .map(|index| ProcessorId::new(index as u32))
29     }
30 
31     /// 获取CpuMask中的最后一个被置位的cpu
32     pub fn last(&self) -> Option<ProcessorId> {
33         self.bmp
34             .last_index()
35             .map(|index| ProcessorId::new(index as u32))
36     }
37 
38     /// 获取指定cpu之后第一个为1的位的cpu
39     pub fn next_index(&self, cpu: ProcessorId) -> Option<ProcessorId> {
40         self.bmp
41             .next_index(cpu.data() as usize)
42             .map(|index| ProcessorId::new(index as u32))
43     }
44 
45     /// 获取指定cpu之后第一个为未被置位的cpu
46     pub fn next_zero_index(&self, cpu: ProcessorId) -> Option<ProcessorId> {
47         self.bmp
48             .next_false_index(cpu.data() as usize)
49             .map(|index| ProcessorId::new(index as u32))
50     }
51 
52     pub fn set(&mut self, cpu: ProcessorId, value: bool) -> Option<bool> {
53         self.bmp.set(cpu.data() as usize, value)
54     }
55 
56     pub fn get(&self, cpu: ProcessorId) -> Option<bool> {
57         self.bmp.get(cpu.data() as usize)
58     }
59 
60     pub fn is_empty(&self) -> bool {
61         self.bmp.is_empty()
62     }
63 
64     /// 迭代所有被置位的cpu
65     pub fn iter_cpu(&self) -> CpuMaskIter {
66         CpuMaskIter {
67             mask: self,
68             index: ProcessorId::new(0),
69             set: true,
70         }
71     }
72 
73     /// 迭代所有未被置位的cpu
74     pub fn iter_zero_cpu(&self) -> CpuMaskIter {
75         CpuMaskIter {
76             mask: self,
77             index: ProcessorId::new(0),
78             set: false,
79         }
80     }
81 }
82 
83 pub struct CpuMaskIter<'a> {
84     mask: &'a CpuMask,
85     index: ProcessorId,
86     set: bool,
87 }
88 
89 impl<'a> Iterator for CpuMaskIter<'a> {
90     type Item = ProcessorId;
91 
92     fn next(&mut self) -> Option<ProcessorId> {
93         if self.index.data() == 0 {
94             if self.set {
95                 self.index = self.mask.first()?;
96             } else {
97                 self.index = self.mask.first_zero()?;
98             }
99         }
100 
101         if self.set {
102             self.index = self.mask.next_index(self.index)?;
103         } else {
104             self.index = self.mask.next_zero_index(self.index)?;
105         }
106         Some(self.index)
107     }
108 }
109 
110 impl core::fmt::Debug for CpuMask {
111     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
112         f.debug_struct("CpuMask")
113             .field("bmp", &format!("size: {}", self.bmp.size()))
114             .finish()
115     }
116 }
117