Rollup merge of #153376 - Zoxc:cycle-waiters-iter, r=nnethercote

Replace `visit_waiters` with `abstracted_waiters_of`

This replaces `visit_waiters` which uses closures to visit waiters with `abstracted_waiters_of` which returns a list of waiters. This makes the control flow of their users a bit clearer.

Additionally `abstracted_waiters_of` includes non-query waiters unlike `visit_waiters`. This means that `connected_to_root` now considers resumable waiters from the compiler root as being connected to root, while previously only non-resumable waiters counted. This can result in some additional entry points being found. Similarly in the loop detecting entry points we now consider queries in the cycle with direct resumable waiters from the compiler root as being entry points.

When looking for entry points we now look at waiters until we found a query to populate the `EntryPoint.waiter` field instead of stopping when we determined it to be an entry point.

cc @petrochenkov @nnethercote

Author: Zalathar

compiler/rustc_middle/src/query/job.rs

@@ -65,7 +65,7 @@ impl<'tcx> QueryJob<'tcx> {
 
 #[derive(Debug)]
 pub struct QueryWaiter<'tcx> {
-    pub query: Option<QueryJobId>,
+    pub parent: Option<QueryJobId>,
     pub condvar: Condvar,
     pub span: Span,
     pub cycle: Mutex<Option<CycleError<'tcx>>>,
@@ -94,8 +94,12 @@ impl<'tcx> QueryLatch<'tcx> {
             return Ok(()); // already complete
         };
 
-        let waiter =
-            Arc::new(QueryWaiter { query, span, cycle: Mutex::new(None), condvar: Condvar::new() });
+        let waiter = Arc::new(QueryWaiter {
+            parent: query,
+            span,
+            cycle: Mutex::new(None),
+            condvar: Condvar::new(),
+        });
 
         // We push the waiter on to the `waiters` list. It can be accessed inside
         // the `wait` call below, by 1) the `set` method or 2) by deadlock detection.

compiler/rustc_query_impl/src/job.rs

@@ -111,38 +111,44 @@ pub(crate) fn find_dep_kind_root<'tcx>(
     last_layout
 }
 
