fastbit 0.11.1

//! Iterator implementation for bit collections.
//!
//! This module provides an efficient iterator that traverses only the set bits (1s)
//! in a bit collection, yielding their indices. The implementation uses bit manipulation
//! techniques to quickly locate set bits and skip over unset regions.

use std::{marker::PhantomData, ptr::NonNull};

use crate::traits::BitWord;

/// An iterator over the set bits in a bit collection.
///
/// This iterator yields the indices of all bits that are set to 1.
/// It efficiently skips over words that contain no set bits and uses
/// bit manipulation to quickly find the next set bit within a word.
///
/// The type parameter `W` represents the word type used for storage
/// (e.g., `u32`, `u64`), which must implement the `BitWord` trait.
/// An iterator over the set bits in a bit collection.
///
/// This iterator yields the indices of all bits that are set to 1.
/// It efficiently skips over words that contain no set bits and uses
/// bit manipulation to quickly find the next set bit within a word.
///
/// The type parameter `W` represents the word type used for storage
/// (e.g., `u32`, `u64`), which must implement the `BitWord` trait.
pub struct Iter<'a, W: BitWord> {
    /// Pointer to the first word in the bit collection.
    ptr: NonNull<W>,
    /// Total number of bits in the collection.
    len: usize,
    /// Current bit position within the iteration.
    word_offset: usize,
    /// Current word being processed.
    curr: W,
    /// Phantom data to tie the lifetime of the iterator to the bit collection.
    _marker: PhantomData<&'a W>,
}

impl<'a, W: BitWord> Iter<'a, W> {
    /// Creates a new iterator from a non-null pointer to the first word and the total bit length.
    ///
    /// This constructor initializes the iterator state, reading the first word if the collection
    /// is not empty.
    ///
    /// # Safety
    ///
    /// The caller must ensure that:
    /// - The pointer is valid for reads of size `len`
    /// - The memory referenced by `ptr` remains valid for the lifetime `'a`
    pub(crate) fn new(ptr: NonNull<W>, len: usize) -> Self {
        let curr = if len == 0 {
            W::from(0)
        } else {
            unsafe { ptr.read() }
        };
        Self {
            ptr,
            len,
            word_offset: 0,
            curr,
            _marker: PhantomData,
        }
    }

    /// Creates a new iterator from raw parts: a pointer to the first word and the total bit length.
    ///
    /// This constructor handles null pointers by using a dangling pointer when appropriate.
    /// It initializes the iterator state, reading the first word if the collection is not empty.
    ///
    /// # Safety
    ///
    /// The caller must ensure that:
    /// - If non-null, the pointer is valid for reads of size `len`
    /// - The memory referenced by `ptr` remains valid for the lifetime `'a`
    pub(crate) fn from_raw_parts(ptr: *const W, len: usize) -> Self {
        let curr = if len == 0 {
            W::from(0)
        } else {
            unsafe { *ptr }
        };
        Self {
            ptr: NonNull::new(ptr as *mut _).unwrap_or(NonNull::dangling()),
            len,
            word_offset: 0,
            curr,
            _marker: PhantomData,
        }
    }
}

/// Implementation of the `Iterator` trait for `Iter<'a, W>`.
///
/// This iterator yields the indices of set bits in the bit collection.
impl<'a, W: BitWord> Iterator for Iter<'a, W> {
    /// The type of items returned by this iterator (indices of set bits).
    type Item = usize;

    /// Advances the iterator and returns the next set bit index.
    ///
    /// This method efficiently finds the next set bit by:
    /// 1. Skipping over words that are all zeros
    /// 2. Using `trailing_zeros()` to find the position of the next set bit
    /// 3. Clearing that bit in the current word to prepare for the next iteration
    ///
    /// Returns `None` when there are no more set bits to iterate over.
    fn next(&mut self) -> Option<Self::Item> {
        while self.word_offset < self.len {
            if self.curr == W::from(0) {
                self.word_offset += W::BITS;
                if self.word_offset >= self.len {
                    break;
                }
                self.curr = unsafe { self.ptr.add(self.word_offset / W::BITS).read() };
                continue;
            }

            let tz = self.curr.trailing_zeros();
            if self.word_offset + tz < self.len {
                self.curr &= !(W::from(1) << tz);
                return Some(self.word_offset + tz);
            } else {
                break;
            }
        }
        None
    }
}

#[cfg(test)]
mod tests {
    use crate::{BitRead, BitVec, traits::BitWrite};

    #[test]
    fn test_iterator_empty() {
        let bv: BitVec<u32> = BitVec::new(0);
        let mut iter = bv.iter();
        assert_eq!(iter.next(), None);
    }

    #[test]
    fn test_iterator_single_bit() {
        let mut bv: BitVec<u64> = BitVec::new(1);
        assert_eq!(bv.iter().collect::<Vec<_>>(), vec![]);
        bv.set(0);
        assert_eq!(bv.iter().collect::<Vec<_>>(), vec![0]);
    }

    #[test]
    fn test_iterator_multiple_bits() {
        let mut bv: BitVec<u8> = BitVec::new(10);
        bv.set(2);
        bv.set(5);
        bv.set(9);
        let bits: Vec<_> = bv.iter().collect();
        assert_eq!(bits, vec![2, 5, 9]);
    }

    #[test]
    fn test_iterator_all_bits_set() {
        let mut bv: BitVec<u16> = BitVec::new(8);
        bv.fill();
        let bits: Vec<_> = bv.iter().collect();
        assert_eq!(bits, (0..8).collect::<Vec<_>>());
    }

    #[test]
    fn test_iterator_after_clear_and_set() {
        let mut bv: BitVec<u32> = BitVec::new(16);
        bv.fill();
        bv.clear();
        assert_eq!(bv.iter().collect::<Vec<_>>(), vec![]);
        bv.set(7);
        bv.set(15);
        assert!(bv.test(7));
        assert!(bv.test(15));
        assert_eq!(bv.iter().collect::<Vec<_>>(), vec![7, 15]);
    }

    #[test]
    fn test_iterator_sparse_bits() {
        let mut bv: BitVec<u64> = BitVec::new(64);
        bv.set(0);
        bv.set(31);
        bv.set(32);
        bv.set(63);
        let bits: Vec<_> = bv.iter().collect();
        assert_eq!(bits, vec![0, 31, 32, 63]);
    }

    #[test]
    fn test_iterator_large_bitvec() {
        let mut bv: BitVec<usize> = BitVec::new(130);
        for i in (0..130).step_by(13) {
            bv.set(i);
        }
        let expected: Vec<_> = (0..130).step_by(13).collect();
        let actual: Vec<_> = bv.iter().collect();
        assert_eq!(actual, expected);
    }

    #[test]
    fn test_iterator_partial_last_word() {
        let mut bv: BitVec<u8> = BitVec::new(70); // 70 bits, not a multiple of 64 (on 64-bit)
        bv.set(0);
        bv.set(63);
        bv.set(64);
        bv.set(69);
        let bits: Vec<_> = bv.iter().collect();
        assert_eq!(bits, vec![0, 63, 64, 69]);
    }
}