async-ecs/async-ecs/src/misc/bit_iter.rs

use hibitset::BitSetLike;

#[derive(Debug, Clone)]
pub struct BitIter<T> {
    pub(crate) set: T,
    pub(crate) masks: [usize; LAYERS],
    pub(crate) prefix: [u32; LAYERS - 1],
}

impl<T> BitIter<T> {
    pub fn new(set: T, masks: [usize; LAYERS], prefix: [u32; LAYERS - 1]) -> Self {
        Self { set, masks, prefix }
    }
}

impl<T: BitSetLike> BitIter<T> {
    pub fn contains(&self, i: u32) -> bool {
        self.set.contains(i)
    }
}

#[derive(PartialEq)]
pub(crate) enum State {
    Empty,
    Continue,
    Value(u32),
}

impl<T> Iterator for BitIter<T>
where
    T: BitSetLike,
{
    type Item = u32;

    fn next(&mut self) -> Option<Self::Item> {
        use self::State::*;

        'find: loop {
            for level in 0..LAYERS {
                match self.handle_level(level) {
                    Value(v) => return Some(v),
                    Continue => continue 'find,
                    Empty => {}
                }
            }

            return None;
        }
    }
}

impl<T: BitSetLike> BitIter<T> {
    pub(crate) fn handle_level(&mut self, level: usize) -> State {
        use self::State::*;

        if self.masks[level] == 0 {
            Empty
        } else {
            let first_bit = self.masks[level].trailing_zeros();
            self.masks[level] &= !(1 << first_bit);

            let idx = self.prefix.get(level).cloned().unwrap_or(0) | first_bit;

            if level == 0 {
                Value(idx)
            } else {
                self.masks[level - 1] = self.set.get_from_layer(level - 1, idx as usize);
                self.prefix[level - 1] = idx << BITS;

                Continue
            }
        }
    }
}

const LAYERS: usize = 4;

#[cfg(target_pointer_width = "64")]
pub const BITS: usize = 6;

#[cfg(target_pointer_width = "32")]
pub const BITS: usize = 5;

#[cfg(test)]
mod tests {
    use hibitset::{BitSet, BitSetLike};

    #[test]
    fn iterator_clone() {
        let mut set = BitSet::new();

        set.add(1);
        set.add(3);

        let iter = set.iter().skip(1);
        for (a, b) in iter.clone().zip(iter) {
            assert_eq!(a, b);
        }
    }
}