// Eryn Wells <eryn@erynwells.me>

macro_rules! bit_scanner {
    ($name:ident) => {
        pub(crate) struct $name {
            bits: u64,
            shift: usize,
        }

        impl $name {
            pub(crate) fn new(bits: u64) -> Self {
                Self { bits, shift: 0 }
            }
        }
    };
}

bit_scanner!(LeadingBitScanner);
bit_scanner!(TrailingBitScanner);

impl Iterator for LeadingBitScanner {
    type Item = usize;

    fn next(&mut self) -> Option<Self::Item> {
        let u64bits = u64::BITS as usize;

        if self.shift == u64bits {
            return None;
        }

        let shifted_bits = self.bits << self.shift;
        let leading_zeros = shifted_bits.leading_zeros() as usize;

        if leading_zeros == u64bits {
            self.shift = leading_zeros;
            return None;
        }

        let position = u64bits - (self.shift + leading_zeros + 1);
        // Shift 1 additional place to account for the 1 that `leading_zeros` found.
        self.shift += leading_zeros + 1;

        Some(position)
    }
}

impl Iterator for TrailingBitScanner {
    type Item = usize;

    fn next(&mut self) -> Option<Self::Item> {
        let u64bits = u64::BITS as usize;

        if self.shift == u64bits {
            return None;
        }

        let shifted_bits = self.bits >> self.shift;
        let trailing_zeros = shifted_bits.trailing_zeros() as usize;

        if trailing_zeros == u64bits {
            self.shift = trailing_zeros;
            return None;
        }

        let position = self.shift + trailing_zeros;
        // Shift 1 additional place to account for the 1 that `leading_zeros` found.
        self.shift += trailing_zeros + 1;

        Some(position)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn leading_zero() {
        let mut scanner = LeadingBitScanner::new(0);
        assert_eq!(scanner.next(), None);
    }

    #[test]
    fn leading_one() {
        let mut scanner = LeadingBitScanner::new(1);
        assert_eq!(scanner.next(), Some(0));
        assert_eq!(scanner.next(), None);
    }

    #[test]
    fn leading_complex() {
        let mut scanner = LeadingBitScanner::new(0b11000101);
        assert_eq!(scanner.next(), Some(7));
        assert_eq!(scanner.next(), Some(6));
        assert_eq!(scanner.next(), Some(2));
        assert_eq!(scanner.next(), Some(0));
        assert_eq!(scanner.next(), None);
    }

    #[test]
    fn trailing_zero() {
        let mut scanner = TrailingBitScanner::new(0);
        assert_eq!(scanner.next(), None);
    }

    #[test]
    fn trailing_one() {
        let mut scanner = TrailingBitScanner::new(1);
        assert_eq!(scanner.next(), Some(0));
        assert_eq!(scanner.next(), None);
    }

    #[test]
    fn trailing_complex() {
        let mut scanner = TrailingBitScanner::new(0b11000101);
        assert_eq!(scanner.next(), Some(0));
        assert_eq!(scanner.next(), Some(2));
        assert_eq!(scanner.next(), Some(6));
        assert_eq!(scanner.next(), Some(7));
        assert_eq!(scanner.next(), None);
    }
}