Auto merge of #96578 - tmiasko:chunked-bit-set-fmt, r=nnethercote

Fix -Zdump-mir-dataflow by implementing DebugWithContext for ChunkedBitSet

`DebugWithContext` is used to format changes to dataflow state along with MIR
in graphviz dot files. In the case of `ChunkedBitSet` it was left unimplemented,
so attempts to use `-Zdump-mir-dataflow -Zdump-mir=all` resulted in an ICE:

> thread 'rustc' panicked at 'not implemented: implement when/if needed',

Provide the missing implementation.

r? `@nnethercote`

3 files changed, 144 insertions, 36 deletions

diff --git a/compiler/rustc_index/src/bit_set.rs b/compiler/rustc_index/src/bit_set.rs
index 33a8d6c11ff..059755a743b 100644
--- a/compiler/rustc_index/src/bit_set.rs
+++ b/compiler/rustc_index/src/bit_set.rs
@@ -479,6 +479,11 @@ impl<T: Idx> ChunkedBitSet<T> {
         }
     }
 
+    #[inline]
+    pub fn iter(&self) -> ChunkedBitIter<'_, T> {
+        ChunkedBitIter::new(self)
+    }
+
     /// Insert `elem`. Returns whether the set has changed.
     pub fn insert(&mut self, elem: T) -> bool {
         assert!(elem.index() < self.domain_size);
@@ -697,6 +702,49 @@ impl<T> Clone for ChunkedBitSet<T> {
     }
 }
 
