perf: Iterate dataflow blocks in traversal order [pathological loop/await CFGs -52%/-82%; primary ~neutral]

compiler/rustc_data_structures/src/lib.rs

@@ -87,7 +87,6 @@ pub mod unhash;
 pub mod union_find;
 pub mod unord;
 pub mod vec_cache;
-pub mod work_queue;
 
 mod atomic_ref;
 

compiler/rustc_index/src/bit_set.rs

@@ -275,6 +275,34 @@ impl<T: Idx> DenseBitSet<T> {
         BitIter::new(&self.words)
     }
 
+    /// Returns the smallest set bit within `range`, if any.
+    pub fn first_set_in(&self, range: impl RangeBounds<T>) -> Option<T> {
+        let (start, end) = inclusive_start_end(range, self.domain_size)?;
+        let (start_word_index, start_mask) = word_index_and_mask(start);
+        let (end_word_index, _) = word_index_and_mask(end);
+
+        // Mask off the bits before `start`; they are outside the range.
+        let start_word = self.words[start_word_index] & !(start_mask - 1);
+        if start_word != 0 {
+            let pos = min_bit(start_word) + WORD_BITS * start_word_index;
+            // A hit beyond `end` can only occur in the last word of the range, so there is nothing
+            // left to search.
+            return if pos <= end { Some(T::new(pos)) } else { None };
+        }
+
+        if let Some(offset) =
+            self.words[start_word_index + 1..=end_word_index].iter().position(|&w| w != 0)
+        {
+            let word_idx = start_word_index + 1 + offset;
+            let pos = min_bit(self.words[word_idx]) + WORD_BITS * word_idx;
+            if pos <= end {
+                return Some(T::new(pos));
+            }
+        }
+
+        None
+    }
+
     pub fn last_set_in(&self, range: impl RangeBounds<T>) -> Option<T> {
         let (start, end) = inclusive_start_end(range, self.domain_size)?;
         let (start_word_index, _) = word_index_and_mask(start);
@@ -1748,6 +1776,11 @@ fn clear_excess_bits_in_final_word(domain_size: usize, words: &mut [Word]) {
     }
 }
 
+#[inline]
+fn min_bit(word: Word) -> usize {
+    word.trailing_zeros() as usize
+}
+
 #[inline]
 fn max_bit(word: Word) -> usize {
     WORD_BITS - 1 - word.leading_zeros() as usize

compiler/rustc_index/src/bit_set/tests.rs

@@ -791,6 +791,49 @@ fn dense_insert_range() {
     }
 }
 
+#[test]
+fn dense_first_set_in() {
+    fn easy(set: &DenseBitSet<usize>, needle: impl RangeBounds<usize>) -> Option<usize> {
+        set.iter().find(|e| needle.contains(e))
+    }
+
+    #[track_caller]
+    fn cmp(set: &DenseBitSet<usize>, needle: impl RangeBounds<usize> + Clone + std::fmt::Debug) {
+        assert_eq!(
+            set.first_set_in(needle.clone()),
+            easy(set, needle.clone()),
+            "{:?} in {:?}",
+            needle,
+            set
+        );
+    }
+    let mut set = DenseBitSet::new_empty(300);
+    cmp(&set, 50..=50);
+    set.insert(WORD_BITS);
+    cmp(&set, WORD_BITS..=WORD_BITS);
+    set.insert(WORD_BITS - 1);
+    cmp(&set, 0..=WORD_BITS - 1);
+    cmp(&set, 0..=5);
+    cmp(&set, 10..100);
+    set.insert(100);
+    cmp(&set, 100..110);
+    cmp(&set, 99..100);
+    cmp(&set, 99..=100);
+
+    for i in 0..=WORD_BITS * 2 {
+        for j in i..=WORD_BITS * 2 {
+            for k in 0..WORD_BITS * 2 {
+                let mut set = DenseBitSet::new_empty(300);
+                cmp(&set, i..j);
+                cmp(&set, i..=j);
+                set.insert(k);
+                cmp(&set, i..j);
+                cmp(&set, i..=j);
+            }
+        }
+    }
+}
+
 #[test]
 fn dense_last_set_before() {
     fn easy(set: &DenseBitSet<usize>, needle: impl RangeBounds<usize>) -> Option<usize> {

compiler/rustc_mir_dataflow/src/framework/mod.rs

@@ -34,11 +34,10 @@
 
 use std::cmp::Ordering;
 
-use rustc_data_structures::work_queue::WorkQueue;
 use rustc_index::bit_set::{DenseBitSet, MixedBitSet};
 use rustc_index::{Idx, IndexVec};
 use rustc_middle::bug;
-use rustc_middle::mir::{self, BasicBlock, CallReturnPlaces, Location, TerminatorEdges, traversal};
+use rustc_middle::mir::{self, BasicBlock, CallReturnPlaces, Location, TerminatorEdges};
 use rustc_middle::ty::TyCtxt;
 use tracing::error;
 
@@ -260,19 +259,7 @@ pub trait Analysis<'tcx> {
             bug!("`initialize_start_block` is not yet supported for backward dataflow analyses");
         }
 
-        let mut dirty_queue: WorkQueue<BasicBlock> = WorkQueue::with_none(body.basic_blocks.len());
-
-        if Self::Direction::IS_FORWARD {
-            for (bb, _) in traversal::reverse_postorder(body) {
-                dirty_queue.insert(bb);
-            }
-        } else {
-            // Reverse post-order on the reverse CFG may generate a better iteration order for
-            // backward dataflow analyses, but probably not enough to matter.
-            for (bb, _) in traversal::postorder(body) {
-                dirty_queue.insert(bb);
-            }
-        }
+        let mut dirty_queue = DirtyQueue::new(body, Self::Direction::IS_BACKWARD);
 
         // `state` is not actually used between iterations;
         // this is just an optimization to avoid reallocating
@@ -311,6 +298,113 @@ pub trait Analysis<'tcx> {
     }
 }
 
