Prune unreachable packages from --universal output (#7209)

## Summary

Closes https://github.com/astral-sh/uv/issues/7196.

crates/uv-resolver/src/graph_ops.rs

@@ -2,7 +2,7 @@ use pep508_rs::MarkerTree;
 use petgraph::graph::NodeIndex;
 use petgraph::visit::EdgeRef;
 use petgraph::{Direction, Graph};
-use rustc_hash::FxHashMap;
+use rustc_hash::{FxBuildHasher, FxHashMap};
 use std::collections::hash_map::Entry;
 
 /// Determine the markers under which a package is reachable in the dependency tree.
@@ -18,7 +18,7 @@ pub(crate) fn marker_reachability<T>(
 ) -> 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();
+    let mut reachability = FxHashMap::with_capacity_and_hasher(graph.node_count(), FxBuildHasher);
 
     // Collect the root nodes.
     //

crates/uv-resolver/src/lock/mod.rs

@@ -12,7 +12,6 @@ use std::path::{Path, PathBuf};
 use std::str::FromStr;
 use std::sync::LazyLock;
 use toml_edit::{value, Array, ArrayOfTables, InlineTable, Item, Table, Value};
-use tracing::debug;
 use url::Url;
 
 use cache_key::RepositoryUrl;
@@ -115,10 +114,6 @@ impl Lock {
             if !dist.is_base() {
                 continue;
             }
-            if graph.reachability[&node_index].is_false() {
-                debug!("Removing unreachable package: `{}`", dist.package_id());
-                continue;
-            }
             let fork_markers = graph
                 .fork_markers(dist.name(), &dist.version, dist.dist.version_or_url().url())
                 .cloned()
@@ -132,10 +127,6 @@ impl Lock {
                 else {
                     continue;
                 };
-                // Prune edges leading to unreachable nodes.
-                if graph.reachability[&edge.target()].is_false() {
-                    continue;
-                }
                 let marker = edge.weight().clone();
                 package.add_dependency(&requires_python, dependency_dist, marker, root)?;
             }
@@ -154,9 +145,6 @@ impl Lock {
             let ResolutionGraphNode::Dist(dist) = &graph.petgraph[node_index] else {
                 continue;
             };
-            if graph.reachability[&node_index].is_false() {
-                continue;
-            }
             if let Some(extra) = dist.extra.as_ref() {
                 let id = PackageId::from_annotated_dist(dist, root)?;
                 let Some(package) = packages.get_mut(&id) else {
@@ -171,10 +159,6 @@ impl Lock {
                     else {
                         continue;
                     };
-                    // Prune edges leading to unreachable nodes.
-                    if graph.reachability[&edge.target()].is_false() {
-                        continue;
-                    }
                     let marker = edge.weight().clone();
                     package.add_optional_dependency(
                         &requires_python,
@@ -199,10 +183,6 @@ impl Lock {
                     else {
                         continue;
                     };
-                    // Prune edges leading to unreachable nodes.
-                    if graph.reachability[&edge.target()].is_false() {
-                        continue;
-                    }
                     let marker = edge.weight().clone();
                     package.add_dev_dependency(
                         &requires_python,
@@ -268,7 +248,6 @@ impl Lock {
             // `(A ∩ (B ∩ C) = ∅) => ((A ∩ B = ∅) or (A ∩ C = ∅))`
             // a single disjointness check with the intersection is sufficient, so we have one
             // constant per platform.
-            let reachability_markers = &graph.reachability[&node_index];
             let platform_tags = &wheel.filename.platform_tag;
             if platform_tags.iter().all(|tag| {
                 linux_tags.into_iter().any(|linux_tag| {
@@ -276,14 +255,20 @@ impl Lock {
                     tag.starts_with(linux_tag) || tag == "linux_armv6l" || tag == "linux_armv7l"
                 })
             }) {
-                !reachability_markers.is_disjoint(&LINUX_MARKERS)
+                !graph.petgraph[node_index]
+                    .marker()
+                    .is_disjoint(&LINUX_MARKERS)
             } else if platform_tags
                 .iter()
                 .all(|tag| windows_tags.contains(&&**tag))
             {
-                !graph.reachability[&node_index].is_disjoint(&WINDOWS_MARKERS)
+                !graph.petgraph[node_index]
+                    .marker()
+                    .is_disjoint(&WINDOWS_MARKERS)
             } else if platform_tags.iter().all(|tag| tag.starts_with("macosx_")) {
-                !graph.reachability[&node_index].is_disjoint(&MAC_MARKERS)
+                !graph.petgraph[node_index]
+                    .marker()
+                    .is_disjoint(&MAC_MARKERS)
             } else {
                 true
             }

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

@@ -1,7 +1,6 @@
 use std::collections::BTreeSet;
 
 use owo_colors::OwoColorize;
-use petgraph::graph::NodeIndex;
 use petgraph::visit::EdgeRef;
 use petgraph::Direction;
 use rustc_hash::{FxBuildHasher, FxHashMap};
@@ -45,6 +44,15 @@ enum DisplayResolutionGraphNode {
     Dist(RequirementsTxtDist),
 }
 
+impl DisplayResolutionGraphNode {
+    fn markers(&self) -> &MarkerTree {
+        match self {
+            DisplayResolutionGraphNode::Root => &MarkerTree::TRUE,
+            DisplayResolutionGraphNode::Dist(dist) => &dist.markers,
+        }
+    }
+}
+
 impl<'a> DisplayResolutionGraph<'a> {
     /// Create a new [`DisplayResolutionGraph`] for the given graph.
     #[allow(clippy::fn_params_excessive_bools)]
@@ -136,11 +144,10 @@ impl std::fmt::Display for DisplayResolutionGraph<'_> {
         // that for each package tells us if it should be installed on the current platform, without
         // looking at which packages depend on it.
         let petgraph = self.resolution.petgraph.map(
-            |index, node| match node {
+            |_index, node| match node {
                 ResolutionGraphNode::Root => DisplayResolutionGraphNode::Root,
                 ResolutionGraphNode::Dist(dist) => {
-                    let reachability = self.resolution.reachability[&index].clone();
+                    let dist = RequirementsTxtDist::from_annotated_dist(dist);
-                    let dist = RequirementsTxtDist::from_annotated_dist(dist, reachability);
                     DisplayResolutionGraphNode::Dist(dist)
                 }
             },
@@ -151,7 +158,7 @@ impl std::fmt::Display for DisplayResolutionGraph<'_> {
 
         // Reduce the graph, such that all nodes for a single package are combined, regardless of
         // the extras.
-        let petgraph = combine_extras(&petgraph, &self.resolution.reachability);
+        let petgraph = combine_extras(&petgraph);
 
         // Collect all packages.
         let mut nodes = petgraph
@@ -318,10 +325,7 @@ type RequirementsTxtGraph = petgraph::graph::Graph<RequirementsTxtDist, (), petg
 /// node.
 ///
 /// We also remove the root node, to simplify the graph structure.
-fn combine_extras(
+fn combine_extras(graph: &IntermediatePetGraph) -> RequirementsTxtGraph {
-    graph: &IntermediatePetGraph,
-    reachability: &FxHashMap<NodeIndex, MarkerTree>,
-) -> RequirementsTxtGraph {
     let mut next = RequirementsTxtGraph::with_capacity(graph.node_count(), graph.edge_count());
     let mut inverse = FxHashMap::with_capacity_and_hasher(graph.node_count(), FxBuildHasher);
 
@@ -361,7 +365,7 @@ fn combine_extras(
             };
             let combined_markers = next[inverse[&dist.version_id()]].markers.clone();
             let mut package_markers = combined_markers.clone();
-            package_markers.or(reachability[&index].clone());
+            package_markers.or(graph[index].markers().clone());
             assert_eq!(
                 package_markers,
                 combined_markers,

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

@@ -55,16 +55,23 @@ pub struct ResolutionGraph {
     /// If there are multiple options for a package, track which fork they belong to so we
     /// can write that to the lockfile and later get the correct preference per fork back.
     pub(crate) package_markers: FxHashMap<PackageName, MarkersForDistribution>,
-    /// The markers under which a package is reachable in the dependency tree.
-    pub(crate) reachability: FxHashMap<NodeIndex, MarkerTree>,
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub(crate) enum ResolutionGraphNode {
     Root,
     Dist(AnnotatedDist),
 }
 
+impl ResolutionGraphNode {
+    pub(crate) fn marker(&self) -> &MarkerTree {
+        match self {
+            ResolutionGraphNode::Root => &MarkerTree::TRUE,
+            ResolutionGraphNode::Dist(dist) => &dist.marker,
+        }
+    }
+}
+
 impl Display for ResolutionGraphNode {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         match self {
@@ -209,7 +216,18 @@ impl ResolutionGraph {
                 .collect()
         };
 
-        let reachability = marker_reachability(&petgraph, &fork_markers);
+        // Compute and apply the marker reachability.
+        let mut reachability = marker_reachability(&petgraph, &fork_markers);
+
+        // Apply the reachability to the graph.
+        for index in petgraph.node_indices() {
+            if let ResolutionGraphNode::Dist(dist) = &mut petgraph[index] {
+                dist.marker = reachability.remove(&index).unwrap_or_default();
+            }
+        }
+
+        // Discard any unreachable nodes.
+        petgraph.retain_nodes(|graph, node| !graph[node].marker().is_false());
 
         if matches!(resolution_strategy, ResolutionStrategy::Lowest) {
             report_missing_lower_bounds(&petgraph, &mut diagnostics);
@@ -225,7 +243,6 @@ impl ResolutionGraph {
             overrides: overrides.clone(),
             options,
             fork_markers,
-            reachability,
         })
     }
 
@@ -331,6 +348,7 @@ impl ResolutionGraph {
             dev: dev.clone(),
             hashes,
             metadata,
+            marker: MarkerTree::TRUE,
         }));
         inverse.insert(
             PackageRef {

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

@@ -5,6 +5,7 @@ use distribution_types::{
     VersionOrUrlRef,
 };
 use pep440_rs::Version;
+use pep508_rs::MarkerTree;
 use pypi_types::HashDigest;
 use uv_distribution::Metadata;
 use uv_normalize::{ExtraName, GroupName, PackageName};
@@ -30,6 +31,7 @@ pub(crate) struct AnnotatedDist {
     pub(crate) dev: Option<GroupName>,
     pub(crate) hashes: Vec<HashDigest>,
     pub(crate) metadata: Option<Metadata>,
+    pub(crate) marker: MarkerTree,
 }
 
 impl AnnotatedDist {

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

@@ -165,7 +165,7 @@ impl RequirementsTxtDist {
         }
     }
 
-    pub(crate) fn from_annotated_dist(annotated: &AnnotatedDist, markers: MarkerTree) -> Self {
+    pub(crate) fn from_annotated_dist(annotated: &AnnotatedDist) -> Self {
         Self {
             dist: annotated.dist.clone(),
             version: annotated.version.clone(),
@@ -175,7 +175,7 @@ impl RequirementsTxtDist {
                 vec![]
             },
             hashes: annotated.hashes.clone(),
-            markers,
+            markers: annotated.marker.clone(),
         }
     }
 }

crates/uv/tests/lock_scenarios.rs

@@ -4444,7 +4444,7 @@ fn unreachable_package() -> Result<()> {
     ----- stdout -----
 
     ----- stderr -----
-    Resolved 3 packages in [TIME]
+    Resolved 2 packages in [TIME]
     "###
     );
 

crates/uv/tests/pip_compile.rs

@@ -12339,3 +12339,46 @@ fn importlib_metadata_not_repeated() -> Result<()> {
 
     Ok(())
 }
+
+/// Regression test for: <https://github.com/astral-sh/uv/issues/6836>
+#[test]
+fn prune_unreachable() -> Result<()> {
+    let context = TestContext::new("3.12");
+
+    let requirements_in = context.temp_dir.child("requirements.in");
+    requirements_in.write_str("argcomplete ; python_version >= '3.8'")?;
+
+    let filters: Vec<_> = [
+        // 3.7 may not be installed
+        (
+            "warning: The requested Python version 3.7 is not available; .* will be used to build dependencies instead.\n",
+            "",
+        ),
+    ]
+        .into_iter()
+        .chain(context.filters())
+        .collect();
+
+    uv_snapshot!(
+        filters,
+        context
+            .pip_compile()
+            .arg("requirements.in")
+            .arg("--universal")
+            .arg("-p")
+            .arg("3.7"),
+        @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+    # This file was autogenerated by uv via the following command:
+    #    uv pip compile --cache-dir [CACHE_DIR] requirements.in --universal -p 3.7
+    argcomplete==3.1.2 ; python_full_version >= '3.8'
+        # via -r requirements.in
+
+    ----- stderr -----
+    Resolved 1 package in [TIME]
+    "###);
+
+    Ok(())
+}

crates/uv/tests/snapshots/ecosystem__transformers-uv-lock-output.snap

@@ -7,4 +7,4 @@ exit_code: 0
 ----- stdout -----
 
 ----- stderr -----
-Resolved 298 packages in [TIME]
+Resolved 296 packages in [TIME]