Auto merge of #153754 - Zalathar:would-change, r=<try>

bors · bors · commit 047dc4e3e9b4 · 2026-03-12T03:09:09.000Z
Avoid the would-change check for bitset chunks that are unique
diff --git a/compiler/rustc_index/src/bit_set.rs b/compiler/rustc_index/src/bit_set.rs
@@ -743,6 +743,39 @@ impl<T: Idx> ChunkedBitSet<T> {
     bit_relations_inherent_impls! {}
 }
 
+/// Behaves like [`Rc::make_mut`], except that if an inner clone would be
+/// required but `would_change` returns false, then no inner clone occurs and
+/// None is returned instead.
+///
+/// This function assumes that `would_change` is cheaper than an inner clone,
+/// but expensive enough to be worth avoiding if an inner clone isn't required.
+///
+/// In other words, this function returns:
+/// - Some (without cloning) if [`Rc::get_mut`] succeeds
+/// - None (without cloning) if `would_change` returns false
+/// - Some (after cloning) otherwise
+fn make_chunk_words_mut_for_change<'a>(
+    self_words: &'a mut Rc<[Word; CHUNK_WORDS]>,
+    would_change: impl Fn(&[Word; CHUNK_WORDS]) -> bool,
+) -> Option<&'a mut [Word; CHUNK_WORDS]> {
+    // If the Rc is already unique, we can skip the would-change check entirely.
+    if let Some(words) = Rc::get_mut(self_words) {
+        // Round-trip through a raw pointer to work around a well-known limitation of the
+        // non-Polonius borrow checker <https://github.com/rust-lang/rust/issues/54663>.
+        // SAFETY: Equivalent to `Some(words)`, which is sound but only accepted by Polonius;
+        // see `tests/ui/nll/polonius/nll-problem-case-3-issue-54663.rs`.
+        return Some(unsafe { &mut *(words as *mut _) });
+    };
+
+    // At this point we would need to clone the chunk if we wanted to return `&mut`,
+    // so first check if the intended operation would actually modify the words.
+    if !would_change(self_words) {
+        return None;
+    }
+
+    Some(Rc::make_mut(self_words))
+}
+
 impl<T: Idx> BitRelations<ChunkedBitSet<T>> for ChunkedBitSet<T> {
     fn union(&mut self, other: &ChunkedBitSet<T>) -> bool {
         assert_eq!(self.domain_size, other.domain_size);
@@ -780,38 +813,36 @@ impl<T: Idx> BitRelations<ChunkedBitSet<T>> for ChunkedBitSet<T> {
                     // performance win. Also, we only need to operate on the
                     // in-use words, hence the slicing.
                     let num_words = num_words(chunk_domain_size as usize);
-
-                    // If both sides are the same, nothing will change. This
-                    // case is very common and it's a pretty fast check, so
-                    // it's a performance win to do it.
-                    if self_chunk_words[0..num_words] == other_chunk_words[0..num_words] {
-                        continue;
-                    }
-
-                    // Do a more precise "will anything change?" test. Also a
-                    // performance win.
                     let op = |a, b| a | b;
-                    if !bitwise_changes(
-                        &self_chunk_words[0..num_words],
-                        &other_chunk_words[0..num_words],
-                        op,
-                    ) {
+
+                    // Try to get a mutable reference to our words, but only if doing so is
+                    // cheap, or if we're definitely going to modify the words.
+                    let Some(self_chunk_words) =
+                        make_chunk_words_mut_for_change(self_chunk_words, |self_chunk_words| {
+                            let lhs = &self_chunk_words[0..num_words];
+                            let rhs = &other_chunk_words[0..num_words];
+                            // If both sides are the same, nothing will change. That case is very
+                            // common, so it's a performance win to do an extremely fast equality
+                            // check before the full would-change check.
+                            (lhs != rhs) && bitwise_changes(lhs, rhs, op)
+                        })
+                    else {
                         continue;
-                    }
+                    };
 
-                    // If we reach here, `self_chunk_words` is definitely changing.
-                    let self_chunk_words = Rc::make_mut(self_chunk_words);
                     let has_changed = bitwise(
                         &mut self_chunk_words[0..num_words],
                         &other_chunk_words[0..num_words],
                         op,
                     );
-                    debug_assert!(has_changed);
-                    *self_chunk_count = count_ones(&self_chunk_words[0..num_words]) as ChunkSize;
-                    if *self_chunk_count == chunk_domain_size {
-                        *self_chunk = Ones;
+                    if has_changed {
+                        *self_chunk_count =
+                            count_ones(&self_chunk_words[0..num_words]) as ChunkSize;
+                        if *self_chunk_count == chunk_domain_size {
+                            *self_chunk = Ones;
+                        }
+                        changed = true;
                     }
-                    changed = true;
                 }
             }
         }
@@ -868,26 +899,32 @@ impl<T: Idx> BitRelations<ChunkedBitSet<T>> for ChunkedBitSet<T> {
                     // See `ChunkedBitSet::union` for details on what is happening here.
                     let num_words = num_words(chunk_domain_size as usize);
                     let op = |a: Word, b: Word| a & !b;
-                    if !bitwise_changes(
-                        &self_chunk_words[0..num_words],
-                        &other_chunk_words[0..num_words],
-                        op,
-                    ) {
+
+                    let Some(self_chunk_words) =
+                        make_chunk_words_mut_for_change(self_chunk_words, |self_chunk_words| {
+                            bitwise_changes(
+                                &self_chunk_words[0..num_words],
+                                &other_chunk_words[0..num_words],
+                                op,
+                            )
+                        })
+                    else {
                         continue;
-                    }
+                    };
 
-                    let self_chunk_words = Rc::make_mut(self_chunk_words);
                     let has_changed = bitwise(
                         &mut self_chunk_words[0..num_words],
                         &other_chunk_words[0..num_words],
                         op,
                     );
-                    debug_assert!(has_changed);
-                    *self_chunk_count = count_ones(&self_chunk_words[0..num_words]) as ChunkSize;
-                    if *self_chunk_count == 0 {
-                        *self_chunk = Zeros;
+                    if has_changed {
+                        *self_chunk_count =
+                            count_ones(&self_chunk_words[0..num_words]) as ChunkSize;
+                        if *self_chunk_count == 0 {
+                            *self_chunk = Zeros;
+                        }
+                        changed = true;
                     }
-                    changed = true;
                 }
             }
         }
