bmap 0.3.0

//! Container of bits adressable by an index.

use std::{cmp, fmt};

mod err;

pub use err::Error;

/// Number of bits per word in the bits vector.
const NUM_WORD_BITS: usize = usize::BITS as usize;


/// A compact list of bits.
pub struct CountedBitmap {
  bits: Vec<usize>,

  /// Number of bits in bit map.
  length: usize,

  /// Number of remaining bits.  This is functionally equivalent to "total
  /// number of zeroes in the container".
  remain: usize
}

impl fmt::Debug for CountedBitmap {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    let blocks: Vec<String> = self
      .iter_ones_block()
      .map(|(first, last)| {
        if first == last {
          format!("{first}")
        } else {
          format!("{first}-{last}")
        }
      })
      .collect();
    let set_blocks = blocks.join(",");
    write!(
      f,
      "BMap {{ length: {}, remain: {}, [{}] }}",
      self.length, self.remain, set_blocks
    )
  }
}


impl CountedBitmap {
  /// Create a new bit vector which can hold `bitcount` number of bits.
  ///
  /// All bits will be initialized to `0`.
  #[must_use]
  pub fn new(bitcount: usize) -> Self {
    let align = NUM_WORD_BITS; // number of bits in bitvec slot
    let num = bitcount;
    let len = num.div_ceil(align);

    let bits = vec![0; len];
    Self {
      bits,
      length: bitcount,
      remain: bitcount
    }
  }

  /// Return the number of bits in bitmap.
  ///
  /// ```
  /// use bmap::CountedBitmap;
  ///
  /// let bmap = CountedBitmap::new(42);
  /// assert_eq!(bmap.len(), 42);
  /// ```
  #[inline]
  #[must_use]
  pub const fn len(&self) -> usize {
    self.length
  }

  /// Change the length of the bitmap.
  ///
  /// Automatically adjust the internal counter as appropriate.
  ///
  /// If the length change is immediately followed by a call to
  /// [`CountedBitmap::clear()`], then it's better to use
  /// [`CountedBitmap::set_len_clear()`], which will change the length and
  /// clear the bitmap in a single call (and will generally be faster).
  ///
  /// # Shrink example
  /// ```
  /// use bmap::CountedBitmap;
  ///
  /// let mut bmap = CountedBitmap::new(2);
  /// bmap.set(0).unwrap();
  /// assert_eq!(bmap.remain(), 1);
  ///
  /// bmap.set_len(1);
  /// assert_eq!(bmap.remain(), 0);
  /// assert!(bmap.is_finished());
  /// ```
  ///
  /// # Grow example
  /// ```
  /// use bmap::CountedBitmap;
  ///
  /// let mut bmap = CountedBitmap::new(2);
  /// bmap.set(0).unwrap();
  /// bmap.set(1).unwrap();
  /// assert_eq!(bmap.remain(), 0);
  /// assert!(bmap.is_finished());  // Completed
  ///
  /// bmap.set_len(4);
  /// assert_eq!(bmap.remain(), 2);  // No longer completed after growth
  /// assert!(!bmap.is_finished());
  /// ```
  ///
  /// # Implementation details
  /// If the bitmap is shunk, only the length of the internal storage is
  /// changed (its _capacity_  is unchanged).
  pub fn set_len(&mut self, bitcount: usize) {
    if bitcount > self.length {
      // The bitmap grew

      let new_len = bitcount.div_ceil(NUM_WORD_BITS);
      if new_len > self.bits.len() {
        self.bits.resize(new_len, 0);
      }

      let added = bitcount - self.length;
      self.remain += added;

      self.length = bitcount;
    } else if bitcount < self.length {
      // The bitmap shrank

      let new_len = bitcount.div_ceil(NUM_WORD_BITS);

      // Reduce `remain` for every zero bit being removed
      for idx in bitcount..self.length {
        if !self.is_set_unchecked(idx) {
          self.remain -= 1;
        }
      }

      // Resize the length of the internal storage. (This changed the
      // length, not the capacity, of the Vec).
      self.bits.truncate(new_len);

      // Clear trailing bits
      if bitcount > 0 {
        let last_word_idx = new_len - 1;
        let bits_in_last_word = bitcount % NUM_WORD_BITS;
        if bits_in_last_word != 0 {
          let mask = (1 << bits_in_last_word) - 1;
          self.bits[last_word_idx] &= mask;
        }
      }

      self.length = bitcount;
    }
  }

