Skip to content

Commit

Permalink
Unify resolutions only during graph building (#5479)
Browse files Browse the repository at this point in the history 
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
#5180 (comment)
  • Loading branch information
@konstin
konstin authored Jul 26, 2024
1 parent 77b0052 commit cb505d2
Show file tree
Hide file tree
Showing 3 changed files with 47 additions and 74 deletions.
11 changes: 0 additions & 11 deletions crates/uv-resolver/src/pins.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}
}
49 changes: 31 additions & 18 deletions crates/uv-resolver/src/resolution/graph.rs
View file
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -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,
Expand All @@ -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.
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,
Expand All @@ -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.
Expand Down Expand Up @@ -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 {
Expand All @@ -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;
Expand Down
61 changes: 16 additions & 45 deletions crates/uv-resolver/src/resolver/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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());
}
}

Expand All @@ -409,6 +407,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
continue 'FORK;
}

Self::trace_resolution(&resolution);
resolutions.push(resolution);
continue 'FORK;
};
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -2307,7 +2303,7 @@ impl ForkState {
to_dev: dependency_dev.clone(),
marker: None,
};
dependencies.insert(edge);
edges.insert(edge);
}

PubGrubPackageInner::Marker {
Expand All @@ -2332,7 +2328,7 @@ impl ForkState {
to_dev: None,
marker: Some(dependency_marker.clone()),
};
dependencies.insert(edge);
edges.insert(edge);
}

PubGrubPackageInner::Extra {
Expand All @@ -2358,7 +2354,7 @@ impl ForkState {
to_dev: None,
marker: dependency_marker.clone(),
};
dependencies.insert(edge);
edges.insert(edge);
}

PubGrubPackageInner::Dev {
Expand All @@ -2384,15 +2380,15 @@ impl ForkState {
to_dev: Some(dependency_dev.clone()),
marker: dependency_marker.clone(),
};
dependencies.insert(edge);
edges.insert(edge);
}

_ => {}
}
}
}

let packages = solution
let nodes = solution
.into_iter()
.filter_map(|(package, version)| {
if let PubGrubPackageInner::Package {
Expand All @@ -2409,7 +2405,7 @@ impl ForkState {
dev: dev.clone(),
url: self.fork_urls.get(name).cloned(),
},
FxHashSet::from_iter([version]),
version,
))
} else {
None
Expand All @@ -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>,
Expand Down Expand Up @@ -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.
///
Expand All @@ -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
Expand Down

0 comments on commit cb505d2

Please sign in to comment.