Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Unify resolutions only during graph building #5479

Merged
merged 1 commit into from
Jul 26, 2024
Merged

Unify resolutions only during graph building #5479

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 0 additions & 11 deletions crates/uv-resolver/src/pins.rs
View file
Open in desktop
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
Open in desktop
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
Open in desktop
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>,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah nice, because each resolution is only allowed one version of a package, and Resolution no longer needs to represent the union of all resolutions since we merge them right into the graph?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yep

/// 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
Loading