  /// Change length of bitmap and clear it.
  ///
  /// # Grow example
  /// ```
  /// use bmap::CountedBitmap;
  ///
  /// let mut bmap = CountedBitmap::new(2);
  /// bmap.set(0).unwrap();
  /// assert_eq!(bmap.remain(), 1);
  ///
  /// bmap.set_len_clear(4);
  /// assert_eq!(bmap.remain(), 4);
  /// ```
  ///
  /// # Shrink example
  /// ```
  /// use bmap::CountedBitmap;
  ///
  /// let mut bmap = CountedBitmap::new(4);
  /// bmap.set(0).unwrap();
  /// bmap.set(3).unwrap();
  /// assert_eq!(bmap.remain(), 2);
  ///
  /// bmap.set_len_clear(2);
  /// assert_eq!(bmap.remain(), 2);
  /// ```
  pub fn set_len_clear(&mut self, bitcount: usize) {
    match bitcount.cmp(&self.length) {
      cmp::Ordering::Less => {
        // First shrink buffer ..
        let new_len = bitcount.div_ceil(NUM_WORD_BITS);
        self.bits.truncate(new_len);

        // .. then clear it
        self.bits.as_mut_slice().fill(0);
      }
      cmp::Ordering::Equal => {
        // Length unchganged -- just clear everything
        self.bits.as_mut_slice().fill(0);
      }
      cmp::Ordering::Greater => {
        // First clear existing ..
        self.bits.as_mut_slice().fill(0);

        // .. then grow, with zeroes
        let new_len = bitcount.div_ceil(NUM_WORD_BITS);
        self.bits.resize(new_len, 0);
      }
    }

    self.length = bitcount;

    // Reset count
    self.remain = self.length;
  }

  /// Returns true if no bits have been set.
  ///
  /// ```
  /// use bmap::CountedBitmap;
  ///
  /// let mut bmap = CountedBitmap::new(42);
  /// assert!(bmap.is_empty());
  ///
  /// bmap.set(11);
  /// assert!(!bmap.is_empty());
  ///
  /// // Special case -- a zero-length bitmap is considered to be empty
  /// let bmap = CountedBitmap::new(0);
  /// assert!(bmap.is_empty());
  /// ```
  #[must_use]
  pub const fn is_empty(&self) -> bool {
    self.remain == self.length
  }

  /// Return number of bits that have not been set.  (I.e. number of bits
  /// remaining to be set).
  ///
  /// ```
  /// use bmap::CountedBitmap;
  ///
  /// let mut bmap = CountedBitmap::new(4);
  /// assert_eq!(bmap.remain(), 4);
  /// bmap.set(1);
  /// assert_eq!(bmap.remain(), 3);
  /// ```
  #[inline]
  #[must_use]
  pub const fn remain(&self) -> usize {
    self.remain
  }

  /// Return how many many bits are set to 1 in relative terms.  The returned
  /// value will be a value between 0.0 and 1.0 (inclusive-inclusive).  For a
  /// 0-length bit vector the value 1.0 will be returned.
  #[must_use]
  #[expect(clippy::cast_precision_loss)]
  pub fn progress(&self) -> f64 {
    if self.is_empty() {
      1.0
    } else {
      let done = (self.length - self.remain) as f64;
      done / self.len() as f64
    }
  }

