Replace visit_waiters with abstracted_waiters_of

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);