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