  /// Return a boolean indicating whether all bits have been set.
  ///
  /// ```
  /// use bmap::CountedBitmap;
  ///
  /// let mut bmap = CountedBitmap::new(2);
  /// assert!(bmap.is_finished() == false);
  /// bmap.set(0);
  /// assert!(bmap.is_finished() == false);
  /// bmap.set(1);
  /// assert!(bmap.is_finished() == true);
  ///
  /// // Special case: zero-length bitmap is always finished.
  /// let mut bmap = CountedBitmap::new(0);
  /// assert!(bmap.is_finished() == true);
  /// ```
  #[inline]
  #[must_use]
  pub const fn is_finished(&self) -> bool {
    self.remain == 0
  }

  /// Given a bit index, return `true` is the corresponding bit is set.  Return
  /// `false` otherwise.
  ///
  /// ```
  /// use bmap::CountedBitmap;
  ///
  /// let mut bmap = CountedBitmap::new(2);
  /// assert_eq!(bmap.is_set(0).unwrap(), false);
  /// bmap.set(0).unwrap();
  /// assert_eq!(bmap.is_set(0).unwrap(), true);
  /// ```
  ///
  /// # Errors
  /// If `idx` is out of bounds then [`Error::OutOfBounds`] is returned.
  #[inline]
  pub fn is_set(&self, idx: usize) -> Result<bool, Error> {
    let (iword, bitval) = self.get_bidx(idx)?;

    // SAFETY: The index has been validated already.
    let v = unsafe { self.bits.get_unchecked(iword) };

    Ok(v & bitval != 0)
  }

  /// Set a bit.
  ///
  /// ```
  /// use bmap::CountedBitmap;
  ///
  /// let mut bmap = CountedBitmap::new(4);
  /// bmap.set(2).unwrap();
  /// ```
  ///
  /// # Errors
  /// If `idx` is out of bounds then [`Error::OutOfBounds`] is returned.
  #[inline]
  pub fn set(&mut self, idx: usize) -> Result<(), Error> {
    let (iword, bitval) = self.get_bidx(idx)?;

    // SAFETY: The index has been validated already.
    let v = unsafe { self.bits.get_unchecked_mut(iword) };
    if *v & bitval == 0 {
      *v |= bitval;
      self.remain -= 1;
    }
    Ok(())
  }

  /// Set a bit, where it is assumed that the caller has validated that the
  /// `idx` parameter is valid.
  ///
  /// ```
  /// use bmap::CountedBitmap;
  ///
  /// let mut bmap = CountedBitmap::new(4);
  /// unsafe { bmap.set_unchecked(2) };
  /// assert_eq!(bmap.is_set(2).unwrap(), true);
  /// ```
  ///
  /// # Safety
  /// The caller must ensure that `idx` is within bounds.
  #[inline]
  pub unsafe fn set_unchecked(&mut self, idx: usize) {
    let (iword, bitval) = Self::get_bidx_unchecked(idx);
    let v = unsafe { self.bits.get_unchecked_mut(iword) };
    if *v & bitval == 0 {
      *v |= bitval;
      self.remain -= 1;
    }
  }

  /// Clear a bit.
  ///
  /// ```
  /// use bmap::CountedBitmap;
  ///
  /// let mut bmap = CountedBitmap::new(4);
  /// bmap.set(2).unwrap();
  /// assert_eq!(bmap.is_set(2).unwrap(), true);
  /// bmap.unset(2).unwrap();
  /// assert_eq!(bmap.is_set(2).unwrap(), false);
  /// ```
  ///
  /// # Errors
  /// If `idx` is out of bounds then [`Error::OutOfBounds`] is returned.
  #[inline]
  pub fn unset(&mut self, idx: usize) -> Result<(), Error> {
    let (iword, bitval) = self.get_bidx(idx)?;

    // SAFETY: The index has been validated already.
    let v = unsafe { self.bits.get_unchecked_mut(iword) };
    if *v & bitval != 0 {
      *v &= !bitval;
      self.remain += 1;
    }
    Ok(())
  }

  /// Clear the bitmap.
  ///
  /// Sets all bits to zero and resets the remaining counter.
  pub fn clear(&mut self) {
    self.bits.as_mut_slice().fill(0);
    self.remain = self.length;
  }