+/// Placeholder position for a block that is not part of the visit order, i.e. one that is
+/// unreachable from the start block.
+const NO_POSITION: u32 = u32::MAX;
+
+/// Ditto, for a block that is currently queued in `DirtyQueue::unreachable`.
+const NO_POSITION_DIRTY: u32 = u32::MAX - 1;
+
+/// The set of blocks whose entry state may still change, used by [`Analysis::iterate_to_fixpoint`].
+///
+/// The queue repeatedly sweeps the analysis's visit order (reverse postorder for forward analyses,
+/// postorder for backward ones) and yields the dirty blocks it passes, instead of yielding blocks
+/// in insertion order. A sweep computes every input of a block except those coming from loop back
+/// edges before the block itself, so the number of sweeps needed is bounded by the loop nesting
+/// depth rather than by the size of the CFG. Insertion order gives no such bound: it interleaves
+/// the loops of the CFG arbitrarily, and on a function with many loops in sequence (a coroutine
+/// with many `.await`s, say) the number of visits per block grows with the number of loops.
+struct DirtyQueue<'a> {
+    /// The blocks to visit, in reverse postorder.
+    ///
+    /// Postorder is this same sequence walked back to front, so backward analyses reuse it in
+    /// reverse rather than storing an order of their own. (Reverse postorder on the reverse CFG
+    /// may be a better order for them, but probably not enough to matter.)
+    rpo: &'a [BasicBlock],
+
+    /// Whether `rpo` is to be walked back to front.
+    backward: bool,
+
+    /// The position of each block within the visit order, or `NO_POSITION`/`NO_POSITION_DIRTY` for
+    /// blocks that have none.
+    positions: IndexVec<BasicBlock, u32>,
+
+    /// The dirty blocks, keyed by their position in the visit order.
+    dirty: DenseBitSet<usize>,
+
+    /// How far the current sweep has gotten.
+    cursor: usize,
+
+    /// Dirty blocks that have no position. Backward analyses reach these, as an unreachable block
+    /// can be a predecessor of a reachable one.
+    unreachable: Vec<BasicBlock>,
+}
+
+impl<'a> DirtyQueue<'a> {
+    /// Creates a queue in which every block of the visit order is dirty.
+    fn new(body: &'a mir::Body<'_>, backward: bool) -> Self {
+        let rpo = body.basic_blocks.reverse_postorder();
+        let mut positions = IndexVec::from_elem_n(NO_POSITION, body.basic_blocks.len());
+        for (i, &bb) in rpo.iter().enumerate() {
+            positions[bb] = if backward { (rpo.len() - 1 - i) as u32 } else { i as u32 };
+        }
+
+        DirtyQueue {
+            rpo,
+            backward,
+            positions,
+            dirty: DenseBitSet::new_filled(rpo.len()),
+            cursor: 0,
+            unreachable: Vec::new(),
+        }
+    }
+
+    #[inline]
+    fn block_at(&self, pos: usize) -> BasicBlock {
+        if self.backward { self.rpo[self.rpo.len() - 1 - pos] } else { self.rpo[pos] }
+    }
+
+    #[inline]
+    fn insert(&mut self, bb: BasicBlock) {
+        match self.positions[bb] {
+            NO_POSITION => {
+                self.positions[bb] = NO_POSITION_DIRTY;
+                self.unreachable.push(bb);
+            }
+            NO_POSITION_DIRTY => {}
+            pos => {
+                self.dirty.insert(pos as usize);
+            }
+        }
+    }
+
+    /// Yields the next dirty block of the current sweep, starting a new sweep once the visit order
+    /// runs out.
+    ///
+    /// Restarting the sweep at a block dirtied behind the cursor as soon as that happens would be a
+    /// pessimization, not an improvement: a block joined from many back edges, such as the resume
+    /// dispatch of a coroutine, would then be recomputed once per back edge instead of once per
+    /// sweep.
+    #[inline]
+    fn pop(&mut self) -> Option<BasicBlock> {
+        let next = self
+            .dirty
+            .first_set_in(self.cursor..)
+            .or_else(|| if self.cursor > 0 { self.dirty.first_set_in(0..) } else { None });
+
+        if let Some(pos) = next {
+            self.dirty.remove(pos);
+            self.cursor = pos;
+            Some(self.block_at(pos))
+        } else if let Some(bb) = self.unreachable.pop() {
+            self.positions[bb] = NO_POSITION;
+            Some(bb)
+        } else {
+            None
+        }
+    }
+}
+
 #[derive(Debug, Clone, Copy)]
 pub enum SwitchTargetIndex {
     // Index of a normal switch target.