xref: /DragonOS/kernel/src/libs/cpumask.rs (revision 9a0802fd2ddda39e96342997abbfc30bf65f1f0e)
1 use core::ops::BitAnd;
2 
3 use bitmap::{traits::BitMapOps, AllocBitmap};
4 
5 use crate::{mm::percpu::PerCpu, smp::cpu::ProcessorId};
6 
7 #[derive(Clone)]
8 pub struct CpuMask {
9     bmp: AllocBitmap,
10 }
11 
12 #[allow(dead_code)]
13 impl CpuMask {
14     pub fn new() -> Self {
15         let bmp = AllocBitmap::new(PerCpu::MAX_CPU_NUM as usize);
16         Self { bmp }
17     }
18 
19     /// # from_cpu - 从指定的CPU创建CPU掩码
20     ///
21     /// 该函数用于根据给定的CPU标识创建一个CPU掩码,只有指定的CPU被设置为激活状态。
22     ///
23     /// ## 参数
24     /// - `cpu`: `ProcessorId`,指定要设置为激活状态的CPU。
25     ///
26     /// ## 返回值
27     /// - `Self`: 返回一个新的`CpuMask`实例,其中只有指定的CPU被设置为激活状态。
28     pub fn from_cpu(cpu: ProcessorId) -> Self {
29         let mut mask = Self::new();
30         mask.set(cpu, true);
31         mask
32     }
33 
34     /// 获取CpuMask中的第一个cpu
35     pub fn first(&self) -> Option<ProcessorId> {
36         self.bmp
37             .first_index()
38             .map(|index| ProcessorId::new(index as u32))
39     }
40 
41     /// 获取CpuMask中第一个未被置位的cpu
42     pub fn first_zero(&self) -> Option<ProcessorId> {
43         self.bmp
44             .first_false_index()
45             .map(|index| ProcessorId::new(index as u32))
46     }
47 
48     /// 获取CpuMask中的最后一个被置位的cpu
49     pub fn last(&self) -> Option<ProcessorId> {
50         self.bmp
51             .last_index()
52             .map(|index| ProcessorId::new(index as u32))
53     }
54 
55     /// 获取指定cpu之后第一个为1的位的cpu
56     pub fn next_index(&self, cpu: ProcessorId) -> Option<ProcessorId> {
57         self.bmp
58             .next_index(cpu.data() as usize)
59             .map(|index| ProcessorId::new(index as u32))
60     }
61 
62     /// 获取指定cpu之后第一个为未被置位的cpu
63     pub fn next_zero_index(&self, cpu: ProcessorId) -> Option<ProcessorId> {
64         self.bmp
65             .next_false_index(cpu.data() as usize)
66             .map(|index| ProcessorId::new(index as u32))
67     }
68 
69     pub fn set(&mut self, cpu: ProcessorId, value: bool) -> Option<bool> {
70         self.bmp.set(cpu.data() as usize, value)
71     }
72 
73     pub fn get(&self, cpu: ProcessorId) -> Option<bool> {
74         self.bmp.get(cpu.data() as usize)
75     }
76 
77     pub fn is_empty(&self) -> bool {
78         self.bmp.is_empty()
79     }
80 
81     /// 迭代所有被置位的cpu
82     pub fn iter_cpu(&self) -> CpuMaskIter {
83         CpuMaskIter {
84             mask: self,
85             index: None,
86             set: true,
87             begin: true,
88         }
89     }
90 
91     /// 迭代所有未被置位的cpu
92     pub fn iter_zero_cpu(&self) -> CpuMaskIter {
93         CpuMaskIter {
94             mask: self,
95             index: None,
96             set: false,
97             begin: true,
98         }
99     }
100 
101     pub fn inner(&self) -> &AllocBitmap {
102         &self.bmp
103     }
104 
105     pub fn bitand_assign(&mut self, rhs: &CpuMask) {
106         self.bmp.bitand_assign(&rhs.bmp);
107     }
108 }
109 
110 impl BitAnd for &CpuMask {
111     type Output = CpuMask;
112 
113     fn bitand(self, rhs: &CpuMask) -> Self::Output {
114         let bmp = &self.bmp & &rhs.bmp;
115         CpuMask { bmp }
116     }
117 }
118 
119 impl Default for CpuMask {
120     fn default() -> Self {
121         Self::new()
122     }
123 }
124 
125 pub struct CpuMaskIter<'a> {
126     mask: &'a CpuMask,
127     index: Option<ProcessorId>,
128     set: bool,
129     begin: bool,
130 }
131 
132 impl<'a> Iterator for CpuMaskIter<'a> {
133     type Item = ProcessorId;
134 
135     fn next(&mut self) -> Option<ProcessorId> {
136         if self.index.is_none() && self.begin {
137             if self.set {
138                 self.index = self.mask.first();
139             } else {
140                 self.index = self.mask.first_zero();
141             }
142 
143             self.begin = false;
144         }
145         let result = self.index;
146         if self.set {
147             self.index = self.mask.next_index(self.index?);
148         } else {
149             self.index = self.mask.next_zero_index(self.index?);
150         }
151 
152         result
153     }
154 }
155 
156 impl core::fmt::Debug for CpuMask {
157     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
158         f.debug_struct("CpuMask")
159             .field("bmp", &format!("size: {}", self.bmp.size()))
160             .finish()
161     }
162 }
163