  /// If bit at the specified index is not set, then call closure.  If closure
  /// returns `true` then set bit.
  ///
  /// ```
  /// use bmap::CountedBitmap;
  ///
  /// let mut bmap = CountedBitmap::new(1);
  /// let mut set_to_true = false;
  /// let ret = bmap
  ///   .cond_set(0, || {
  ///     set_to_true = true;
  ///     true
  ///   })
  ///   .unwrap();
  /// assert_eq!(ret, true);
  /// assert_eq!(bmap.is_finished(), true);
  /// assert_eq!(set_to_true, true);
  /// ```
  ///
  /// # Errors
  /// If `idx` is out of bounds then [`Error::OutOfBounds`] is returned.
  pub fn cond_set<F>(&mut self, idx: usize, mut f: F) -> Result<bool, Error>
  where
    F: FnMut() -> bool
  {
    let (iword, bitval) = self.get_bidx(idx)?;

    // SAFETY: The index has already been validated.
    let v = unsafe { self.bits.get_unchecked_mut(iword) };
    if *v & bitval == 0 && f() {
      *v |= bitval;
      self.remain -= 1;
      Ok(true)
    } else {
      Ok(false)
    }
  }
}


/// Iterators.
impl CountedBitmap {
  /// Return an iterator that returns each of the container's bit values as
  /// booleans.
  ///
  /// ```
  /// use bmap::CountedBitmap;
  /// let mut bmap = CountedBitmap::new(4);
  /// bmap.set(1).unwrap();
  /// bmap.set(2).unwrap();
  /// let mut it = bmap.iter().enumerate();
  /// assert_eq!(it.next(), Some((0, false)));
  /// assert_eq!(it.next(), Some((1, true)));
  /// assert_eq!(it.next(), Some((2, true)));
  /// assert_eq!(it.next(), Some((3, false)));
  /// assert_eq!(it.next(), None);
  /// ```
  #[must_use]
  pub const fn iter(&self) -> BitIter<'_> {
    BitIter { bmap: self, idx: 0 }
  }

  /// Create an iterator that will return the indexes of all zeroes in the bit
  /// map.
  ///
  /// ```
  /// use bmap::CountedBitmap;
  /// let mut bmap = CountedBitmap::new(4);
  /// bmap.set(1).unwrap();
  /// bmap.set(2).unwrap();
  /// let mut it = bmap.iter_zeroes();
  /// assert_eq!(it.next(), Some(0));
  /// assert_eq!(it.next(), Some(3));
  /// assert_eq!(it.next(), None);
  /// ```
  #[must_use]
  pub const fn iter_zeroes(&self) -> BitValIter<'_> {
    BitValIter {
      bmap: self,
      idx: 0,
      set: false
    }
  }

  /// Create an iterator that will return the indexes of all ones in the bit
  /// map.
  ///
  /// ```
  /// use bmap::CountedBitmap;
  /// let mut bmap = CountedBitmap::new(4);
  /// bmap.set(1).unwrap();
  /// bmap.set(2).unwrap();
  /// let mut it = bmap.iter_ones();
  /// assert_eq!(it.next(), Some(1));
  /// assert_eq!(it.next(), Some(2));
  /// assert_eq!(it.next(), None);
  /// ```
  #[must_use]
  pub const fn iter_ones(&self) -> BitValIter<'_> {
    BitValIter {
      bmap: self,
      idx: 0,
      set: true
    }
  }

  /// Create an iterator that will return index ranges of contiguous blocks of
  /// zeroes.
  ///
  /// ```
  /// use bmap::CountedBitmap;
  /// let mut bmap = CountedBitmap::new(4);
  /// bmap.set(0).unwrap();
  /// bmap.set(1).unwrap();
  /// let mut it = bmap.iter_zeroes_block();
  /// assert_eq!(it.next(), Some((2, 3)));
  /// assert_eq!(it.next(), None);
  /// ```
  #[must_use]
  pub const fn iter_zeroes_block(&self) -> BitBlocksIter<'_> {
    BitBlocksIter {
      bmap: self,
      idx: 0,
      set: false
    }
  }

  /// Create an iterator that will return index ranges of contiguous blocks of
  /// ones.
  ///
  ///
  /// ```
  /// use bmap::CountedBitmap;
  /// let mut bmap = CountedBitmap::new(4);
  /// bmap.set(0).unwrap();
  /// bmap.set(1).unwrap();
  /// let mut it = bmap.iter_ones_block();
  /// assert_eq!(it.next(), Some((0, 1)));
  /// assert_eq!(it.next(), None);
  /// ```
  #[must_use]
  pub const fn iter_ones_block(&self) -> BitBlocksIter<'_> {
    BitBlocksIter {
      bmap: self,
      idx: 0,
      set: true
    }
  }
}

