xref: /DragonOS/kernel/src/libs/cpumask.rs (revision da152319797436368304cbc3f85a3b9ec049134b)
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 
85 pub struct CpuMaskIter<'a> {
86     mask: &'a CpuMask,
87     index: Option<ProcessorId>,
88     set: bool,
89     begin: bool,
90 }
91 
92 impl<'a> Iterator for CpuMaskIter<'a> {
93     type Item = ProcessorId;
94 
95     fn next(&mut self) -> Option<ProcessorId> {
96         if self.index.is_none() && self.begin {
97             if self.set {
98                 self.index = self.mask.first();
99             } else {
100                 self.index = self.mask.first_zero();
101             }
102 
103             self.begin = false;
104         }
105         let result = self.index;
106         if self.set {
107             self.index = self.mask.next_index(self.index?);
108         } else {
109             self.index = self.mask.next_zero_index(self.index?);
110         }
111 
112         result
113     }
114 }
115 
116 impl core::fmt::Debug for CpuMask {
117     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
118         f.debug_struct("CpuMask")
119             .field("bmp", &format!("size: {}", self.bmp.size()))
120             .finish()
121     }
122 }
123