Move reachability function (#7091)

Split out from #7076

crates/uv-resolver/src/graph_ops.rs

@@ -1,7 +1,10 @@
 use pep508_rs::MarkerTree;
 use petgraph::algo::greedy_feedback_arc_set;
+use petgraph::graph::NodeIndex;
 use petgraph::visit::{EdgeRef, Topo};
 use petgraph::{Directed, Direction, Graph};
+use rustc_hash::FxHashMap;
+use std::collections::hash_map::Entry;
 
 /// A trait for a graph node that can be annotated with a [`MarkerTree`].
 pub(crate) trait Markers {
@@ -79,3 +82,77 @@ pub(crate) fn propagate_markers<T: Markers>(
 
     graph
 }
+
+/// Determine the markers under which a package is reachable in the dependency tree.
+///
+/// The algorithm is a variant of Dijkstra's algorithm for not totally ordered distances:
+/// Whenever we find a shorter distance to a node (a marker that is not a subset of the existing
+/// marker), we re-queue the node and update all its children. This implicitly handles cycles,
+/// whenever we re-reach a node through a cycle the marker we have is a more
+/// specific marker/longer path, so we don't update the node and don't re-queue it.
+pub(crate) fn marker_reachability<T>(
+    graph: &Graph<T, MarkerTree>,
+    fork_markers: &[MarkerTree],
+) -> FxHashMap<NodeIndex, MarkerTree> {
+    // Note that we build including the virtual packages due to how we propagate markers through
+    // the graph, even though we then only read the markers for base packages.
+    let mut reachability = FxHashMap::default();
+
+    // Collect the root nodes.
+    //
+    // Besides the actual virtual root node, virtual dev dependencies packages are also root
+    // nodes since the edges don't cover dev dependencies.
+    let mut queue: Vec<_> = graph
+        .node_indices()
+        .filter(|node_index| {
+            graph
+                .edges_directed(*node_index, Direction::Incoming)
+                .next()
+                .is_none()
+        })
+        .collect();
+
+    // The root nodes are always applicable, unless the user has restricted resolver
+    // environments with `tool.uv.environments`.
+    let root_markers: MarkerTree = if fork_markers.is_empty() {
+        MarkerTree::TRUE
+    } else {
+        fork_markers
+            .iter()
+            .fold(MarkerTree::FALSE, |mut acc, marker| {
+                acc.or(marker.clone());
+                acc
+            })
+    };
+    for root_index in &queue {
+        reachability.insert(*root_index, root_markers.clone());
+    }
+
+    // Propagate all markers through the graph, so that the eventual marker for each node is the
+    // union of the markers of each path we can reach the node by.
+    while let Some(parent_index) = queue.pop() {
+        let marker = reachability[&parent_index].clone();
+        for child_edge in graph.edges_directed(parent_index, Direction::Outgoing) {
+            // The marker for all paths to the child through the parent.
+            let mut child_marker = child_edge.weight().clone();
+            child_marker.and(marker.clone());
+            match reachability.entry(child_edge.target()) {
+                Entry::Occupied(mut existing) => {
+                    // If the marker is a subset of the existing marker (A ⊆ B exactly if
+                    // A ∪ B = A), updating the child wouldn't change child's marker.
+                    child_marker.or(existing.get().clone());
+                    if &child_marker != existing.get() {
+                        existing.insert(child_marker);
+                        queue.push(child_edge.target());
+                    }
+                }
+                Entry::Vacant(vacant) => {
+                    vacant.insert(child_marker.clone());
+                    queue.push(child_edge.target());
+                }
+            }
+        }
+    }
+
+    reachability
+}

crates/uv-resolver/src/resolution/graph.rs

@@ -5,21 +5,19 @@ use distribution_types::{
 use indexmap::IndexSet;
 use pep440_rs::{Version, VersionSpecifier};
 use pep508_rs::{MarkerEnvironment, MarkerTree, MarkerTreeKind, VerbatimUrl};
-use petgraph::prelude::EdgeRef;
 use petgraph::{
     graph::{Graph, NodeIndex},
     Directed, Direction,
 };
 use pypi_types::{HashDigest, ParsedUrlError, Requirement, VerbatimParsedUrl, Yanked};
 use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet};
-use std::collections::hash_map::Entry;
-use std::collections::HashMap;
 use std::fmt::{Display, Formatter};
 use uv_configuration::{Constraints, Overrides};
 use uv_distribution::Metadata;
 use uv_git::GitResolver;
 use uv_normalize::{ExtraName, GroupName, PackageName};
 
+use crate::graph_ops::marker_reachability;
 use crate::pins::FilePins;
 use crate::preferences::Preferences;
 use crate::redirect::url_to_precise;
@@ -211,7 +209,7 @@ impl ResolutionGraph {
                 .collect()
         };
 
-        let reachability = Self::reachability(&petgraph, &fork_markers);
+        let reachability = marker_reachability(&petgraph, &fork_markers);
 
         if matches!(resolution_strategy, ResolutionStrategy::Lowest) {
             report_missing_lower_bounds(&petgraph, &mut diagnostics);
@@ -231,80 +229,6 @@ impl ResolutionGraph {
         })
     }
 
-    /// Determine the markers under which a package is reachable in the dependency tree.
-    ///
-    /// The algorithm is a variant of Dijkstra's algorithm for not totally ordered distances:
-    /// Whenever we find a shorter distance to a node (a marker that is not a subset of the existing
-    /// marker), we re-queue the node and update all its children. This implicitly handles cycles,
-    /// whenever we re-reach a node through a cycle the marker we have is a more
-    /// specific marker/longer path, so we don't update the node and don't re-queue it.
-    fn reachability(
-        petgraph: &Graph<ResolutionGraphNode, MarkerTree>,
-        fork_markers: &[MarkerTree],
-    ) -> HashMap<NodeIndex, MarkerTree, FxBuildHasher> {
-        // Note that we build including the virtual packages due to how we propagate markers through
-        // the graph, even though we then only read the markers for base packages.
-        let mut reachability = FxHashMap::default();
-
-        // Collect the root nodes.
-        //
-        // Besides the actual virtual root node, virtual dev dependencies packages are also root
-        // nodes since the edges don't cover dev dependencies.
-        let mut queue: Vec<_> = petgraph
-            .node_indices()
-            .filter(|node_index| {
-                petgraph
-                    .edges_directed(*node_index, Direction::Incoming)
-                    .next()
-                    .is_none()
-            })
-            .collect();
-
-        // The root nodes are always applicable, unless the user has restricted resolver
-        // environments with `tool.uv.environments`.
-        let root_markers: MarkerTree = if fork_markers.is_empty() {
-            MarkerTree::TRUE
-        } else {
-            fork_markers
-                .iter()
-                .fold(MarkerTree::FALSE, |mut acc, marker| {
-                    acc.or(marker.clone());
-                    acc
-                })
-        };
-        for root_index in &queue {
-            reachability.insert(*root_index, root_markers.clone());
-        }
-
-        // Propagate all markers through the graph, so that the eventual marker for each node is the
-        // union of the markers of each path we can reach the node by.
-        while let Some(parent_index) = queue.pop() {
-            let marker = reachability[&parent_index].clone();
-            for child_edge in petgraph.edges_directed(parent_index, Direction::Outgoing) {
-                // The marker for all paths to the child through the parent.
-                let mut child_marker = child_edge.weight().clone();
-                child_marker.and(marker.clone());
-                match reachability.entry(child_edge.target()) {
-                    Entry::Occupied(mut existing) => {
-                        // If the marker is a subset of the existing marker (A ⊆ B exactly if
-                        // A ∪ B = A), updating the child wouldn't change child's marker.
-                        child_marker.or(existing.get().clone());
-                        if &child_marker != existing.get() {
-                            existing.insert(child_marker);
-                            queue.push(child_edge.target());
-                        }
-                    }
-                    Entry::Vacant(vacant) => {
-                        vacant.insert(child_marker.clone());
-                        queue.push(child_edge.target());
-                    }
-                }
-            }
-        }
-
-        reachability
-    }
-
     fn add_edge(
         petgraph: &mut Graph<ResolutionGraphNode, MarkerTree>,
         inverse: &mut FxHashMap<PackageRef<'_>, NodeIndex>,