add bit_set_intervals function

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "baa"
-version = "0.14.4"
+version = "0.14.5"
 edition = "2021"
 authors = ["Kevin Laeufer <laeufer@cornell.edu>"]
 description = "BitVector and Array Arithmetic"

src/bv/arithmetic.rs

@@ -435,3 +435,62 @@ pub(crate) fn assert_unused_bits_zero(value: &[Word], width: WidthInt) {
         assert_eq!(unused, 0, "unused msb bits need to be zero!")
     }
 }
+
+pub(crate) fn find_ranges_of_ones(words: &[Word]) -> Vec<std::ops::Range<WidthInt>> {
+    // the actual width does not matter since we assume that all unused bits in the msb are set to zero
+    let mut out = vec![];
+    let mut range_start: Option<WidthInt> = None;
+    for (word_ii, word) in words.iter().enumerate() {
+        let lsb_ii = word_ii as WidthInt * Word::BITS;
+        let mut word = *word;
+        let mut bits_consumed = 0;
+
+        // handle open range from previous word
+        if let Some(start) = range_start {
+            let ones = word.trailing_ones();
+            bits_consumed += ones;
+            word >>= ones;
+            if ones < Word::BITS {
+                range_start = None;
+                out.push(start..lsb_ii + bits_consumed);
+            }
+        }
+
+        // find ranges in this word
+        while bits_consumed < Word::BITS {
+            debug_assert!(range_start.is_none());
+            if word == 0 {
+                // done
+                bits_consumed = Word::BITS;
+            } else {
+                let zeros = word.trailing_zeros();
+                bits_consumed += zeros;
+                word >>= zeros;
+                let start = bits_consumed;
+                let ones = word.trailing_ones();
+                bits_consumed += ones;
+                word = word.overflowing_shr(ones).0;
+                match bits_consumed.cmp(&Word::BITS) {
+                    Ordering::Less => {
+                        let end = bits_consumed;
+                        out.push(lsb_ii + start..lsb_ii + end);
+                    }
+                    Ordering::Equal => {
+                        // done, range might expand to next word
+                        range_start = Some(start + lsb_ii);
+                    }
+                    Ordering::Greater => {
+                        unreachable!("")
+                    }
+                }
+            }
+        }
+    }
+    // finish open range
+    if let Some(start) = range_start {
+        let end = words.len() as WidthInt * Word::BITS;
+        out.push(start..end);
+    }
+
+    out
+}

src/bv/ops.rs

@@ -237,6 +237,19 @@ pub trait BitVecOps {
         bit_pos
     }
 
+    /// Computes all ranges for which the bits are one.
+    fn bit_set_intervals(&self) -> Vec<std::ops::Range<WidthInt>> {
+        match self.width() {
+            0 => vec![],
+            1 if self.is_zero() => vec![],
+            1 => {
+                let range = 0..1;
+                vec![range]
+            }
+            _ => crate::bv::arithmetic::find_ranges_of_ones(self.words()),
+        }
+    }
+
     declare_arith_bin_fn!(add);
     declare_arith_bin_fn!(sub);
     declare_arith_bin_fn!(shift_left);
@@ -659,4 +672,28 @@ mod tests {
             assert_eq!(a.is_pow_2(), Some(bit));
         }
     }
+
+    #[test]
+    fn test_bit_set_intervals() {
+        let a = BitVecValue::zero(1456);
+        assert_eq!(a.bit_set_intervals(), []);
+        let a = BitVecValue::from_u64(Word::MAX, Word::BITS);
+        assert_eq!(a.bit_set_intervals(), [0..Word::BITS]);
+        let mut a = BitVecValue::from_u64((1 as Word) << (Word::BITS - 1), 127);
+        assert_eq!(a.bit_set_intervals(), [(Word::BITS - 1)..Word::BITS]);
+        a.set_bit(Word::BITS);
+        assert_eq!(a.bit_set_intervals(), [(Word::BITS - 1)..(Word::BITS + 1)]);
+        a.clear_bit(Word::BITS - 1);
+        assert_eq!(a.bit_set_intervals(), [Word::BITS..(Word::BITS + 1)]);
+        a.set_bit(13);
+        assert_eq!(
+            a.bit_set_intervals(),
+            [13..14, Word::BITS..(Word::BITS + 1)]
+        );
+        a.set_bit(14);
+        assert_eq!(
+            a.bit_set_intervals(),
+            [13..15, Word::BITS..(Word::BITS + 1)]
+        );
+    }
 }