Unify resolutions only during graph building (#5479)

With our previous eager union, we were losing the fork markers. We now
carry this information into the resolution graph construction and, in
the next step, can read the markers there.

Part of
https://github.com/astral-sh/uv/issues/5180#issuecomment-2247696198

crates/uv-resolver/src/pins.rs

@@ -29,15 +29,4 @@ impl FilePins {
     ) -> Option<&ResolvedDist> {
         self.0.get(name)?.get(version)
     }
-
-    /// Add the pins in `other` to `self`.
-    ///
-    /// This assumes that if a version for a particular package exists in
-    /// both `self` and `other`, then they will both correspond to identical
-    /// distributions.
-    pub(crate) fn union(&mut self, other: FilePins) {
-        for (name, versions) in other.0 {
-            self.0.entry(name).or_default().extend(versions);
-        }
-    }
 }

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

@@ -1,17 +1,16 @@
-use indexmap::IndexSet;
-use petgraph::{
-    graph::{Graph, NodeIndex},
-    Directed,
-};
-use rustc_hash::{FxBuildHasher, FxHashMap};
-
 use distribution_types::{
     Dist, DistributionMetadata, Name, ResolutionDiagnostic, ResolvedDist, VersionId,
     VersionOrUrlRef,
 };
+use indexmap::IndexSet;
 use pep440_rs::{Version, VersionSpecifier};
 use pep508_rs::{MarkerEnvironment, MarkerTree};
+use petgraph::{
+    graph::{Graph, NodeIndex},
+    Directed,
+};
 use pypi_types::{HashDigest, ParsedUrlError, Requirement, VerbatimParsedUrl, Yanked};
+use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet};
 use uv_configuration::{Constraints, Overrides};
 use uv_distribution::Metadata;
 use uv_git::GitResolver;
