lewis.moe
Nothing to see here, move along
1use crate::types::MAX_CPUS;
2
3const WORDS: usize = MAX_CPUS / 64;
4
5#[derive(Debug, Clone)]
6pub struct CpuSet {
7 bits: [u64; WORDS],
8}
9
10impl CpuSet {
11 pub const fn empty() -> Self {
12 Self {
13 bits: [0u64; WORDS],
14 }
15 }
16
17 pub fn set(&mut self, cpu: usize) {
18 assert!(cpu < MAX_CPUS);
19 self.bits[cpu / 64] |= 1u64 << (cpu % 64);
20 }
21
22 pub fn clear(&mut self, cpu: usize) {
23 assert!(cpu < MAX_CPUS);
24 self.bits[cpu / 64] &= !(1u64 << (cpu % 64));
25 }
26
27 pub fn test(&self, cpu: usize) -> bool {
28 if cpu < MAX_CPUS {
29 self.bits[cpu / 64] & (1u64 << (cpu % 64)) != 0
30 } else {
31 false
32 }
33 }
34
35 pub fn is_empty(&self) -> bool {
36 self.bits.iter().all(|&w| w == 0)
37 }
38
39 pub fn count(&self) -> usize {
40 self.bits.iter().map(|w| w.count_ones() as usize).sum()
41 }
42
43 pub fn iter_set(&self) -> CpuSetIter<'_> {
44 CpuSetIter {
45 set: self,
46 word_idx: 0,
47 remaining: self.bits[0],
48 }
49 }
50}
51
52pub struct CpuSetIter<'a> {
53 set: &'a CpuSet,
54 word_idx: usize,
55 remaining: u64,
56}
57
58impl Iterator for CpuSetIter<'_> {
59 type Item = usize;
60
61 fn next(&mut self) -> Option<usize> {
62 loop {
63 if self.remaining != 0 {
64 let bit = self.remaining.trailing_zeros() as usize;
65 self.remaining &= self.remaining - 1;
66 return Some(self.word_idx * 64 + bit);
67 }
68 self.word_idx += 1;
69 if self.word_idx >= WORDS {
70 return None;
71 }
72 self.remaining = self.set.bits[self.word_idx];
73 }
74 }
75}