Skip to content

Experiment with lazy allocation for empty DenseBitSet storage #157024

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
Jamesbarford wants to merge 1 commit into
base: main
Choose a base branch
from
+194 −44
Draft

Experiment with lazy allocation for empty DenseBitSet storage #157024

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
183 changes: 139 additions & 44 deletions compiler/rustc_index/src/bit_set.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use std::hash::{Hash, Hasher};
use std::marker::PhantomData;
use std::ops::{Bound, Range, RangeBounds};
use std::rc::Rc;
Expand Down Expand Up @@ -112,10 +113,10 @@ macro_rules! bit_relations_inherent_impls {
/// will panic if the bitsets have differing domain sizes.
///
#[cfg_attr(feature = "nightly", derive(Decodable_NoContext, Encodable_NoContext))]
#[derive(Eq, PartialEq, Hash)]
pub struct DenseBitSet<T> {
domain_size: usize,
words: Vec<Word>,
// Lazily allocated. `None` represents a full set of zero words.
words: Option<Vec<Word>>,
marker: PhantomData<T>,
}

Expand All @@ -124,61 +125,90 @@ impl<T> DenseBitSet<T> {
pub fn domain_size(&self) -> usize {
self.domain_size
}

#[inline]
fn num_words(&self) -> usize {
num_words(self.domain_size)
}

#[inline]
fn words(&self) -> &[Word] {
self.words.as_deref().unwrap_or(&[])
}

#[inline]
fn ensure_words(&mut self) -> &mut Vec<Word> {
let num_words = self.num_words();
self.words.get_or_insert_with(|| vec![0; num_words])
}

#[inline]
fn word(&self, word_index: usize) -> Word {
self.words().get(word_index).copied().unwrap_or(0)
}
}

impl<T: Idx> DenseBitSet<T> {
/// Creates a new, empty bitset with a given `domain_size`.
#[inline]
pub fn new_empty(domain_size: usize) -> DenseBitSet<T> {
let num_words = num_words(domain_size);
DenseBitSet { domain_size, words: vec![0; num_words], marker: PhantomData }
DenseBitSet { domain_size, words: None, marker: PhantomData }
}

/// Creates a new, filled bitset with a given `domain_size`.
#[inline]
pub fn new_filled(domain_size: usize) -> DenseBitSet<T> {
let num_words = num_words(domain_size);
let mut result =
DenseBitSet { domain_size, words: vec![!0; num_words], marker: PhantomData };
let words = if num_words == 0 { None } else { Some(vec![!0; num_words]) };
let mut result = DenseBitSet { domain_size, words, marker: PhantomData };
result.clear_excess_bits();
result
}

/// Clear all elements.
#[inline]
pub fn clear(&mut self) {
self.words.fill(0);
if let Some(words) = &mut self.words {
words.fill(0);
}
}

/// Clear excess bits in the final word.
fn clear_excess_bits(&mut self) {
clear_excess_bits_in_final_word(self.domain_size, &mut self.words);
if let Some(words) = &mut self.words {
clear_excess_bits_in_final_word(self.domain_size, words);
}
}

/// Count the number of set bits in the set.
pub fn count(&self) -> usize {
count_ones(&self.words)
count_ones(self.words())
}

/// Returns `true` if `self` contains `elem`.
#[inline]
pub fn contains(&self, elem: T) -> bool {
assert!(elem.index() < self.domain_size);
let (word_index, mask) = word_index_and_mask(elem);
(self.words[word_index] & mask) != 0
(self.word(word_index) & mask) != 0
}

/// Is `self` is a (non-strict) superset of `other`?
#[inline]
pub fn superset(&self, other: &DenseBitSet<T>) -> bool {
assert_eq!(self.domain_size, other.domain_size);
self.words.iter().zip(&other.words).all(|(a, b)| (a & b) == *b)
match (&self.words, &other.words) {
(_, None) => true,
(None, Some(other_words)) => other_words.iter().all(|&word| word == 0),
(Some(_), Some(_)) => (0..self.num_words())
.all(|index| (self.word(index) & other.word(index)) == other.word(index)),
}
}

/// Is the set empty?
#[inline]
pub fn is_empty(&self) -> bool {
self.words.iter().all(|a| *a == 0)
self.words().iter().all(|a| *a == 0)
}

/// Insert `elem`. Returns whether the set has changed.
Expand All @@ -191,7 +221,7 @@ impl<T: Idx> DenseBitSet<T> {
self.domain_size,
);
let (word_index, mask) = word_index_and_mask(elem);
let word_ref = &mut self.words[word_index];
let word_ref = &mut self.ensure_words()[word_index];
let word = *word_ref;
let new_word = word | mask;
*word_ref = new_word;
Expand All @@ -206,28 +236,32 @@ impl<T: Idx> DenseBitSet<T> {

let (start_word_index, start_mask) = word_index_and_mask(start);
let (end_word_index, end_mask) = word_index_and_mask(end);
let words = self.ensure_words();

// Set all words in between start and end (exclusively of both).
for word_index in (start_word_index + 1)..end_word_index {
self.words[word_index] = !0;
words[word_index] = !0;
}

if start_word_index != end_word_index {
// Start and end are in different words, so we handle each in turn.
//
// We set all leading bits. This includes the start_mask bit.
self.words[start_word_index] |= !(start_mask - 1);
words[start_word_index] |= !(start_mask - 1);
// And all trailing bits (i.e. from 0..=end) in the end word,
// including the end.
self.words[end_word_index] |= end_mask | (end_mask - 1);
words[end_word_index] |= end_mask | (end_mask - 1);
} else {
self.words[start_word_index] |= end_mask | (end_mask - start_mask);
words[start_word_index] |= end_mask | (end_mask - start_mask);
}
}

/// Sets all bits to true.
pub fn insert_all(&mut self) {
self.words.fill(!0);
if self.domain_size == 0 {
return;
}
self.ensure_words().fill(!0);
self.clear_excess_bits();
}

Expand All @@ -241,16 +275,16 @@ impl<T: Idx> DenseBitSet<T> {
let (end_word_index, end_mask) = word_index_and_mask(end);

if start_word_index == end_word_index {
self.words[start_word_index] & (end_mask | (end_mask - start_mask)) != 0
self.word(start_word_index) & (end_mask | (end_mask - start_mask)) != 0
} else {
if self.words[start_word_index] & !(start_mask - 1) != 0 {
if self.word(start_word_index) & !(start_mask - 1) != 0 {
return true;
}

let remaining = start_word_index + 1..end_word_index;
if remaining.start <= remaining.end {
self.words[remaining].iter().any(|&w| w != 0)
|| self.words[end_word_index] & (end_mask | (end_mask - 1)) != 0
self.words().get(remaining).is_some_and(|words| words.iter().any(|&w| w != 0))
|| self.word(end_word_index) & (end_mask | (end_mask - 1)) != 0
} else {
false
}
Expand All @@ -261,8 +295,11 @@ impl<T: Idx> DenseBitSet<T> {
#[inline]
pub fn remove(&mut self, elem: T) -> bool {
assert!(elem.index() < self.domain_size);
let Some(words) = &mut self.words else {
return false;
};
let (word_index, mask) = word_index_and_mask(elem);
let word_ref = &mut self.words[word_index];
let word_ref = &mut words[word_index];
let word = *word_ref;
let new_word = word & !mask;
*word_ref = new_word;
Expand All @@ -272,15 +309,18 @@ impl<T: Idx> DenseBitSet<T> {
/// Iterates over the indices of set bits in a sorted order.
#[inline]
pub fn iter(&self) -> BitIter<'_, T> {
BitIter::new(&self.words)
BitIter::new(self.words())
}

pub fn last_set_in(&self, range: impl RangeBounds<T>) -> Option<T> {
let (start, end) = inclusive_start_end(range, self.domain_size)?;
let Some(words) = &self.words else {
return None;
};
let (start_word_index, _) = word_index_and_mask(start);
let (end_word_index, end_mask) = word_index_and_mask(end);

let end_word = self.words[end_word_index] & (end_mask | (end_mask - 1));
let end_word = words[end_word_index] & (end_mask | (end_mask - 1));
if end_word != 0 {
let pos = max_bit(end_word) + WORD_BITS * end_word_index;
if start <= pos {
Expand All @@ -291,11 +331,10 @@ impl<T: Idx> DenseBitSet<T> {
// We exclude end_word_index from the range here, because we don't want
// to limit ourselves to *just* the last word: the bits set it in may be
// after `end`, so it may not work out.
if let Some(offset) =
self.words[start_word_index..end_word_index].iter().rposition(|&w| w != 0)
if let Some(offset) = words[start_word_index..end_word_index].iter().rposition(|&w| w != 0)
{
let word_idx = start_word_index + offset;
let start_word = self.words[word_idx];
let start_word = words[word_idx];
let pos = max_bit(start_word) + WORD_BITS * word_idx;
if start <= pos {
return Some(T::new(pos));
Expand All @@ -319,7 +358,13 @@ impl<T: Idx> DenseBitSet<T> {
// quickly and accurately detect whether the update changed anything.
// But that's only worth doing if there's an actual use-case.

update_words(&mut self.words, &other.words, |a, b| a | !b);
let other_words = other.words();
let words = self.ensure_words();
if other_words.is_empty() {
words.fill(!0);
} else {
update_words(words, other_words, |a, b| a | !b);
}
// The bitwise update `a | !b` can result in the last word containing
// out-of-domain bits, so we need to clear them.
self.clear_excess_bits();
Expand All @@ -330,17 +375,38 @@ impl<T: Idx> DenseBitSet<T> {
impl<T: Idx> BitRelations<DenseBitSet<T>> for DenseBitSet<T> {
fn union(&mut self, other: &DenseBitSet<T>) -> bool {
assert_eq!(self.domain_size, other.domain_size);
update_words(&mut self.words, &other.words, |a, b| a | b)
let Some(other_words) = &other.words else {
return false;
};
let Some(words) = &mut self.words else {
if other_words.iter().all(|&word| word == 0) {
return false;
}
self.words = Some(other_words.clone());
return true;
};
update_words(words, other_words, |a, b| a | b)
}

fn subtract(&mut self, other: &DenseBitSet<T>) -> bool {
assert_eq!(self.domain_size, other.domain_size);
update_words(&mut self.words, &other.words, |a, b| a & !b)
let (Some(words), Some(other_words)) = (&mut self.words, &other.words) else {
return false;
};
update_words(words, other_words, |a, b| a & !b)
}

fn intersect(&mut self, other: &DenseBitSet<T>) -> bool {
assert_eq!(self.domain_size, other.domain_size);
update_words(&mut self.words, &other.words, |a, b| a & b)
let Some(words) = &mut self.words else {
return false;
};
let Some(other_words) = &other.words else {
let changed = words.iter().any(|&word| word != 0);
words.fill(0);
return changed;
};
update_words(words, other_words, |a, b| a & b)
}
}

Expand All @@ -365,6 +431,27 @@ impl<T> Clone for DenseBitSet<T> {
}
}

impl<T> PartialEq for DenseBitSet<T> {
fn eq(&self, other: &Self) -> bool {
self.domain_size == other.domain_size
&& (0..self.num_words()).all(|index| self.word(index) == other.word(index))
}
}

impl<T> Eq for DenseBitSet<T> {}

impl<T> Hash for DenseBitSet<T> {
fn hash<H: Hasher>(&self, state: &mut H) {
self.domain_size.hash(state);

let num_words = self.num_words();
num_words.hash(state);
for index in 0..num_words {
self.word(index).hash(state);
}
}
}

impl<T: Idx> fmt::Debug for DenseBitSet<T> {
fn fmt(&self, w: &mut fmt::Formatter<'_>) -> fmt::Result {
w.debug_list().entries(self.iter()).finish()
Expand All @@ -380,8 +467,8 @@ impl<T: Idx> ToString for DenseBitSet<T> {

// i tracks how many bits we have printed so far.
let mut i = 0;
for word in &self.words {
let mut word = *word;
for word_index in 0..self.num_words() {
let mut word = self.word(word_index);
for _ in 0..WORD_BYTES {
// for each byte in `word`:
let remain = self.domain_size - i;
Expand Down Expand Up @@ -1306,8 +1393,10 @@ impl<T: Idx> GrowableBitSet<T> {
}

let min_num_words = num_words(min_domain_size);
if self.bit_set.words.len() < min_num_words {
self.bit_set.words.resize(min_num_words, 0)
if let Some(words) = &mut self.bit_set.words {
if words.len() < min_num_words {
words.resize(min_num_words, 0)
}
}
}

Expand Down Expand Up @@ -1357,7 +1446,7 @@ impl<T: Idx> GrowableBitSet<T> {
#[inline]
pub fn contains(&self, elem: T) -> bool {
let (word_index, mask) = word_index_and_mask(elem);
self.bit_set.words.get(word_index).is_some_and(|word| (word & mask) != 0)
self.bit_set.words().get(word_index).is_some_and(|word| (word & mask) != 0)
}

#[inline]
Expand Down Expand Up @@ -1419,13 +1508,16 @@ impl<R: Idx, C: Idx> BitMatrix<R, C> {
pub fn from_row_n(row: &DenseBitSet<C>, num_rows: usize) -> BitMatrix<R, C> {
let num_columns = row.domain_size();
let words_per_row = num_words(num_columns);
assert_eq!(words_per_row, row.words.len());
BitMatrix {
num_rows,
num_columns,
words: iter::repeat_n(&row.words, num_rows).flatten().cloned().collect(),
marker: PhantomData,
let mut words = vec![0; num_rows * words_per_row];
if let Some(row_words) = &row.words {
assert_eq!(words_per_row, row_words.len());
if words_per_row > 0 {
for matrix_row in words.chunks_exact_mut(words_per_row) {
matrix_row.copy_from_slice(row_words);
}
}
}
BitMatrix { num_rows, num_columns, words, marker: PhantomData }
}

pub fn rows(&self) -> impl Iterator<Item = R> {
Expand Down Expand Up @@ -1517,8 +1609,11 @@ impl<R: Idx, C: Idx> BitMatrix<R, C> {
pub fn union_row_with(&mut self, with: &DenseBitSet<C>, write: R) -> bool {
assert!(write.index() < self.num_rows);
assert_eq!(with.domain_size(), self.num_columns);
let Some(with_words) = &with.words else {
return false;
};
let (write_start, write_end) = self.range(write);
update_words(&mut self.words[write_start..write_end], &with.words, |a, b| a | b)
update_words(&mut self.words[write_start..write_end], with_words, |a, b| a | b)
}

/// Sets every cell in `row` to true.
Expand Down
Loading
Loading