@@ -926,26 +963,32 @@ impl<T: Idx> BitRelations<ChunkedBitSet<T>> for ChunkedBitSet<T> {
                     // See `ChunkedBitSet::union` for details on what is happening here.
                     let num_words = num_words(chunk_domain_size as usize);
                     let op = |a, b| a & b;
-                    if !bitwise_changes(
-                        &self_chunk_words[0..num_words],
-                        &other_chunk_words[0..num_words],
-                        op,
-                    ) {
+
+                    let Some(self_chunk_words) =
+                        make_chunk_words_mut_for_change(self_chunk_words, |self_chunk_words| {
+                            bitwise_changes(
+                                &self_chunk_words[0..num_words],
+                                &other_chunk_words[0..num_words],
+                                op,
+                            )
+                        })
+                    else {
                         continue;
-                    }
+                    };
 
-                    let self_chunk_words = Rc::make_mut(self_chunk_words);
                     let has_changed = bitwise(
                         &mut self_chunk_words[0..num_words],
                         &other_chunk_words[0..num_words],
                         op,
                     );
-                    debug_assert!(has_changed);
-                    *self_chunk_count = count_ones(&self_chunk_words[0..num_words]) as ChunkSize;
-                    if *self_chunk_count == 0 {
-                        *self_chunk = Zeros;
+                    if has_changed {
+                        *self_chunk_count =
+                            count_ones(&self_chunk_words[0..num_words]) as ChunkSize;
+                        if *self_chunk_count == 0 {
+                            *self_chunk = Zeros;
+                        }
+                        changed = true;
                     }
-                    changed = true;
                 }
             }
         }
diff --git a/compiler/rustc_index/src/bit_set/tests.rs b/compiler/rustc_index/src/bit_set/tests.rs
@@ -298,6 +298,108 @@ fn chunked_bitset() {
     b10000b.assert_valid();
 }
 