@@ -67,7 +66,7 @@ struct PackageRef<'a> {
 impl ResolutionGraph {
     /// Create a new graph from the resolved PubGrub state.
     pub(crate) fn from_state(
-        resolution: Resolution,
+        resolutions: &[Resolution],
         requirements: &[Requirement],
         constraints: &Constraints,
         overrides: &Overrides,
@@ -77,18 +76,24 @@ impl ResolutionGraph {
         python: &PythonRequirement,
         options: Options,
     ) -> Result<Self, ResolveError> {
+        let size_guess = resolutions[0].nodes.len();
         let mut petgraph: Graph<ResolutionGraphNode, Option<MarkerTree>, Directed> =
-            Graph::with_capacity(resolution.nodes.len(), resolution.nodes.len());
+            Graph::with_capacity(size_guess, size_guess);
         let mut inverse: FxHashMap<PackageRef, NodeIndex<u32>> =
-            FxHashMap::with_capacity_and_hasher(resolution.nodes.len(), FxBuildHasher);
+            FxHashMap::with_capacity_and_hasher(size_guess, FxBuildHasher);
         let mut diagnostics = Vec::new();
 
         // Add the root node.
         let root_index = petgraph.add_node(ResolutionGraphNode::Root);
 
-        // Add every package to the graph.
+        let mut seen = FxHashSet::default();
-        for (package, versions) in &resolution.nodes {
+        for resolution in resolutions {
-            for version in versions {
+            // Add every package to the graph.
+            for (package, version) in &resolution.nodes {
+                if !seen.insert((package, version)) {
+                    // Insert each node only once.
+                    continue;
+                }
                 Self::add_version(
                     &mut petgraph,
                     &mut inverse,
@@ -102,10 +107,17 @@ impl ResolutionGraph {
                 )?;
             }
         }
+        let mut seen = FxHashSet::default();
+        for resolution in resolutions {
+            // Add every edge to the graph.
+            for edge in &resolution.edges {
+                if !seen.insert(edge) {
+                    // Insert each node only once.
+                    continue;
+                }
 
-        // Add every edge to the graph.
+                Self::add_edge(&mut petgraph, &mut inverse, root_index, edge);
-        for edge in resolution.edges {
+            }
-            Self::add_edge(&mut petgraph, &mut inverse, root_index, edge);
         }
 
         // Extract the `Requires-Python` range, if provided.
@@ -141,7 +153,7 @@ impl ResolutionGraph {
         petgraph: &mut Graph<ResolutionGraphNode, Option<MarkerTree>>,
         inverse: &mut FxHashMap<PackageRef<'_>, NodeIndex>,
         root_index: NodeIndex,
-        edge: ResolutionDependencyEdge,
+        edge: &ResolutionDependencyEdge,
     ) {
         let from_index = edge.from.as_ref().map_or(root_index, |from| {
             inverse[&PackageRef {
@@ -166,7 +178,8 @@ impl ResolutionGraph {
         {
             // If either the existing marker or new marker is `None`, then the dependency is
             // included unconditionally, and so the combined marker should be `None`.
-            if let (Some(marker), Some(ref version_marker)) = (marker.as_mut(), edge.marker) {
+            if let (Some(marker), Some(ref version_marker)) = (marker.as_mut(), edge.marker.clone())
+            {
                 marker.or(version_marker.clone());
             } else {
                 *marker = None;

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

@@ -380,11 +380,9 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
                     let resolution = state.into_resolution();
 
                     // Walk over the selected versions, and mark them as preferences.
-                    for (package, versions) in &resolution.nodes {
+                    for (package, version) in &resolution.nodes {
                         if let Entry::Vacant(entry) = preferences.entry(package.name.clone()) {
-                            if let Some(version) = versions.iter().next() {
+                            entry.insert(version.clone().into());
-                                entry.insert(version.clone().into());
-                            }
                         }
                     }
 
@@ -409,6 +407,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
                         continue 'FORK;
                     }
 
+                    Self::trace_resolution(&resolution);
                     resolutions.push(resolution);
                     continue 'FORK;
                 };
@@ -591,25 +590,22 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
                 }
             }
         }
-        let mut combined = Resolution::universal();
         if resolutions.len() > 1 {
             info!(
                 "Solved your requirements for {} environments",
                 resolutions.len()
             );
         }
-        for resolution in resolutions {
+        for resolution in &resolutions {
             if let Some(markers) = resolution.markers.fork_markers() {
                 debug!(
                     "Distinct solution for ({markers}) with {} packages",
                     resolution.nodes.len()
                 );
             }
-            combined.union(resolution);
         }
-        Self::trace_resolution(&combined);
         ResolutionGraph::from_state(
-            combined,
+            &resolutions,
             &self.requirements,
             &self.constraints,
             &self.overrides,
@@ -2244,7 +2240,7 @@ impl ForkState {
 
     fn into_resolution(self) -> Resolution {
         let solution = self.pubgrub.partial_solution.extract_solution();
-        let mut dependencies: FxHashSet<ResolutionDependencyEdge> = FxHashSet::default();
+        let mut edges: FxHashSet<ResolutionDependencyEdge> = FxHashSet::default();
         for (package, self_version) in &solution {
             for id in &self.pubgrub.incompatibilities[package] {
                 let pubgrub::solver::Kind::FromDependencyOf(
@@ -2307,7 +2303,7 @@ impl ForkState {
                             to_dev: dependency_dev.clone(),
                             marker: None,
                         };
-                        dependencies.insert(edge);
+                        edges.insert(edge);
                     }
 
                     PubGrubPackageInner::Marker {
@@ -2332,7 +2328,7 @@ impl ForkState {
                             to_dev: None,
                             marker: Some(dependency_marker.clone()),
                         };
-                        dependencies.insert(edge);
+                        edges.insert(edge);
                     }
 
                     PubGrubPackageInner::Extra {
@@ -2358,7 +2354,7 @@ impl ForkState {
                             to_dev: None,
                             marker: dependency_marker.clone(),
                         };
-                        dependencies.insert(edge);
+                        edges.insert(edge);
                     }
 
                     PubGrubPackageInner::Dev {
@@ -2384,7 +2380,7 @@ impl ForkState {
                             to_dev: Some(dependency_dev.clone()),
                             marker: dependency_marker.clone(),
                         };
-                        dependencies.insert(edge);
+                        edges.insert(edge);
                     }
 
                     _ => {}
@@ -2392,7 +2388,7 @@ impl ForkState {
             }
         }
 
-        let packages = solution
+        let nodes = solution
             .into_iter()
             .filter_map(|(package, version)| {
                 if let PubGrubPackageInner::Package {
@@ -2409,7 +2405,7 @@ impl ForkState {
                             dev: dev.clone(),
                             url: self.fork_urls.get(name).cloned(),
                         },
-                        FxHashSet::from_iter([version]),
+                        version,
                     ))
                 } else {
                     None
@@ -2418,21 +2414,18 @@ impl ForkState {
             .collect();
 
         Resolution {
-            nodes: packages,
+            nodes,
-            edges: dependencies,
+            edges,
             pins: self.pins,
             markers: self.markers,
         }
     }
 }
 
-/// The resolution from one or more forks including the virtual packages and the edges between them.
+/// The resolution from a single fork including the virtual packages and the edges between them.
-///
-/// Each package can have multiple versions and each edge between two packages can have multiple
-/// version specifiers to support diverging versions and requirements in different forks.
 #[derive(Debug)]
 pub(crate) struct Resolution {
-    pub(crate) nodes: FxHashMap<ResolutionPackage, FxHashSet<Version>>,
+    pub(crate) nodes: FxHashMap<ResolutionPackage, Version>,
     /// The directed connections between the nodes, where the marker is the node weight. We don't
     /// store the requirement itself, but it can be retrieved from the package metadata.
     pub(crate) edges: FxHashSet<ResolutionDependencyEdge>,
@@ -2471,17 +2464,6 @@ pub(crate) struct ResolutionDependencyEdge {
     pub(crate) marker: Option<MarkerTree>,
 }
 
-impl Resolution {
-    fn universal() -> Self {
-        Self {
-            nodes: FxHashMap::default(),
-            edges: FxHashSet::default(),
-            pins: FilePins::default(),
-            markers: ResolverMarkers::Universal,
-        }
-    }
-}
-
 impl Resolution {
     /// Whether we got two identical resolutions in two separate forks.
     ///
@@ -2495,17 +2477,6 @@ impl Resolution {
         // change which packages and versions are installed.
         self.nodes == other.nodes && self.edges == other.edges
     }
-
-    fn union(&mut self, other: Resolution) {
-        for (other_package, other_versions) in other.nodes {
-            self.nodes
-                .entry(other_package)
-                .or_default()
-                .extend(other_versions);
-        }
-        self.edges.extend(other.edges);
-        self.pins.union(other.pins);
-    }
 }
 
 /// Fetch the metadata for an item