From 6ce0300e608cbee9730205f1f809adc1cb5ab300 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kevin=20L=C3=A4ufer?= <laeufer@cornell.edu>
Date: Thu, 21 Nov 2024 14:56:33 -0500
Subject: [PATCH] add bit_set_intervals function

---
 Cargo.toml           |  2 +-
 src/bv/arithmetic.rs | 59 ++++++++++++++++++++++++++++++++++++++++++++
 src/bv/ops.rs        | 37 +++++++++++++++++++++++++++
 3 files changed, 97 insertions(+), 1 deletion(-)

diff --git a/Cargo.toml b/Cargo.toml
index c89b462..55726f6 100644
--- a/Cargo.toml
+++ b/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"
diff --git a/src/bv/arithmetic.rs b/src/bv/arithmetic.rs
index d0fea0d..4d2a225 100644
--- a/src/bv/arithmetic.rs
+++ b/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
+}
diff --git a/src/bv/ops.rs b/src/bv/ops.rs
index d67675d..a19e120 100644
--- a/src/bv/ops.rs
+++ b/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)]
+        );
+    }
 }