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);
 
+        let mut seen = FxHashSet::default();
+        for resolution in resolutions {
             // Add every package to the graph.
-        for (package, versions) in &resolution.nodes {
-            for version in versions {
+            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,11 +107,18 @@ impl ResolutionGraph {
                 )?;
             }
         }
-
+        let mut seen = FxHashSet::default();
+        for resolution in resolutions {
             // Add every edge to the graph.
-        for edge in resolution.edges {
+            for edge in &resolution.edges {
+                if !seen.insert(edge) {
+                    // Insert each node only once.
+                    continue;
+                }
+
                 Self::add_edge(&mut petgraph, &mut inverse, root_index, edge);
             }
+        }
 
         // Extract the `Requires-Python` range, if provided.
         // TODO(charlie): Infer the supported Python range from the `Requires-Python` of the
@@ -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,13 +380,11 @@ 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());
                         }
                     }
-                    }
 
                     // If another fork had the same resolution, merge into that fork instead.
                     if let Some(existing_resolution) = resolutions
@@ -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,
-            edges: dependencies,
+            nodes,
+            edges,
             pins: self.pins,
             markers: self.markers,
         }
     }
 }
 
-/// The resolution from one or more forks 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.
+/// The resolution from a single fork including the virtual packages and the edges between them.
 #[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