+pub struct ChunkedBitIter<'a, T: Idx> {
+    index: usize,
+    bitset: &'a ChunkedBitSet<T>,
+}
+
+impl<'a, T: Idx> ChunkedBitIter<'a, T> {
+    #[inline]
+    fn new(bitset: &'a ChunkedBitSet<T>) -> ChunkedBitIter<'a, T> {
+        ChunkedBitIter { index: 0, bitset }
+    }
+}
+
+impl<'a, T: Idx> Iterator for ChunkedBitIter<'a, T> {
+    type Item = T;
+    fn next(&mut self) -> Option<T> {
+        while self.index < self.bitset.domain_size() {
+            let elem = T::new(self.index);
+            let chunk = &self.bitset.chunks[chunk_index(elem)];
+            match &chunk {
+                Zeros(chunk_domain_size) => {
+                    self.index += *chunk_domain_size as usize;
+                }
+                Ones(_chunk_domain_size) => {
+                    self.index += 1;
+                    return Some(elem);
+                }
+                Mixed(_chunk_domain_size, _, words) => loop {
+                    let elem = T::new(self.index);
+                    self.index += 1;
+                    let (word_index, mask) = chunk_word_index_and_mask(elem);
+                    if (words[word_index] & mask) != 0 {
+                        return Some(elem);
+                    }
+                    if self.index % CHUNK_BITS == 0 {
+                        break;
+                    }
+                },
+            }
+        }
+        None
+    }
+}
+
 impl Chunk {
     #[cfg(test)]
     fn assert_valid(&self) {
diff --git a/compiler/rustc_index/src/bit_set/tests.rs b/compiler/rustc_index/src/bit_set/tests.rs
index eec7dab5189..cfc891e97a3 100644
--- a/compiler/rustc_index/src/bit_set/tests.rs
+++ b/compiler/rustc_index/src/bit_set/tests.rs
@@ -343,6 +343,40 @@ fn chunked_bitset() {
 }
 
 #[test]
+fn chunked_bitset_iter() {
+    fn with_elements(elements: &[usize], domain_size: usize) -> ChunkedBitSet<usize> {
+        let mut s = ChunkedBitSet::new_empty(domain_size);
+        for &e in elements {
+            s.insert(e);
+        }
+        s
+    }
+
+    // Empty
+    let vec: Vec<usize> = Vec::new();
+    let bit = with_elements(&vec, 9000);
+    assert_eq!(vec, bit.iter().collect::<Vec<_>>());
+
+    // Filled
+    let n = 10000;
+    let vec: Vec<usize> = (0..n).collect();
+    let bit = with_elements(&vec, n);
+    assert_eq!(vec, bit.iter().collect::<Vec<_>>());
+
+    // Filled with trailing zeros
+    let n = 10000;
+    let vec: Vec<usize> = (0..n).collect();
+    let bit = with_elements(&vec, 2 * n);
+    assert_eq!(vec, bit.iter().collect::<Vec<_>>());
+
+    // Mixed
+    let n = 12345;
+    let vec: Vec<usize> = vec![0, 1, 2, 2010, 2047, 2099, 6000, 6002, 6004];
+    let bit = with_elements(&vec, n);
+    assert_eq!(vec, bit.iter().collect::<Vec<_>>());
+}
+
+#[test]
 fn grow() {
     let mut set: GrowableBitSet<usize> = GrowableBitSet::with_capacity(65);
     for index in 0..65 {
diff --git a/compiler/rustc_mir_dataflow/src/framework/fmt.rs b/compiler/rustc_mir_dataflow/src/framework/fmt.rs
index 99735673f4d..209e6f7ac9f 100644
--- a/compiler/rustc_mir_dataflow/src/framework/fmt.rs
+++ b/compiler/rustc_mir_dataflow/src/framework/fmt.rs
@@ -93,57 +93,83 @@ where
             };
         }
 
-        let mut first = true;
-        for idx in set_in_self.iter() {
-            let delim = if first {
-                "\u{001f}+"
-            } else if f.alternate() {
-                "\n\u{001f}+"
-            } else {
-                ", "
-            };
+        fmt_diff(&set_in_self, &cleared_in_self, ctxt, f)
+    }
+}
 
-            write!(f, "{}", delim)?;
-            idx.fmt_with(ctxt, f)?;
-            first = false;
-        }
+impl<T, C> DebugWithContext<C> for ChunkedBitSet<T>
+where
+    T: Idx + DebugWithContext<C>,
+{
+    fn fmt_with(&self, ctxt: &C, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_set().entries(self.iter().map(|i| DebugWithAdapter { this: i, ctxt })).finish()
+    }
 
-        if !f.alternate() {
-            first = true;
-            if !set_in_self.is_empty() && !cleared_in_self.is_empty() {
-                write!(f, "\t")?;
-            }
-        }
+    fn fmt_diff_with(&self, old: &Self, ctxt: &C, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let size = self.domain_size();
+        assert_eq!(size, old.domain_size());
 
-        for idx in cleared_in_self.iter() {
-            let delim = if first {
-                "\u{001f}-"
-            } else if f.alternate() {
-                "\n\u{001f}-"
-            } else {
-                ", "
-            };
+        let mut set_in_self = HybridBitSet::new_empty(size);
+        let mut cleared_in_self = HybridBitSet::new_empty(size);
 
-            write!(f, "{}", delim)?;
-            idx.fmt_with(ctxt, f)?;
-            first = false;
+        for i in (0..size).map(T::new) {
+            match (self.contains(i), old.contains(i)) {
+                (true, false) => set_in_self.insert(i),
+                (false, true) => cleared_in_self.insert(i),
+                _ => continue,
+            };
         }
 
-        Ok(())
+        fmt_diff(&set_in_self, &cleared_in_self, ctxt, f)
     }
 }
 
-impl<T, C> DebugWithContext<C> for ChunkedBitSet<T>
+fn fmt_diff<T, C>(
+    inserted: &HybridBitSet<T>,
+    removed: &HybridBitSet<T>,
+    ctxt: &C,
+    f: &mut fmt::Formatter<'_>,
+) -> fmt::Result
 where
     T: Idx + DebugWithContext<C>,
 {
-    fn fmt_with(&self, _ctxt: &C, _f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        unimplemented!("implement when/if needed");
+    let mut first = true;
+    for idx in inserted.iter() {
+        let delim = if first {
+            "\u{001f}+"
+        } else if f.alternate() {
+            "\n\u{001f}+"
+        } else {
+            ", "
+        };
+
+        write!(f, "{}", delim)?;
+        idx.fmt_with(ctxt, f)?;
+        first = false;
+    }
+
+    if !f.alternate() {
+        first = true;
+        if !inserted.is_empty() && !removed.is_empty() {
+            write!(f, "\t")?;
+        }
     }
 
-    fn fmt_diff_with(&self, _old: &Self, _ctxt: &C, _f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        unimplemented!("implement when/if needed");
+    for idx in removed.iter() {
+        let delim = if first {
+            "\u{001f}-"
+        } else if f.alternate() {
+            "\n\u{001f}-"
+        } else {
+            ", "
+        };
+
+        write!(f, "{}", delim)?;
+        idx.fmt_with(ctxt, f)?;
+        first = false;
     }
+
+    Ok(())
 }
 
 impl<T, C> DebugWithContext<C> for &'_ T