impl<'a> IntoIterator for &'a CountedBitmap {
  type Item = bool;
  type IntoIter = BitIter<'a>;
  fn into_iter(self) -> Self::IntoIter {
    self.iter()
  }
}

/// An [`Iterator`] which iterates over bit map and yields `bool` values where
/// `false` represents `0` and `true` represents `1`.
pub struct BitIter<'a> {
  bmap: &'a CountedBitmap,
  idx: usize
}

impl Iterator for BitIter<'_> {
  type Item = bool;
  fn next(&mut self) -> Option<Self::Item> {
    if self.idx < self.bmap.len() {
      let val = self.bmap.is_set_unchecked(self.idx);
      self.idx += 1;
      Some(val)
    } else {
      None
    }
  }
}

impl ExactSizeIterator for BitIter<'_> {
  fn len(&self) -> usize {
    self.bmap.len()
  }
}


/// An [`Iterator`] which yields the indexes of either zeroes or ones,
/// depending on configuration.
pub struct BitValIter<'a> {
  bmap: &'a CountedBitmap,
  idx: usize,

  /// 'true' means search for 1's.  `false` means search for 0's.
  set: bool
}

impl Iterator for BitValIter<'_> {
  type Item = usize;
  fn next(&mut self) -> Option<Self::Item> {
    // ToDo: Can skip entire words
    while self.idx < self.bmap.length {
      let val = self.bmap.is_set_unchecked(self.idx);
      if self.set == val {
        let idx = self.idx;
        self.idx += 1;
        return Some(idx);
      }
      self.idx += 1;
    }
    None
  }
}

impl ExactSizeIterator for BitValIter<'_> {
  fn len(&self) -> usize {
    if self.set {
      // seaching for 1's
      self.bmap.length - self.bmap.remain
    } else {
      // seaching for 0's
      self.bmap.remain
    }
  }
}


/// An [`Iterator`] which yields range blocks of contiguous zeroes or ones.
pub struct BitBlocksIter<'a> {
  bmap: &'a CountedBitmap,
  idx: usize,

  /// 'true' means return blocks of 1's.  `false` means return blocks of for
  /// 0's.
  set: bool
}

impl Iterator for BitBlocksIter<'_> {
  type Item = (usize, usize);

  fn next(&mut self) -> Option<Self::Item> {
    while self.idx < self.bmap.length {
      // As long as the current bit is _not_ the one of interest, keep
      // increasing index and retrying.  If there are not more blocks of
      // interest the outer while loop will terminate.
      let val = self.bmap.is_set_unchecked(self.idx);
      if self.set != val {
        self.idx += 1;
        continue;
      }

      // If this point has been reached a range will be returned, because a bit
      // of interest has been encountered and it's within the range.

      let low = self.idx;
      self.idx += 1;

      // Keep scanning until either the end has been reached or wrong type of
      // bit is encountered.

      let high = loop {
        if self.idx < self.bmap.length {
          let val = self.bmap.is_set_unchecked(self.idx);
          if self.set == val {
            // Still within the same block -- keep going
            self.idx += 1;
            continue;
          }
          // Found uninteresting bit, so block has ended
          let idx = self.idx - 1;

          // No need to rescan this block next run
          self.idx += 1;
          break idx;
        }
        // end -- return last valid index
        break self.idx - 1;
      };

      return Some((low, high));
    }

    None
  }
}


