17d66c313SLoGin use bitmap::{traits::BitMapOps, AllocBitmap}; 27d66c313SLoGin 3e2841179SLoGin use crate::{mm::percpu::PerCpu, smp::cpu::ProcessorId}; 47d66c313SLoGin 5e2841179SLoGin #[derive(Clone)] 67d66c313SLoGin pub struct CpuMask { 77d66c313SLoGin bmp: AllocBitmap, 87d66c313SLoGin } 97d66c313SLoGin 107d66c313SLoGin #[allow(dead_code)] 117d66c313SLoGin impl CpuMask { 127d66c313SLoGin pub fn new() -> Self { 13e2841179SLoGin let bmp = AllocBitmap::new(PerCpu::MAX_CPU_NUM as usize); 147d66c313SLoGin Self { bmp } 157d66c313SLoGin } 167d66c313SLoGin 177d66c313SLoGin /// 获取CpuMask中的第一个cpu 18e2841179SLoGin pub fn first(&self) -> Option<ProcessorId> { 19e2841179SLoGin self.bmp 20e2841179SLoGin .first_index() 21e2841179SLoGin .map(|index| ProcessorId::new(index as u32)) 227d66c313SLoGin } 237d66c313SLoGin 247d66c313SLoGin /// 获取CpuMask中第一个未被置位的cpu 25e2841179SLoGin pub fn first_zero(&self) -> Option<ProcessorId> { 26e2841179SLoGin self.bmp 27e2841179SLoGin .first_false_index() 28e2841179SLoGin .map(|index| ProcessorId::new(index as u32)) 297d66c313SLoGin } 307d66c313SLoGin 317d66c313SLoGin /// 获取CpuMask中的最后一个被置位的cpu 32e2841179SLoGin pub fn last(&self) -> Option<ProcessorId> { 33e2841179SLoGin self.bmp 34e2841179SLoGin .last_index() 35e2841179SLoGin .map(|index| ProcessorId::new(index as u32)) 367d66c313SLoGin } 377d66c313SLoGin 387d66c313SLoGin /// 获取指定cpu之后第一个为1的位的cpu 39e2841179SLoGin pub fn next_index(&self, cpu: ProcessorId) -> Option<ProcessorId> { 40e2841179SLoGin self.bmp 41e2841179SLoGin .next_index(cpu.data() as usize) 42e2841179SLoGin .map(|index| ProcessorId::new(index as u32)) 437d66c313SLoGin } 447d66c313SLoGin 457d66c313SLoGin /// 获取指定cpu之后第一个为未被置位的cpu 46e2841179SLoGin pub fn next_zero_index(&self, cpu: ProcessorId) -> Option<ProcessorId> { 47e2841179SLoGin self.bmp 48e2841179SLoGin .next_false_index(cpu.data() as usize) 49e2841179SLoGin .map(|index| ProcessorId::new(index as u32)) 50e2841179SLoGin } 51e2841179SLoGin 52e2841179SLoGin pub fn set(&mut self, cpu: ProcessorId, value: bool) -> Option<bool> { 53e2841179SLoGin self.bmp.set(cpu.data() as usize, value) 54e2841179SLoGin } 55e2841179SLoGin 56e2841179SLoGin pub fn get(&self, cpu: ProcessorId) -> Option<bool> { 57e2841179SLoGin self.bmp.get(cpu.data() as usize) 58e2841179SLoGin } 59e2841179SLoGin 60e2841179SLoGin pub fn is_empty(&self) -> bool { 61e2841179SLoGin self.bmp.is_empty() 627d66c313SLoGin } 637d66c313SLoGin 647d66c313SLoGin /// 迭代所有被置位的cpu 657d66c313SLoGin pub fn iter_cpu(&self) -> CpuMaskIter { 667d66c313SLoGin CpuMaskIter { 677d66c313SLoGin mask: self, 688cb2e9b3SLoGin index: None, 697d66c313SLoGin set: true, 708cb2e9b3SLoGin begin: true, 717d66c313SLoGin } 727d66c313SLoGin } 737d66c313SLoGin 747d66c313SLoGin /// 迭代所有未被置位的cpu 757d66c313SLoGin pub fn iter_zero_cpu(&self) -> CpuMaskIter { 767d66c313SLoGin CpuMaskIter { 777d66c313SLoGin mask: self, 788cb2e9b3SLoGin index: None, 797d66c313SLoGin set: false, 808cb2e9b3SLoGin begin: true, 817d66c313SLoGin } 827d66c313SLoGin } 83*9430523bSyuyi2439 84*9430523bSyuyi2439 pub fn inner(&self) -> &AllocBitmap { 85*9430523bSyuyi2439 &self.bmp 86*9430523bSyuyi2439 } 877d66c313SLoGin } 887d66c313SLoGin 897d66c313SLoGin pub struct CpuMaskIter<'a> { 907d66c313SLoGin mask: &'a CpuMask, 918cb2e9b3SLoGin index: Option<ProcessorId>, 927d66c313SLoGin set: bool, 938cb2e9b3SLoGin begin: bool, 947d66c313SLoGin } 957d66c313SLoGin 967d66c313SLoGin impl<'a> Iterator for CpuMaskIter<'a> { 97e2841179SLoGin type Item = ProcessorId; 987d66c313SLoGin 99e2841179SLoGin fn next(&mut self) -> Option<ProcessorId> { 1008cb2e9b3SLoGin if self.index.is_none() && self.begin { 1017d66c313SLoGin if self.set { 1028cb2e9b3SLoGin self.index = self.mask.first(); 1037d66c313SLoGin } else { 1048cb2e9b3SLoGin self.index = self.mask.first_zero(); 1057d66c313SLoGin } 1067d66c313SLoGin 1078cb2e9b3SLoGin self.begin = false; 1087d66c313SLoGin } 1098cb2e9b3SLoGin let result = self.index; 1108cb2e9b3SLoGin if self.set { 1118cb2e9b3SLoGin self.index = self.mask.next_index(self.index?); 1128cb2e9b3SLoGin } else { 1138cb2e9b3SLoGin self.index = self.mask.next_zero_index(self.index?); 1148cb2e9b3SLoGin } 1158cb2e9b3SLoGin 1168cb2e9b3SLoGin result 1177d66c313SLoGin } 1187d66c313SLoGin } 119e2841179SLoGin 120e2841179SLoGin impl core::fmt::Debug for CpuMask { 121e2841179SLoGin fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { 122e2841179SLoGin f.debug_struct("CpuMask") 123e2841179SLoGin .field("bmp", &format!("size: {}", self.bmp.size())) 124e2841179SLoGin .finish() 125e2841179SLoGin } 126e2841179SLoGin } 127