+/// Creates a new `ChunkedBitSet` containing all `i` for which `fill_fn(i)` is true.
+fn chunked_fill_with(domain_size: usize, fill_fn: impl Fn(usize) -> bool) -> ChunkedBitSet<usize> {
+    let mut bitset = ChunkedBitSet::new_empty(domain_size);
+    for i in 0..domain_size {
+        if fill_fn(i) {
+            bitset.insert(i);
+        }
+    }
+    bitset
+}
+
+/// Asserts that for each `i` in `0..bitset.domain_size()`, `bitset.contains(i) == expected_fn(i)`.
+#[track_caller]
+fn assert_chunked_contains(bitset: &ChunkedBitSet<usize>, expected_fn: impl Fn(usize) -> bool) {
+    for i in 0..bitset.domain_size() {
+        let expected = expected_fn(i);
+        assert_eq!(bitset.contains(i), expected, "i = {i}");
+    }
+}
+
+#[test]
+fn chunked_bulk_ops() {
+    struct ChunkedBulkOp {
+        name: &'static str,
+        op_fn: fn(&mut ChunkedBitSet<usize>, &ChunkedBitSet<usize>) -> bool,
+        spec_fn: fn(fn(usize) -> bool, fn(usize) -> bool, usize) -> bool,
+    }
+    let ops = &[
+        ChunkedBulkOp {
+            name: "union",
+            op_fn: ChunkedBitSet::union,
+            spec_fn: |fizz, buzz, i| fizz(i) || buzz(i),
+        },
+        ChunkedBulkOp {
+            name: "subtract",
+            op_fn: ChunkedBitSet::subtract,
+            spec_fn: |fizz, buzz, i| fizz(i) && !buzz(i),
+        },
+        ChunkedBulkOp {
+            name: "intersect",
+            op_fn: ChunkedBitSet::intersect,
+            spec_fn: |fizz, buzz, i| fizz(i) && buzz(i),
+        },
+    ];
+
+    let domain_sizes = [
+        CHUNK_BITS / 7, // Smaller than a full chunk.
+        CHUNK_BITS,
+        (CHUNK_BITS + CHUNK_BITS / 7), // Larger than a full chunk.
+    ];
+
+    for ChunkedBulkOp { name, op_fn, spec_fn } in ops {
+        for domain_size in domain_sizes {
+            // If true, use identical values for LHS and RHS, to test "fizz op fizz".
+            // If false, use different values, to test "fizz op buzz".
+            for identical in [false, true] {
+                // If true, make a clone of LHS before doing the op.
+                // This covers optimizations that depend on whether chunk words are shared or not.
+                for unique in [false, true] {
+                    // Print the current test case, so that we can see which one failed.
+                    println!(
+                        "Testing op={name}, domain_size={domain_size}, identical={identical}, unique={unique} ..."
+                    );
+
+                    let fizz_fn = |i| i % 3 == 0;
+                    let buzz_fn = if identical { fizz_fn } else { |i| i % 5 == 0 };
+
+                    // Check that `fizz op buzz` gives the expected results.
+                    chunked_test_inner(domain_size, unique, fizz_fn, buzz_fn, op_fn, |i| {
+                        spec_fn(fizz_fn, buzz_fn, i)
+                    });
+                }
+            }
+        }
+    }
+}
+
+fn chunked_test_inner(
+    domain_size: usize,
+    unique: bool,
+    fizz_fn: impl Fn(usize) -> bool + Copy,
+    buzz_fn: impl Fn(usize) -> bool + Copy,
+    op_fn: impl Fn(&mut ChunkedBitSet<usize>, &ChunkedBitSet<usize>) -> bool,
+    expected_fn: impl Fn(usize) -> bool + Copy,
+) {
+    // Create two bitsets, "fizz" (LHS) and "buzz" (RHS).
+    let mut fizz = chunked_fill_with(domain_size, fizz_fn);
+    let buzz = chunked_fill_with(domain_size, buzz_fn);
+
+    // If requested, clone `fizz` so that its word Rcs are not uniquely-owned.
+    let _cloned = (!unique).then(|| fizz.clone());
+
+    // Perform the op (e.g. union/subtract/intersect), and verify that the
+    // mutated LHS contains exactly the expected values.
+    let changed = op_fn(&mut fizz, &buzz);
+    assert_chunked_contains(&fizz, expected_fn);
+
+    // Verify that the "changed" return value is correct.
+    let should_change = (0..domain_size).any(|i| fizz_fn(i) != expected_fn(i));
+    assert_eq!(changed, should_change);
+}
+
 fn with_elements_chunked(elements: &[usize], domain_size: usize) -> ChunkedBitSet<usize> {
     let mut s = ChunkedBitSet::new_empty(domain_size);
     for &e in elements {
diff --git a/tests/ui/nll/polonius/nll-problem-case-3-issue-54663.nll.stderr b/tests/ui/nll/polonius/nll-problem-case-3-issue-54663.nll.stderr
@@ -1,5 +1,5 @@
 error[E0499]: cannot borrow `*x` as mutable more than once at a time
-  --> $DIR/nll-problem-case-3-issue-54663.rs:20:9
+  --> $DIR/nll-problem-case-3-issue-54663.rs:22:9
    |
 LL | fn foo(x: &mut u8) -> Option<&u8> {
    |           - let's call the lifetime of this reference `'1`
@@ -11,6 +11,35 @@ LL |     }
 LL |     bar(x)
    |         ^ second mutable borrow occurs here
 
-error: aborting due to 1 previous error
+error[E0502]: cannot borrow `*self_words` as immutable because it is also borrowed as mutable
+  --> $DIR/nll-problem-case-3-issue-54663.rs:38:22
+   |
+LL | fn make_chunk_words_mut_for_change<'a>(
+   |                                    -- lifetime `'a` defined here
+...
+LL |     if let Some(words) = Rc::get_mut(self_words) {
+   |                                      ---------- mutable borrow occurs here
+LL |         return Some(words);
+   |                ----------- returning this value requires that `*self_words` is borrowed for `'a`
+...
+LL |     if !would_change(self_words) {
+   |                      ^^^^^^^^^^ immutable borrow occurs here
+
+error[E0499]: cannot borrow `*self_words` as mutable more than once at a time
+  --> $DIR/nll-problem-case-3-issue-54663.rs:42:23
+   |
+LL | fn make_chunk_words_mut_for_change<'a>(
+   |                                    -- lifetime `'a` defined here
+...
+LL |     if let Some(words) = Rc::get_mut(self_words) {
+   |                                      ---------- first mutable borrow occurs here
+LL |         return Some(words);
+   |                ----------- returning this value requires that `*self_words` is borrowed for `'a`
+...
+LL |     Some(Rc::make_mut(self_words))
+   |                       ^^^^^^^^^^ second mutable borrow occurs here
+
+error: aborting due to 3 previous errors
 
-For more information about this error, try `rustc --explain E0499`.
+Some errors have detailed explanations: E0499, E0502.
+For more information about an error, try `rustc --explain E0499`.
diff --git a/tests/ui/nll/polonius/nll-problem-case-3-issue-54663.rs b/tests/ui/nll/polonius/nll-problem-case-3-issue-54663.rs
@@ -13,6 +13,8 @@
 //@ [legacy] check-pass
 //@ [legacy] compile-flags: -Z polonius=legacy
 
+use std::rc::Rc;
+
 fn foo(x: &mut u8) -> Option<&u8> {
     if let Some(y) = bar(x) {
         return Some(y);
@@ -23,3 +25,19 @@ fn foo(x: &mut u8) -> Option<&u8> {
 fn bar(x: &mut u8) -> Option<&u8> {
     Some(x)
 }
+
+// Adapted from the compiler's `make_chunk_words_mut_for_change` in `bit_set.rs`.
+fn make_chunk_words_mut_for_change<'a>(
+    self_words: &'a mut Rc<[u64; 32]>,
+    would_change: impl Fn(&[u64; 32]) -> bool,
+) -> Option<&'a mut [u64; 32]> {
+    if let Some(words) = Rc::get_mut(self_words) {
+        return Some(words);
+    };
+
+    if !would_change(self_words) {
+        return None;
+    }
+
+    Some(Rc::make_mut(self_words))
+}