/// Internals
impl CountedBitmap {
  /// Number of words in the bitmap.
  #[allow(dead_code)]
  fn word_count(&self) -> usize {
    self.bits.len()
  }

  /// Returns a two-tuple `(index, bitvalue)`.
  #[inline]
  const fn get_bidx(&self, idx: usize) -> Result<(usize, usize), Error> {
    if idx >= self.length {
      return Err(Error::OutOfBounds);
    }

    let t = Self::get_bidx_unchecked(idx);

    Ok(t)
  }

  /// Returns a two-tuple `(index, bitvalue)`.
  #[inline]
  const fn get_bidx_unchecked(idx: usize) -> (usize, usize) {
    // Don't worry about the division and modulo -- as long as NUM_WORD_BITS is
    // a power of two the compiler will generate raw bit operations rather than
    // actual division and modulo.

    let iword = idx / NUM_WORD_BITS;

    let ibit = idx % NUM_WORD_BITS;
    let bitval = 1 << ibit;

    (iword, bitval)
  }

  #[inline]
  fn is_set_unchecked(&self, idx: usize) -> bool {
    let (iword, bitval) = Self::get_bidx_unchecked(idx);
    let v = unsafe { self.bits.get_unchecked(iword) };
    v & bitval != 0
  }
}


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

  #[test]
  fn size() {
    let bmap = CountedBitmap::new(0);
    assert_eq!(bmap.len(), 0);
    assert_eq!(bmap.word_count(), 0);

    let bmap = CountedBitmap::new(1);
    assert_eq!(bmap.len(), 1);
    assert_eq!(bmap.word_count(), 1);

    let idx = NUM_WORD_BITS - 1;
    let bmap = CountedBitmap::new(idx);
    assert_eq!(bmap.len(), idx);
    assert_eq!(bmap.word_count(), 1);

    let bmap = CountedBitmap::new(NUM_WORD_BITS);
    assert_eq!(bmap.len(), NUM_WORD_BITS);
    assert_eq!(bmap.word_count(), 1);

    let idx = NUM_WORD_BITS + 1;
    let bmap = CountedBitmap::new(idx);
    assert_eq!(bmap.len(), idx);
    assert_eq!(bmap.word_count(), 2);
  }

  #[test]
  fn finished() {
    let bmap = CountedBitmap::new(0);
    assert!(bmap.is_finished());

    let mut bmap = CountedBitmap::new(1);
    assert!(!bmap.is_finished());
    bmap.set(0).unwrap();
    assert!(bmap.is_finished());
  }

  #[test]
  fn dbg_output0() {
    let bmap = CountedBitmap::new(3);
    let s = format!("{bmap:?}");
    assert_eq!(s, "BMap { length: 3, remain: 3, [] }");
  }

  #[test]
  fn dbg_output1() {
    let mut bmap = CountedBitmap::new(3);
    bmap.set(1).unwrap();
    let s = format!("{bmap:?}");
    assert_eq!(s, "BMap { length: 3, remain: 2, [1] }");
  }

  #[test]
  fn dbg_output2() {
    let mut bmap = CountedBitmap::new(5);
    bmap.set(1).unwrap();
    bmap.set(2).unwrap();
    let s = format!("{bmap:?}");
    assert_eq!(s, "BMap { length: 5, remain: 3, [1-2] }");
  }

  #[test]
  fn dbg_output3() {
    let mut bmap = CountedBitmap::new(7);
    bmap.set(1).unwrap();
    bmap.set(2).unwrap();

    bmap.set(4).unwrap();
    bmap.set(5).unwrap();

    let s = format!("{bmap:?}");
    assert_eq!(s, "BMap { length: 7, remain: 3, [1-2,4-5] }");
  }
}

// vim: set ft=rust et sw=2 ts=2 sts=2 cinoptions=2 tw=79 :