From cb505d24f80725ccdfb8625acc0ce418d99bf63c Mon Sep 17 00:00:00 2001 From: konsti Date: Fri, 26 Jul 2024 16:29:48 +0200 Subject: [PATCH] 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 | 11 ---- crates/uv-resolver/src/resolution/graph.rs | 49 ++++++++++------- crates/uv-resolver/src/resolver/mod.rs | 61 ++++++---------------- 3 files changed, 47 insertions(+), 74 deletions(-) diff --git a/crates/uv-resolver/src/pins.rs b/crates/uv-resolver/src/pins.rs index 827a8e6f4da1..d471f745a38d 100644 --- a/crates/uv-resolver/src/pins.rs +++ b/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); - } - } } diff --git a/crates/uv-resolver/src/resolution/graph.rs b/crates/uv-resolver/src/resolution/graph.rs index a91d31ed4b49..a2d93e3bcc78 100644 --- a/crates/uv-resolver/src/resolution/graph.rs +++ b/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 { + let size_guess = resolutions[0].nodes.len(); let mut petgraph: Graph, Directed> = - Graph::with_capacity(resolution.nodes.len(), resolution.nodes.len()); + Graph::with_capacity(size_guess, size_guess); let mut inverse: FxHashMap> = - 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. - for (package, versions) in &resolution.nodes { - for version in versions { + let mut seen = FxHashSet::default(); + for resolution in resolutions { + // 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. - for edge in resolution.edges { - Self::add_edge(&mut petgraph, &mut inverse, root_index, edge); + 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>, inverse: &mut FxHashMap, 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; diff --git a/crates/uv-resolver/src/resolver/mod.rs b/crates/uv-resolver/src/resolver/mod.rs index 8d94fcef7f42..961147d3a3c3 100644 --- a/crates/uv-resolver/src/resolver/mod.rs +++ b/crates/uv-resolver/src/resolver/mod.rs @@ -380,11 +380,9 @@ impl ResolverState ResolverState ResolverState 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 = FxHashSet::default(); + let mut edges: FxHashSet = 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>, + pub(crate) nodes: FxHashMap, /// 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, @@ -2471,17 +2464,6 @@ pub(crate) struct ResolutionDependencyEdge { pub(crate) marker: Option, } -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