Skip to content

Commit 76d1d53

Browse files
martingmsmartingms
committed
Optimize ChunkedBitSet dense relations
1 parent b505555 commit 76d1d53

File tree

1 file changed

+161
-6
lines changed

1 file changed

+161
-6
lines changed

‎compiler/rustc_index/src/bit_set.rs

Lines changed: 161 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -651,24 +651,174 @@ impl<T: Idx> BitRelations<ChunkedBitSet<T>> for ChunkedBitSet<T> {
651651

652652
impl<T: Idx> BitRelations<HybridBitSet<T>> for ChunkedBitSet<T> {
653653
fn union(&mut self, other: &HybridBitSet<T>) -> bool {
654-
// FIXME: this is slow if `other` is dense, and could easily be
655-
// improved, but it hasn't been a problem in practice so far.
656654
assert_eq!(self.domain_size, other.domain_size());
657-
sequential_update(|elem| self.insert(elem), other.iter())
655+
656+
match other {
657+
HybridBitSet::Sparse(_) => sequential_update(|elem| self.insert(elem), other.iter()),
658+
HybridBitSet::Dense(dense) => self.union(dense),
659+
}
658660
}
659661

660662
fn subtract(&mut self, other: &HybridBitSet<T>) -> bool {
661-
// FIXME: this is slow if `other` is dense, and could easily be
662-
// improved, but it hasn't been a problem in practice so far.
663663
assert_eq!(self.domain_size, other.domain_size());
664-
sequential_update(|elem| self.remove(elem), other.iter())
664+
665+
match other {
666+
HybridBitSet::Sparse(_) => sequential_update(|elem| self.remove(elem), other.iter()),
667+
HybridBitSet::Dense(dense) => self.subtract(dense),
668+
}
665669
}
666670

667671
fn intersect(&mut self, _other: &HybridBitSet<T>) -> bool {
668672
unimplemented!("implement if/when necessary");
669673
}
670674
}
671675

676+
impl<T: Idx> BitRelations<BitSet<T>> for ChunkedBitSet<T> {
677+
fn union(&mut self, other: &BitSet<T>) -> bool {
678+
assert_eq!(self.domain_size, other.domain_size());
679+
680+
let mut changed = false;
681+
for (chunk, other_words) in self.chunks.iter_mut().zip(other.words().chunks(CHUNK_WORDS)) {
682+
match chunk {
683+
Zeros(chunk_domain_size) => {
684+
if let Some(first_nonzero_index) = first_nonzero(other_words) {
685+
let other_count =
686+
bit_count(&other_words[first_nonzero_index..]) as ChunkSize;
687+
debug_assert!(other_count <= *chunk_domain_size);
688+
if other_count == *chunk_domain_size {
689+
*chunk = Ones(*chunk_domain_size);
690+
changed = true;
691+
} else if other_count != 0 {
692+
// We take some effort to avoid copying the words.
693+
let words = Rc::<[Word; CHUNK_WORDS]>::new_zeroed();
694+
// SAFETY: `words` can safely be all zeroes.
695+
let mut words = unsafe { words.assume_init() };
696+
let words_ref = Rc::get_mut(&mut words).unwrap();
697+
698+
debug_assert_eq!(
699+
num_words(*chunk_domain_size as usize),
700+
other_words.len()
701+
);
702+
words_ref[first_nonzero_index..other_words.len()]
703+
.copy_from_slice(&other_words[first_nonzero_index..]);
704+
705+
*chunk = Mixed(*chunk_domain_size, other_count, words);
706+
changed = true;
707+
}
708+
}
709+
}
710+
Ones(_) => {}
711+
Mixed(chunk_domain_size, chunk_count, chunk_words) => {
712+
if let Some(first_nonzero_index) = first_nonzero(other_words) {
713+
debug_assert_eq!(num_words(*chunk_domain_size as usize), other_words.len());
714+
let op = |a, b| a | b;
715+
if bitwise_changes(
716+
&chunk_words[first_nonzero_index..other_words.len()],
717+
&other_words[first_nonzero_index..],
718+
op,
719+
) {
720+
let chunk_words = Rc::make_mut(chunk_words);
721+
let has_changed = bitwise(
722+
&mut chunk_words[first_nonzero_index..other_words.len()],
723+
&other_words[first_nonzero_index..],
724+
op,
725+
);
726+
debug_assert!(has_changed);
727+
728+
*chunk_count = bit_count(chunk_words) as ChunkSize;
729+
debug_assert!(*chunk_count > 0);
730+
if *chunk_count == *chunk_domain_size {
731+
*chunk = Ones(*chunk_domain_size);
732+
}
733+
changed = true
734+
}
735+
}
736+
}
737+
}
738+
}
739+
changed
740+
}
741+
742+
fn subtract(&mut self, other: &BitSet<T>) -> bool {
743+
assert_eq!(self.domain_size, other.domain_size());
744+
745+
let mut changed = false;
746+
for (chunk, other_words) in self.chunks.iter_mut().zip(other.words().chunks(CHUNK_WORDS)) {
747+
match chunk {
748+
Zeros(_) => {}
749+
Ones(chunk_domain_size) => {
750+
if let Some(first_nonzero_index) = first_nonzero(other_words) {
751+
let other_count =
752+
bit_count(&other_words[first_nonzero_index..]) as ChunkSize;
753+
debug_assert!(other_count <= *chunk_domain_size);
754+
if other_count == *chunk_domain_size {
755+
*chunk = Zeros(*chunk_domain_size);
756+
changed = true;
757+
} else {
758+
// We take some effort to avoid copying the words.
759+
let words = Rc::<[Word; CHUNK_WORDS]>::new_zeroed();
760+
// SAFETY: `words` can safely be all zeroes.
761+
let mut words = unsafe { words.assume_init() };
762+
let words_ref = Rc::get_mut(&mut words).unwrap();
763+
764+
debug_assert_eq!(
765+
num_words(*chunk_domain_size as usize),
766+
other_words.len()
767+
);
768+
for (word, other) in words_ref[first_nonzero_index..]
769+
.iter_mut()
770+
.zip(other_words[first_nonzero_index..].iter())
771+
{
772+
*word = !other;
773+
}
774+
775+
clear_excess_bits_in_final_word(
776+
*chunk_domain_size as usize,
777+
&mut words_ref[..other_words.len()],
778+
);
779+
780+
*chunk =
781+
Mixed(*chunk_domain_size, *chunk_domain_size - other_count, words);
782+
changed = true;
783+
}
784+
}
785+
}
786+
Mixed(chunk_domain_size, chunk_count, chunk_words) => {
787+
if let Some(first_nonzero_index) = first_nonzero(other_words) {
788+
debug_assert_eq!(num_words(*chunk_domain_size as usize), other_words.len());
789+
let op = |a, b: Word| a & !b;
790+
if bitwise_changes(
791+
&chunk_words[first_nonzero_index..other_words.len()],
792+
&other_words[first_nonzero_index..],
793+
op,
794+
) {
795+
let chunk_words = Rc::make_mut(chunk_words);
796+
let has_changed = bitwise(
797+
&mut chunk_words[first_nonzero_index..other_words.len()],
798+
&other_words[first_nonzero_index..],
799+
op,
800+
);
801+
debug_assert!(has_changed);
802+
803+
*chunk_count = bit_count(chunk_words) as ChunkSize;
804+
debug_assert!(chunk_count < chunk_domain_size);
805+
if *chunk_count == 0 {
806+
*chunk = Zeros(*chunk_domain_size);
807+
}
808+
changed = true
809+
}
810+
}
811+
}
812+
}
813+
}
814+
changed
815+
}
816+
817+
fn intersect(&mut self, _other: &BitSet<T>) -> bool {
818+
unimplemented!("implement if/when necessary");
819+
}
820+
}
821+
672822
impl<T> Clone for ChunkedBitSet<T> {
673823
fn clone(&self) -> Self {
674824
ChunkedBitSet {
@@ -1771,6 +1921,11 @@ fn chunk_word_index_and_mask<T: Idx>(elem: T) -> (usize, Word) {
17711921
word_index_and_mask(chunk_elem)
17721922
}
17731923

1924+
#[inline]
1925+
fn first_nonzero(words: &[Word]) -> Option<usize> {
1926+
words.iter().position(|w| *w != 0)
1927+
}
1928+
17741929
fn clear_excess_bits_in_final_word(domain_size: usize, words: &mut [Word]) {
17751930
let num_bits_in_final_word = domain_size % WORD_BITS;
17761931
if num_bits_in_final_word > 0 {

0 commit comments

Comments
 (0)