-/// A resumable waiter of a query. The usize is the index into waiters in the query's latch
-type Waiter = (QueryJobId, usize);
-
-/// Visits all the non-resumable and resumable waiters of a query.
-/// Only waiters in a query are visited.
-/// `visit` is called for every waiter and is passed a query waiting on `query`
-/// and a span indicating the reason the query waited on `query`.
-/// If `visit` returns `Break`, this function also returns `Break`,
-/// and if all `visit` calls returns `Continue` it also returns `Continue`.
-/// For visits of non-resumable waiters it returns the return value of `visit`.
-/// For visits of resumable waiters it returns information required to resume that waiter.
-fn visit_waiters<'tcx>(
-    job_map: &QueryJobMap<'tcx>,
-    query: QueryJobId,
-    mut visit: impl FnMut(Span, QueryJobId) -> ControlFlow<Option<Waiter>>,
-) -> ControlFlow<Option<Waiter>> {
-    // Visit the parent query which is a non-resumable waiter since it's on the same stack
-    if let Some(parent) = job_map.parent_of(query) {
-        visit(job_map.span_of(query), parent)?;
-    }
+/// The locaton of a resumable waiter. The usize is the index into waiters in the query's latch.
+/// We'll use this to remove the waiter using `QueryLatch::extract_waiter` if we're waking it up.
+type ResumableWaiterLocation = (QueryJobId, usize);
+
+/// This abstracts over non-resumable waiters which are found in `QueryJob`'s `parent` field
+/// and resumable waiters are in `latch` field.
+struct AbstractedWaiter {
+    /// The span corresponding to the reason for why we're waiting on this query.
+    span: Span,
+    /// The query which we are waiting from, if none the waiter is from a compiler root.
+    parent: Option<QueryJobId>,
+    resumable: Option<ResumableWaiterLocation>,
+}
+
+/// Returns all the non-resumable and resumable waiters of a query.
+/// This is used so we can uniformly loop over both non-resumable and resumable waiters.
+fn abstracted_waiters_of(job_map: &QueryJobMap<'_>, query: QueryJobId) -> Vec<AbstractedWaiter> {
+    let mut result = Vec::new();
 
-    // Visit the explicit waiters which use condvars and are resumable
+    // Add the parent which is a non-resumable waiter since it's on the same stack
+    result.push(AbstractedWaiter {
+        span: job_map.span_of(query),
+        parent: job_map.parent_of(query),
+        resumable: None,
+    });
+
+    // Add the explicit waiters which use condvars and are resumable
     if let Some(latch) = job_map.latch_of(query) {
         for (i, waiter) in latch.waiters.lock().as_ref().unwrap().iter().enumerate() {
-            if let Some(waiter_query) = waiter.query {
-                // Return a value which indicates that this waiter can be resumed
-                visit(waiter.span, waiter_query).map_break(|_| Some((query, i)))?;
-            }
+            result.push(AbstractedWaiter {
+                span: waiter.span,
+                parent: waiter.parent,
+                resumable: Some((query, i)),
+            });
         }
     }
 
-    ControlFlow::Continue(())
+    result
 }
 
 /// Look for query cycles by doing a depth first search starting at `query`.
@@ -155,13 +161,13 @@ fn cycle_check<'tcx>(
     span: Span,
     stack: &mut Vec<(Span, QueryJobId)>,
     visited: &mut FxHashSet<QueryJobId>,
-) -> ControlFlow<Option<Waiter>> {
+) -> ControlFlow<Option<ResumableWaiterLocation>> {
     if !visited.insert(query) {
-        return if let Some(p) = stack.iter().position(|q| q.1 == query) {
+        return if let Some(pos) = stack.iter().position(|q| q.1 == query) {
             // We detected a query cycle, fix up the initial span and return Some
 
             // Remove previous stack entries
-            stack.drain(0..p);
+            stack.drain(0..pos);
             // Replace the span for the first query with the cycle cause
             stack[0].0 = span;
             ControlFlow::Break(None)
@@ -174,16 +180,23 @@ fn cycle_check<'tcx>(
     stack.push((span, query));
 
     // Visit all the waiters
-    let r = visit_waiters(job_map, query, |span, successor| {
-        cycle_check(job_map, successor, span, stack, visited)
-    });
-
-    // Remove the entry in our stack if we didn't find a cycle
-    if r.is_continue() {
-        stack.pop();
+    for abstracted_waiter in abstracted_waiters_of(job_map, query) {
+        let Some(parent) = abstracted_waiter.parent else {
+            // Skip waiters which are not queries
+            continue;
+        };
+        if let ControlFlow::Break(maybe_resumable) =
+            cycle_check(job_map, parent, abstracted_waiter.span, stack, visited)
+        {
+            // Return the resumable waiter in `waiter.resumable` if present
+            return ControlFlow::Break(abstracted_waiter.resumable.or(maybe_resumable));
+        }
     }
 
-    r
+    // Remove the entry in our stack since we didn't find a cycle
+    stack.pop();
+
+    ControlFlow::Continue(())
 }
 
 /// Finds out if there's a path to the compiler root (aka. code which isn't in a query)
@@ -193,18 +206,26 @@ fn connected_to_root<'tcx>(
     job_map: &QueryJobMap<'tcx>,
     query: QueryJobId,
     visited: &mut FxHashSet<QueryJobId>,
-) -> ControlFlow<Option<Waiter>> {
+) -> bool {
     // We already visited this or we're deliberately ignoring it
     if !visited.insert(query) {
-        return ControlFlow::Continue(());
+        return false;
     }
 
-    // This query is connected to the root (it has no query parent), return true
-    if job_map.parent_of(query).is_none() {
-        return ControlFlow::Break(None);
+    // Visit all the waiters
+    for abstracted_waiter in abstracted_waiters_of(job_map, query) {
+        match abstracted_waiter.parent {
+            // This query is connected to the root
+            None => return true,
+            Some(parent) => {
+                if connected_to_root(job_map, parent, visited) {
+                    return true;
+                }
+            }
+        }
     }
 
-    visit_waiters(job_map, query, |_, successor| connected_to_root(job_map, successor, visited))
+    false
 }
 
 /// Looks for query cycles starting from the last query in `jobs`.
@@ -220,7 +241,7 @@ fn remove_cycle<'tcx>(
     let mut visited = FxHashSet::default();
     let mut stack = Vec::new();
     // Look for a cycle starting with the last query in `jobs`
-    if let ControlFlow::Break(waiter) =
+    if let ControlFlow::Break(resumable) =
         cycle_check(job_map, jobs.pop().unwrap(), DUMMY_SP, &mut stack, &mut visited)
     {
         // The stack is a vector of pairs of spans and queries; reverse it so that
@@ -242,44 +263,44 @@ fn remove_cycle<'tcx>(
 
         struct EntryPoint {
             query_in_cycle: QueryJobId,
-            waiter: Option<(Span, QueryJobId)>,
+            query_waiting_on_cycle: Option<(Span, QueryJobId)>,
         }
 
         // Find the queries in the cycle which are
         // connected to queries outside the cycle
         let entry_points = stack
             .iter()
             .filter_map(|&(_, query_in_cycle)| {
-                if job_map.parent_of(query_in_cycle).is_none() {
-                    // This query is connected to the root (it has no query parent)
-                    Some(EntryPoint { query_in_cycle, waiter: None })
-                } else {
-                    let mut waiter_on_cycle = None;
-                    // Find a direct waiter who leads to the root
-                    let _ = visit_waiters(job_map, query_in_cycle, |span, waiter| {
-                        // Mark all the other queries in the cycle as already visited
-                        let mut visited = FxHashSet::from_iter(stack.iter().map(|q| q.1));
-
-                        if connected_to_root(job_map, waiter, &mut visited).is_break() {
-                            waiter_on_cycle = Some((span, waiter));
-                            ControlFlow::Break(None)
-                        } else {
-                            ControlFlow::Continue(())
-                        }
-                    });
-
-                    waiter_on_cycle.map(|waiter_on_cycle| EntryPoint {
-                        query_in_cycle,
-                        waiter: Some(waiter_on_cycle),
-                    })
+                let mut entrypoint = false;
+                let mut query_waiting_on_cycle = None;
+
+                // Find a direct waiter who leads to the root
+                for abstracted_waiter in abstracted_waiters_of(job_map, query_in_cycle) {
+                    let Some(parent) = abstracted_waiter.parent else {
+                        // The query in the cycle is directly connected to root.
+                        entrypoint = true;
+                        continue;
+                    };
+
+                    // Mark all the other queries in the cycle as already visited,
+                    // so paths to the root through the cycle itself won't count.
+                    let mut visited = FxHashSet::from_iter(stack.iter().map(|q| q.1));
+
+                    if connected_to_root(job_map, parent, &mut visited) {
+                        query_waiting_on_cycle = Some((abstracted_waiter.span, parent));
+                        entrypoint = true;
+                        break;
+                    }
                 }
+
+                entrypoint.then_some(EntryPoint { query_in_cycle, query_waiting_on_cycle })
             })
             .collect::<Vec<EntryPoint>>();
 
         // Pick an entry point, preferring ones with waiters
         let entry_point = entry_points
             .iter()
-            .find(|entry_point| entry_point.waiter.is_some())
+            .find(|entry_point| entry_point.query_waiting_on_cycle.is_some())
             .unwrap_or(&entry_points[0]);
 
         // Shift the stack so that our entry point is first
@@ -289,7 +310,9 @@ fn remove_cycle<'tcx>(
             stack.rotate_left(pos);
         }
 
-        let usage = entry_point.waiter.map(|(span, job)| (span, job_map.frame_of(job).clone()));
+        let usage = entry_point
+            .query_waiting_on_cycle
+            .map(|(span, job)| (span, job_map.frame_of(job).clone()));
 
         // Create the cycle error
         let error = CycleError {
@@ -300,9 +323,9 @@ fn remove_cycle<'tcx>(
                 .collect(),
         };
 
-        // We unwrap `waiter` here since there must always be one
+        // We unwrap `resumable` here since there must always be one
         // edge which is resumable / waited using a query latch
-        let (waitee_query, waiter_idx) = waiter.unwrap();
+        let (waitee_query, waiter_idx) = resumable.unwrap();
 
         // Extract the waiter we want to resume
         let waiter = job_map.latch_of(waitee_query).unwrap().extract_waiter(waiter_idx);