sparsemap 5.6.0

//! Ascending iteration over the set bits of a [`SparseMap`].

use crate::{Chunk, SparseMap, BITS_PER_WORD, CHUNK_BITS, WORDS_PER_CHUNK};
use alloc::collections::btree_map;

/// State for expanding the chunk currently being yielded.
enum Cursor<'a> {
    /// A run yielding `next..=last` (inclusive last, so a run touching
    /// the top of the universe can yield `u64::MAX` without needing an
    /// exclusive end of 2^64).
    Run { next: u64, last: u64, done: bool },
    /// A dense window: `base`, the word index, and the remaining bits
    /// of the current word.
    Dense {
        base: u64,
        word: usize,
        bits: u64,
        words: &'a [u64; WORDS_PER_CHUNK],
    },
}

/// An iterator over the set bits of a [`SparseMap`] in ascending order.
///
/// Created by [`SparseMap::iter`] (and by `&SparseMap`'s
/// [`IntoIterator`]).
#[must_use = "iterators are lazy and do nothing unless consumed"]
pub struct Iter<'a> {
    chunks: btree_map::Iter<'a, u64, Chunk>,
    cur: Option<Cursor<'a>>,
}

impl<'a> Iter<'a> {
    pub(crate) fn new(map: &'a SparseMap) -> Self {
        Iter {
            chunks: map.chunks.iter(),
            cur: None,
        }
    }

    fn load_next_chunk(&mut self) -> bool {
        match self.chunks.next() {
            Some((&base, Chunk::Run(n))) => {
                self.cur = Some(Cursor::Run {
                    next: base,
                    // n*CHUNK_BITS fits in u64 (n < 2^32); the `-1`
                    // groups so `base + span - 1` never reaches 2^64.
                    last: base + (u64::from(*n) * CHUNK_BITS - 1),
                    done: false,
                });
                true
            }
            Some((&base, Chunk::Dense(w))) => {
                self.cur = Some(Cursor::Dense {
                    base,
                    word: 0,
                    bits: w[0],
                    words: w,
                });
                true
            }
            None => false,
        }
    }
}

impl Iterator for Iter<'_> {
    type Item = u64;

    fn next(&mut self) -> Option<u64> {
        loop {
            match &mut self.cur {
                Some(Cursor::Run { next, last, done }) => {
                    if !*done {
                        let v = *next;
                        if *next == *last {
                            // Last bit; stop before `next` would wrap.
                            *done = true;
                        } else {
                            *next += 1;
                        }
                        return Some(v);
                    }
                    self.cur = None;
                }
                Some(Cursor::Dense {
                    base,
                    word,
                    bits,
                    words,
                }) => {
                    while *bits == 0 {
                        *word += 1;
                        if *word >= WORDS_PER_CHUNK {
                            break;
                        }
                        *bits = words[*word];
                    }
                    if *word < WORDS_PER_CHUNK {
                        let tz = bits.trailing_zeros();
                        *bits &= *bits - 1; // clear lowest set bit
                        return Some(*base + *word as u64 * BITS_PER_WORD + u64::from(tz));
                    }
                    self.cur = None;
                }
                None => {
                    if !self.load_next_chunk() {
                        return None;
                    }
                }
            }
        }
    }
}

impl<'a> IntoIterator for &'a SparseMap {
    type Item = u64;
    type IntoIter = Iter<'a>;

    fn into_iter(self) -> Iter<'a> {
        self.iter()
    }
}

impl SparseMap {
    /// Returns an iterator over the set bits, in ascending order.
    ///
    /// ```
    /// use sparsemap::SparseMap;
    /// let m: SparseMap = [3, 1, 4, 1, 5].into_iter().collect();
    /// assert_eq!(m.iter().collect::<Vec<_>>(), vec![1, 3, 4, 5]);
    /// ```
    pub fn iter(&self) -> Iter<'_> {
        Iter::new(self)
    }

    /// Collects the set bits into a `Vec<u64>`, ascending.
    ///
    /// Allocates; provided for parity with the C `sm_to_array`.
    #[must_use]
    pub fn to_vec(&self) -> alloc::vec::Vec<u64> {
        self.iter().collect()
    }

    /// Materializes the index window `[lo, hi)` into a dense
    /// [`bitvec::vec::BitVec`], with bit `i` of the result set iff
    /// `self.contains(lo + i)`.
    ///
    /// `SparseMap` is a *compressed sparse* structure, so it has no
    /// contiguous bit buffer to borrow a `BitSlice` from; this bridge
    /// expands a *bounded* range on demand for code that wants
    /// `bitvec`'s dense slicing.  The result has `hi - lo` bits (empty
    /// if `lo >= hi`).  Choose the window deliberately: a full
    /// `0..u64::MAX` expansion would allocate 2^64 bits.
    ///
    /// Requires the `bitvec` crate feature.
    #[cfg(feature = "bitvec")]
    #[must_use]
    pub fn to_bitvec(&self, lo: u64, hi: u64) -> bitvec::vec::BitVec {
        use bitvec::vec::BitVec;
        if lo >= hi {
            return BitVec::new();
        }
        let len = (hi - lo) as usize;
        let mut bv = BitVec::repeat(false, len);
        // Walk the set bits (sparse) and stamp those inside the window.
        // iter() yields ascending indices, so we can stop once we pass hi.
        for b in self.iter() {
            if b < lo {
                continue;
            }
            if b >= hi {
                break;
            }
            bv.set((b - lo) as usize, true);
        }
        bv
    }
}