xref: /DragonOS/kernel/src/libs/cpumask.rs (revision 731bc2b32d7b37298883d7a15b6dca659b436ee4)
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: None,
69             set: true,
70             begin: true,
71         }
72     }
73 
74     /// 迭代所有未被置位的cpu
75     pub fn iter_zero_cpu(&self) -> CpuMaskIter {
76         CpuMaskIter {
77             mask: self,
78             index: None,
79             set: false,
80             begin: true,
81         }
82     }
83 
84     pub fn inner(&self) -> &AllocBitmap {
85         &self.bmp
86     }
87 }
88 
89 pub struct CpuMaskIter<'a> {
90     mask: &'a CpuMask,
91     index: Option<ProcessorId>,
92     set: bool,
93     begin: bool,
94 }
95 
96 impl<'a> Iterator for CpuMaskIter<'a> {
97     type Item = ProcessorId;
98 
99     fn next(&mut self) -> Option<ProcessorId> {
100         if self.index.is_none() && self.begin {
101             if self.set {
102                 self.index = self.mask.first();
103             } else {
104                 self.index = self.mask.first_zero();
105             }
106 
107             self.begin = false;
108         }
109         let result = self.index;
110         if self.set {
111             self.index = self.mask.next_index(self.index?);
112         } else {
113             self.index = self.mask.next_zero_index(self.index?);
114         }
115 
116         result
117     }
118 }
119 
120 impl core::fmt::Debug for CpuMask {
121     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
122         f.debug_struct("CpuMask")
123             .field("bmp", &format!("size: {}", self.bmp.size()))
124             .finish()
125     }
126 }
127