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 /// 获取CpuMask中的第一个cpu 20 pub fn first(&self) -> Option<ProcessorId> { 21 self.bmp 22 .first_index() 23 .map(|index| ProcessorId::new(index as u32)) 24 } 25 26 /// 获取CpuMask中第一个未被置位的cpu 27 pub fn first_zero(&self) -> Option<ProcessorId> { 28 self.bmp 29 .first_false_index() 30 .map(|index| ProcessorId::new(index as u32)) 31 } 32 33 /// 获取CpuMask中的最后一个被置位的cpu 34 pub fn last(&self) -> Option<ProcessorId> { 35 self.bmp 36 .last_index() 37 .map(|index| ProcessorId::new(index as u32)) 38 } 39 40 /// 获取指定cpu之后第一个为1的位的cpu 41 pub fn next_index(&self, cpu: ProcessorId) -> Option<ProcessorId> { 42 self.bmp 43 .next_index(cpu.data() as usize) 44 .map(|index| ProcessorId::new(index as u32)) 45 } 46 47 /// 获取指定cpu之后第一个为未被置位的cpu 48 pub fn next_zero_index(&self, cpu: ProcessorId) -> Option<ProcessorId> { 49 self.bmp 50 .next_false_index(cpu.data() as usize) 51 .map(|index| ProcessorId::new(index as u32)) 52 } 53 54 pub fn set(&mut self, cpu: ProcessorId, value: bool) -> Option<bool> { 55 self.bmp.set(cpu.data() as usize, value) 56 } 57 58 pub fn get(&self, cpu: ProcessorId) -> Option<bool> { 59 self.bmp.get(cpu.data() as usize) 60 } 61 62 pub fn is_empty(&self) -> bool { 63 self.bmp.is_empty() 64 } 65 66 /// 迭代所有被置位的cpu 67 pub fn iter_cpu(&self) -> CpuMaskIter { 68 CpuMaskIter { 69 mask: self, 70 index: None, 71 set: true, 72 begin: true, 73 } 74 } 75 76 /// 迭代所有未被置位的cpu 77 pub fn iter_zero_cpu(&self) -> CpuMaskIter { 78 CpuMaskIter { 79 mask: self, 80 index: None, 81 set: false, 82 begin: true, 83 } 84 } 85 86 pub fn inner(&self) -> &AllocBitmap { 87 &self.bmp 88 } 89 } 90 91 impl BitAnd for CpuMask { 92 type Output = Self; 93 94 fn bitand(self, rhs: Self) -> Self::Output { 95 let bmp = self.bmp & rhs.bmp; 96 Self { bmp } 97 } 98 } 99 100 pub struct CpuMaskIter<'a> { 101 mask: &'a CpuMask, 102 index: Option<ProcessorId>, 103 set: bool, 104 begin: bool, 105 } 106 107 impl<'a> Iterator for CpuMaskIter<'a> { 108 type Item = ProcessorId; 109 110 fn next(&mut self) -> Option<ProcessorId> { 111 if self.index.is_none() && self.begin { 112 if self.set { 113 self.index = self.mask.first(); 114 } else { 115 self.index = self.mask.first_zero(); 116 } 117 118 self.begin = false; 119 } 120 let result = self.index; 121 if self.set { 122 self.index = self.mask.next_index(self.index?); 123 } else { 124 self.index = self.mask.next_zero_index(self.index?); 125 } 126 127 result 128 } 129 } 130 131 impl core::fmt::Debug for CpuMask { 132 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { 133 f.debug_struct("CpuMask") 134 .field("bmp", &format!("size: {}", self.bmp.size())) 135 .finish() 136 } 137 } 138