1 files changed, 117 insertions, 16 deletions

diff --git a/zjit/src/bitset.rs b/zjit/src/bitset.rs
index 895bac8e33..986d537d9b 100644
--- a/zjit/src/bitset.rs
+++ b/zjit/src/bitset.rs
@@ -1,3 +1,5 @@
+//! Optimized bitset implementation.
+
 type Entry = u128;
 
 const ENTRY_NUM_BITS: usize = Entry::BITS as usize;
@@ -35,6 +37,16 @@ impl<T: Into<usize> + Copy> BitSet<T> {
         }
     }
 
+    /// Clear a bit. Returns whether the bit was previously set.
+    pub fn remove(&mut self, idx: T) -> bool {
+        debug_assert!(idx.into() < self.num_bits);
+        let entry_idx = idx.into() / ENTRY_NUM_BITS;
+        let bit_idx = idx.into() % ENTRY_NUM_BITS;
+        let was_set = (self.entries[entry_idx] & (1 << bit_idx)) != 0;
+        self.entries[entry_idx] &= !(1 << bit_idx);
+        was_set
+    }
+
     pub fn get(&self, idx: T) -> bool {
         debug_assert!(idx.into() < self.num_bits);
         let entry_idx = idx.into() / ENTRY_NUM_BITS;
@@ -55,6 +67,57 @@ impl<T: Into<usize> + Copy> BitSet<T> {
         }
         changed
     }
+
+    /// Modify `self` to have bits set if they are set in either `self` or `other`. Returns true if `self`
+    /// was modified, and false otherwise.
+    /// `self` and `other` must have the same number of bits.
+    pub fn union_with(&mut self, other: &Self) -> bool {
+        assert_eq!(self.num_bits, other.num_bits);
+        let mut changed = false;
+        for i in 0..self.entries.len() {
+            let before = self.entries[i];
+            self.entries[i] |= other.entries[i];
+            changed |= self.entries[i] != before;
+        }
+        changed
+    }
+
+    /// Modify `self` to remove bits that are set in `other`. Returns true if `self`
+    /// was modified, and false otherwise.
+    /// `self` and `other` must have the same number of bits.
+    pub fn difference_with(&mut self, other: &Self) -> bool {
+        assert_eq!(self.num_bits, other.num_bits);
+        let mut changed = false;
+        for i in 0..self.entries.len() {
+            let before = self.entries[i];
+            self.entries[i] &= !other.entries[i];
+            changed |= self.entries[i] != before;
+        }
+        changed
+    }
+
+    /// Check if two BitSets are equal.
+    /// `self` and `other` must have the same number of bits.
+    pub fn equals(&self, other: &Self) -> bool {
+        assert_eq!(self.num_bits, other.num_bits);
+        self.entries == other.entries
+    }
+
+    /// Returns an iterator over the indices of set bits.
+    /// Only iterates over bits that are set, not all possible indices.
+    pub fn iter_set_bits(&self) -> impl Iterator<Item = usize> + '_ {
+        self.entries.iter().enumerate().flat_map(move |(entry_idx, &entry)| {
+            let mut bits = entry;
+            std::iter::from_fn(move || {
+                if bits == 0 {
+                    return None;
+                }
+                let bit_pos = bits.trailing_zeros() as usize;
+                bits &= bits - 1; // Clear the lowest set bit
+                Some(entry_idx * ENTRY_NUM_BITS + bit_pos)
+            })
+        }).filter(move |&idx| idx < self.num_bits)
+    }
 }
 
 #[cfg(test)]
@@ -65,40 +128,40 @@ mod tests {
     #[should_panic]
     fn get_over_capacity_panics() {
         let set = BitSet::with_capacity(0);
-        assert_eq!(set.get(0usize), false);
+        assert!(!set.get(0usize));
     }
 
     #[test]
     fn with_capacity_defaults_to_zero() {
         let set = BitSet::with_capacity(4);
-        assert_eq!(set.get(0usize), false);
-        assert_eq!(set.get(1usize), false);
-        assert_eq!(set.get(2usize), false);
-        assert_eq!(set.get(3usize), false);
+        assert!(!set.get(0usize));
+        assert!(!set.get(1usize));
+        assert!(!set.get(2usize));
+        assert!(!set.get(3usize));
     }
 
     #[test]
     fn insert_sets_bit() {
         let mut set = BitSet::with_capacity(4);
-        assert_eq!(set.insert(1usize), true);
-        assert_eq!(set.get(1usize), true);
+        assert!(set.insert(1usize));
+        assert!(set.get(1usize));
     }
 
     #[test]
     fn insert_with_set_bit_returns_false() {
         let mut set = BitSet::with_capacity(4);
-        assert_eq!(set.insert(1usize), true);
-        assert_eq!(set.insert(1usize), false);
+        assert!(set.insert(1usize));
+        assert!(!set.insert(1usize));
     }
 
     #[test]
     fn insert_all_sets_all_bits() {
         let mut set = BitSet::with_capacity(4);
         set.insert_all();
-        assert_eq!(set.get(0usize), true);
-        assert_eq!(set.get(1usize), true);
-        assert_eq!(set.get(2usize), true);
-        assert_eq!(set.get(3usize), true);
+        assert!(set.get(0usize));
+        assert!(set.get(1usize));
+        assert!(set.get(2usize));
+        assert!(set.get(3usize));
     }
 
     #[test]
@@ -117,8 +180,46 @@ mod tests {
         right.insert(1usize);
         right.insert(2usize);
         left.intersect_with(&right);
-        assert_eq!(left.get(0usize), false);
-        assert_eq!(left.get(1usize), true);
-        assert_eq!(left.get(2usize), false);
+        assert!(!left.get(0usize));
+        assert!(left.get(1usize));
+        assert!(!left.get(2usize));
+    }
+
+    #[test]
+    fn test_iter_set_bits() {
+        let mut set: BitSet<usize> = BitSet::with_capacity(10);
+        set.insert(1usize);
+        set.insert(5usize);
+        set.insert(9usize);
+
+        let set_bits: Vec<usize> = set.iter_set_bits().collect();
+        assert_eq!(set_bits, vec![1, 5, 9]);
+    }
+
+    #[test]
+    fn test_iter_set_bits_empty() {
+        let set: BitSet<usize> = BitSet::with_capacity(10);
+        let set_bits: Vec<usize> = set.iter_set_bits().collect();
+        assert_eq!(set_bits, vec![]);
+    }
+
+    #[test]
+    fn test_iter_set_bits_all() {
+        let mut set: BitSet<usize> = BitSet::with_capacity(5);
+        set.insert_all();
+        let set_bits: Vec<usize> = set.iter_set_bits().collect();
+        assert_eq!(set_bits, vec![0, 1, 2, 3, 4]);
+    }
+
+    #[test]
+    fn test_iter_set_bits_large() {
+        let mut set: BitSet<usize> = BitSet::with_capacity(200);
+        set.insert(0usize);
+        set.insert(127usize);
+        set.insert(128usize);
+        set.insert(199usize);
+
+        let set_bits: Vec<usize> = set.iter_set_bits().collect();
+        assert_eq!(set_bits, vec![0, 127, 128, 199]);
     }
 }