Skip to content

Commit

Permalink
[red-knot] add BitSet::is_empty and BitSet::union (#13333)
Browse files Browse the repository at this point in the history
Add `::is_empty` and `::union` methods to the `BitSet` implementation.

Allowing unused for now, until these methods become used later with the
declared-types implementation.

---------

Co-authored-by: Alex Waygood <[email protected]>
  • Loading branch information
carljm and AlexWaygood authored Sep 12, 2024
1 parent 175d067 commit 43a5922
Showing 1 changed file with 85 additions and 2 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -32,17 +32,26 @@ impl<const B: usize> BitSet<B> {
bitset
}

#[allow(unused)]
pub(super) fn is_empty(&self) -> bool {
self.blocks().iter().all(|&b| b == 0)
}

/// Convert from Inline to Heap, if needed, and resize the Heap vector, if needed.
fn resize(&mut self, value: u32) {
let num_blocks_needed = (value / 64) + 1;
self.resize_blocks(num_blocks_needed as usize);
}

fn resize_blocks(&mut self, num_blocks_needed: usize) {
match self {
Self::Inline(blocks) => {
let mut vec = blocks.to_vec();
vec.resize(num_blocks_needed as usize, 0);
vec.resize(num_blocks_needed, 0);
*self = Self::Heap(vec);
}
Self::Heap(vec) => {
vec.resize(num_blocks_needed as usize, 0);
vec.resize(num_blocks_needed, 0);
}
}
}
Expand Down Expand Up @@ -89,6 +98,20 @@ impl<const B: usize> BitSet<B> {
}
}

/// Union in-place with another [`BitSet`].
#[allow(unused)]
pub(super) fn union(&mut self, other: &BitSet<B>) {
let mut max_len = self.blocks().len();
let other_len = other.blocks().len();
if other_len > max_len {
max_len = other_len;
self.resize_blocks(max_len);
}
for (my_block, other_block) in self.blocks_mut().iter_mut().zip(other.blocks()) {
*my_block |= other_block;
}
}

/// Return an iterator over the values (in ascending order) in this [`BitSet`].
pub(super) fn iter(&self) -> BitSetIterator<'_, B> {
let blocks = self.blocks();
Expand Down Expand Up @@ -218,11 +241,71 @@ mod tests {
assert_bitset(&b1, &[89]);
}

#[test]
fn union() {
let mut b1 = BitSet::<1>::with(2);
let b2 = BitSet::<1>::with(4);

b1.union(&b2);
assert_bitset(&b1, &[2, 4]);
}

#[test]
fn union_mixed_1() {
let mut b1 = BitSet::<1>::with(4);
let mut b2 = BitSet::<1>::with(4);
b1.insert(89);
b2.insert(5);

b1.union(&b2);
assert_bitset(&b1, &[4, 5, 89]);
}

#[test]
fn union_mixed_2() {
let mut b1 = BitSet::<1>::with(4);
let mut b2 = BitSet::<1>::with(4);
b1.insert(23);
b2.insert(89);

b1.union(&b2);
assert_bitset(&b1, &[4, 23, 89]);
}

#[test]
fn union_heap() {
let mut b1 = BitSet::<1>::with(4);
let mut b2 = BitSet::<1>::with(4);
b1.insert(89);
b2.insert(90);

b1.union(&b2);
assert_bitset(&b1, &[4, 89, 90]);
}

#[test]
fn union_heap_2() {
let mut b1 = BitSet::<1>::with(89);
let mut b2 = BitSet::<1>::with(89);
b1.insert(91);
b2.insert(90);

b1.union(&b2);
assert_bitset(&b1, &[89, 90, 91]);
}

#[test]
fn multiple_blocks() {
let mut b = BitSet::<2>::with(120);
b.insert(45);
assert!(matches!(b, BitSet::Inline(_)));
assert_bitset(&b, &[45, 120]);
}

#[test]
fn empty() {
let b = BitSet::<1>::default();

assert!(b.is_empty());
}
}

0 comments on commit 43a5922

Please sign in to comment.