use crate::types::MAX_CPUS;

const WORDS: usize = MAX_CPUS / 64;

#[derive(Debug, Clone)]
pub struct CpuSet {
    bits: [u64; WORDS],
}

impl CpuSet {
    pub const fn empty() -> Self {
        Self {
            bits: [0u64; WORDS],
        }
    }

    pub fn set(&mut self, cpu: usize) {
        assert!(cpu < MAX_CPUS);
        self.bits[cpu / 64] |= 1u64 << (cpu % 64);
    }

    pub fn clear(&mut self, cpu: usize) {
        assert!(cpu < MAX_CPUS);
        self.bits[cpu / 64] &= !(1u64 << (cpu % 64));
    }

    pub fn test(&self, cpu: usize) -> bool {
        if cpu < MAX_CPUS {
            self.bits[cpu / 64] & (1u64 << (cpu % 64)) != 0
        } else {
            false
        }
    }

    pub fn is_empty(&self) -> bool {
        self.bits.iter().all(|&w| w == 0)
    }

    pub fn count(&self) -> usize {
        self.bits.iter().map(|w| w.count_ones() as usize).sum()
    }

    pub fn iter_set(&self) -> CpuSetIter<'_> {
        CpuSetIter {
            set: self,
            word_idx: 0,
            remaining: self.bits[0],
        }
    }
}

pub struct CpuSetIter<'a> {
    set: &'a CpuSet,
    word_idx: usize,
    remaining: u64,
}

impl Iterator for CpuSetIter<'_> {
    type Item = usize;

    fn next(&mut self) -> Option<usize> {
        loop {
            if self.remaining != 0 {
                let bit = self.remaining.trailing_zeros() as usize;
                self.remaining &= self.remaining - 1;
                return Some(self.word_idx * 64 + bit);
            }
            self.word_idx += 1;
            if self.word_idx >= WORDS {
                return None;
            }
            self.remaining = self.set.bits[self.word_idx];
        }
    }
}