Replace visit_waiters with waiters_of

Author: Zoxc

compiler/rustc_query_impl/src/job.rs

@@ -113,38 +113,41 @@ 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);
+
+struct Waiter {
+    span: Span,
+    /// The query which we are waiting from, if none the waiter is from a compiler root.
+    query: Option<QueryJobId>,
+    resumable: Option<ResumableWaiterLocation>,
+}
 
-    // Visit the explicit waiters which use condvars and are resumable
+/// 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 waiters_of(job_map: &QueryJobMap<'_>, query: QueryJobId) -> Vec<Waiter> {
+    let mut result = Vec::new();
+
+    // Add the parent which is a non-resumable waiter since it's on the same stack
+    result.push(Waiter {
+        span: job_map.span_of(query),
+        query: 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.info.lock().waiters.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(Waiter {
+                span: waiter.span,
+                query: waiter.query,
+                resumable: Some((query, i)),
+            });
         }
     }
 
-    ControlFlow::Continue(())
+    result
 }
 
 /// Look for query cycles by doing a depth first search starting at `query`.
@@ -157,7 +160,7 @@ 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) {
             // We detected a query cycle, fix up the initial span and return Some
@@ -176,16 +179,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 waiter in waiters_of(job_map, query) {
+        let Some(waiter_query) = waiter.query else {
+            // Skip waiters which are not queries
+            continue;
+        };
+        if let ControlFlow::Break(maybe_waiter) =
+            cycle_check(job_map, waiter_query, waiter.span, stack, visited)
+        {
+            // Return the resumable waiter in `waiter.resumable` if present
+            return ControlFlow::Break(waiter.resumable.or(maybe_waiter));
+        }
     }
 
-    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)
@@ -195,18 +205,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 waiter in waiters_of(job_map, query) {
+        match waiter.query {
+            // This query is connected to the root
+            None => return true,
+            Some(waiter) => {
+                if connected_to_root(job_map, waiter, 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`.
@@ -252,29 +270,29 @@ fn remove_cycle<'tcx>(
         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 found_waiter = None;
+
+                // Find a direct waiter who leads to the root
+                for waiter in waiters_of(job_map, query_in_cycle) {
+                    let Some(waiter_query) = waiter.query 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, waiter_query, &mut visited) {
+                        found_waiter = Some((waiter.span, waiter_query));
+                        entrypoint = true;
+                        break;
+                    }
                 }
+
+                entrypoint.then_some(EntryPoint { query_in_cycle, waiter: found_waiter })
             })
             .collect::<Vec<EntryPoint>>();