(DO NOT MERGE) Extract insert and remove (tweaked)

Author: Zalathar

compiler/rustc_index/src/bit_set.rs

@@ -185,21 +185,11 @@ impl<T: Idx> DenseBitSet<T> {
         self.words.iter().all(|a| *a == 0)
     }
 
-    /// Insert `elem`. Returns whether the set has changed.
+    /// Inserts `value` into the set, and returns true if the set has changed
+    /// (i.e. the set did not contain the value).
     #[inline]
-    pub fn insert(&mut self, elem: T) -> bool {
-        assert!(
-            elem.index() < self.domain_size,
-            "inserting element at index {} but domain size is {}",
-            elem.index(),
-            self.domain_size,
-        );
-        let (word_index, mask) = word_index_and_mask(elem);
-        let word_ref = &mut self.words[word_index];
-        let word = *word_ref;
-        let new_word = word | mask;
-        *word_ref = new_word;
-        new_word != word
+    pub fn insert(&mut self, value: T) -> bool {
+        raw::insert(self.domain_size, &mut self.words, value.index())
     }
 
     #[inline]
@@ -223,16 +213,11 @@ impl<T: Idx> DenseBitSet<T> {
         raw::contains_any(self.domain_size, &self.words, (start, end))
     }
 
-    /// Returns `true` if the set has changed.
+    /// Removes `value` from the set, and returns true if the set has changed
+    /// (i.e. the set did contain the value).
     #[inline]
-    pub fn remove(&mut self, elem: T) -> bool {
-        assert!(elem.index() < self.domain_size);
-        let (word_index, mask) = word_index_and_mask(elem);
-        let word_ref = &mut self.words[word_index];
-        let word = *word_ref;
-        let new_word = word & !mask;
-        *word_ref = new_word;
-        new_word != word
+    pub fn remove(&mut self, value: T) -> bool {
+        raw::remove(self.domain_size, &mut self.words, value.index())
     }
 
     /// Iterates over the indices of set bits in a sorted order.
@@ -1608,12 +1593,9 @@ fn num_chunks<T: Idx>(domain_size: T) -> usize {
     domain_size.index().div_ceil(CHUNK_BITS)
 }
 
-#[inline]
+#[inline(always)]
 fn word_index_and_mask<T: Idx>(elem: T) -> (usize, Word) {
-    let elem = elem.index();
-    let word_index = elem / WORD_BITS;
-    let mask = 1 << (elem % WORD_BITS);
-    (word_index, mask)
+    raw::word_index_and_mask(elem.index())
 }
 
 #[inline]

compiler/rustc_index/src/bit_set/raw.rs

@@ -1,7 +1,7 @@
 use std::ops::Bound;
 use std::slice;
 
-use crate::bit_set::{WORD_BITS, Word, inclusive_start_end, max_bit, word_index_and_mask};
+use crate::bit_set::{WORD_BITS, Word, inclusive_start_end, max_bit};
 
 #[inline]
 pub(crate) fn contains_any(
@@ -67,6 +67,40 @@ pub(crate) fn last_set_in(
     None
 }
 
+#[inline(always)]
+pub(crate) fn insert(domain_size: usize, words: &mut [Word], index: usize) -> bool {
+    if index >= domain_size {
+        index_not_in_domain("inserting", index, domain_size);
+    }
+
+    let (word_index, mask) = word_index_and_mask(index);
+    modify_word(&mut words[word_index], |w| w | mask)
+}
+
+#[inline(always)]
+pub(crate) fn remove(domain_size: usize, words: &mut [Word], index: usize) -> bool {
+    if index >= domain_size {
+        index_not_in_domain("removing", index, domain_size);
+    }
+
+    let (word_index, mask) = word_index_and_mask(index);
+    modify_word(&mut words[word_index], |w| w & !mask)
+}
+
+#[cold]
+#[inline(never)]
+fn index_not_in_domain(verb: &'static str, index: usize, domain_size: usize) -> ! {
+    panic!("{verb} at index {index} but domain size is {domain_size}");
+}
+
+/// Updates a `&mut Word` using the given function, and returns true if the value was changed.
+#[inline(always)]
+fn modify_word(word: &mut Word, modify_fn: impl Fn(Word) -> Word) -> bool {
+    let old = *word;
+    *word = modify_fn(old);
+    old != *word
+}
+
 #[inline]
 pub(crate) fn insert_range(
     domain_size: usize,
@@ -149,3 +183,10 @@ impl<'a> Iterator for RawBitIter<'a> {
         }
     }
 }
+
+#[inline(always)]
+pub(crate) fn word_index_and_mask(bit_index: usize) -> (usize, Word) {
+    let word_index = bit_index / WORD_BITS;
+    let mask = 1 << (bit_index % WORD_BITS);
+    (word_index, mask)
+}