Skip to content

Commit

Permalink
Simplify Resolution type (#5301)
Browse files Browse the repository at this point in the history
Looking at how to merge identical forks, i found that the `Resolution`
can be simplified by treating it as a nodes and edges store (plus pins,
they are separate since they are per name-version, not per
(virtual-)package-version). This should also make #5294 more apparent,
which i didn't touch here.

I additionally added some doc comments to the `Resolution` types.
  • Loading branch information
konstin authored Jul 22, 2024
1 parent dc108e3 commit 0a95b7a
Show file tree
Hide file tree
Showing 2 changed files with 90 additions and 102 deletions.
64 changes: 30 additions & 34 deletions crates/uv-resolver/src/resolution/graph.rs
Original file line number Diff line number Diff line change
Expand Up @@ -73,16 +73,16 @@ impl ResolutionGraph {
);

let mut petgraph: Graph<ResolutionGraphNode, Option<MarkerTree>, Directed> =
Graph::with_capacity(resolution.packages.len(), resolution.packages.len());
Graph::with_capacity(resolution.nodes.len(), resolution.nodes.len());
let mut inverse: FxHashMap<NodeKey, NodeIndex<u32>> =
FxHashMap::with_capacity_and_hasher(resolution.packages.len(), FxBuildHasher);
FxHashMap::with_capacity_and_hasher(resolution.nodes.len(), 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.packages {
for (package, versions) in &resolution.nodes {
let ResolutionPackage {
name,
extra,
Expand Down Expand Up @@ -250,39 +250,35 @@ impl ResolutionGraph {
}

// Add every edge to the graph.
for (names, version_set) in resolution.dependencies {
for versions in version_set {
let from_index = names.from.as_ref().map_or(root_index, |from| {
inverse[&(
from,
&versions.from_version,
versions.from_extra.as_ref(),
versions.from_dev.as_ref(),
)]
});
let to_index = inverse[&(
&names.to,
&versions.to_version,
versions.to_extra.as_ref(),
versions.to_dev.as_ref(),
)];

if let Some(edge) = petgraph
.find_edge(from_index, to_index)
.and_then(|edge| petgraph.edge_weight_mut(edge))
{
// 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)) =
(edge.as_mut(), versions.marker)
{
marker.or(version_marker.clone());
} else {
*edge = None;
}
for edge in resolution.edges {
let from_index = edge.from.as_ref().map_or(root_index, |from| {
inverse[&(
from,
&edge.from_version,
edge.from_extra.as_ref(),
edge.from_dev.as_ref(),
)]
});
let to_index = inverse[&(
&edge.to,
&edge.to_version,
edge.to_extra.as_ref(),
edge.to_dev.as_ref(),
)];

if let Some(marker) = petgraph
.find_edge(from_index, to_index)
.and_then(|edge| petgraph.edge_weight_mut(edge))
{
// 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) {
marker.or(version_marker.clone());
} else {
petgraph.update_edge(from_index, to_index, versions.marker.clone());
*marker = None;
}
} else {
petgraph.update_edge(from_index, to_index, edge.marker.clone());
}
}

Expand Down
128 changes: 60 additions & 68 deletions crates/uv-resolver/src/resolver/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -382,7 +382,7 @@ 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.packages {
for (package, versions) 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());
Expand Down Expand Up @@ -597,42 +597,39 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
if !tracing::enabled!(Level::TRACE) {
return;
}
for (names, versions) in &combined.dependencies {
for edge in &combined.edges {
trace!(
"Resolution: {} -> {}",
names
.from
edge.from
.as_ref()
.map(PackageName::as_str)
.unwrap_or("ROOT"),
names.to,
edge.to,
);
for v in versions {
// The unwraps below are OK because `write`ing to
// a String can never fail (except for OOM).
let mut msg = String::new();
write!(msg, "{}", v.from_version).unwrap();
if let Some(ref extra) = v.from_extra {
write!(msg, " (extra: {extra})").unwrap();
}
if let Some(ref dev) = v.from_dev {
write!(msg, " (group: {dev})").unwrap();
}
// The unwraps below are OK because `write`ing to
// a String can never fail (except for OOM).
let mut msg = String::new();
write!(msg, "{}", edge.from_version).unwrap();
if let Some(ref extra) = edge.from_extra {
write!(msg, " (extra: {extra})").unwrap();
}
if let Some(ref dev) = edge.from_dev {
write!(msg, " (group: {dev})").unwrap();
}

write!(msg, " -> ").unwrap();
write!(msg, " -> ").unwrap();

write!(msg, "{}", v.to_version).unwrap();
if let Some(ref extra) = v.to_extra {
write!(msg, " (extra: {extra})").unwrap();
}
if let Some(ref dev) = v.to_dev {
write!(msg, " (group: {dev})").unwrap();
}
if let Some(ref marker) = v.marker {
write!(msg, " ; {marker}").unwrap();
}
trace!("Resolution: {msg}");
write!(msg, "{}", edge.to_version).unwrap();
if let Some(ref extra) = edge.to_extra {
write!(msg, " (extra: {extra})").unwrap();
}
if let Some(ref dev) = edge.to_dev {
write!(msg, " (group: {dev})").unwrap();
}
if let Some(ref marker) = edge.marker {
write!(msg, " ; {marker}").unwrap();
}
trace!("Resolution: {msg}");
}
}

Expand Down Expand Up @@ -2193,10 +2190,7 @@ impl ForkState {

fn into_resolution(self) -> Resolution {
let solution = self.pubgrub.partial_solution.extract_solution();
let mut dependencies: FxHashMap<
ResolutionDependencyNames,
FxHashSet<ResolutionDependencyVersions>,
> = FxHashMap::default();
let mut dependencies: 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 @@ -2244,20 +2238,18 @@ impl ForkState {
if self_name.is_some_and(|self_name| self_name == dependency_name) {
continue;
}
let names = ResolutionDependencyNames {
let edge = ResolutionDependencyEdge {
from: self_name.cloned(),
to: dependency_name.clone(),
};
let versions = ResolutionDependencyVersions {
from_version: self_version.clone(),
from_extra: self_extra.cloned(),
from_dev: self_dev.cloned(),
to: dependency_name.clone(),
to_version: dependency_version.clone(),
to_extra: dependency_extra.clone(),
to_dev: dependency_dev.clone(),
marker: None,
};
dependencies.entry(names).or_default().insert(versions);
dependencies.insert(edge);
}

PubGrubPackageInner::Marker {
Expand All @@ -2268,20 +2260,18 @@ impl ForkState {
if self_name.is_some_and(|self_name| self_name == dependency_name) {
continue;
}
let names = ResolutionDependencyNames {
let edge = ResolutionDependencyEdge {
from: self_name.cloned(),
to: dependency_name.clone(),
};
let versions = ResolutionDependencyVersions {
from_version: self_version.clone(),
from_extra: self_extra.cloned(),
from_dev: self_dev.cloned(),
to: dependency_name.clone(),
to_version: dependency_version.clone(),
to_extra: None,
to_dev: None,
marker: Some(dependency_marker.clone()),
};
dependencies.entry(names).or_default().insert(versions);
dependencies.insert(edge);
}

PubGrubPackageInner::Extra {
Expand All @@ -2293,20 +2283,18 @@ impl ForkState {
if self_name.is_some_and(|self_name| self_name == dependency_name) {
continue;
}
let names = ResolutionDependencyNames {
let edge = ResolutionDependencyEdge {
from: self_name.cloned(),
to: dependency_name.clone(),
};
let versions = ResolutionDependencyVersions {
from_version: self_version.clone(),
from_extra: self_extra.cloned(),
from_dev: self_dev.cloned(),
to: dependency_name.clone(),
to_version: dependency_version.clone(),
to_extra: Some(dependency_extra.clone()),
to_dev: None,
marker: dependency_marker.clone(),
};
dependencies.entry(names).or_default().insert(versions);
dependencies.insert(edge);
}

PubGrubPackageInner::Dev {
Expand All @@ -2318,20 +2306,18 @@ impl ForkState {
if self_name.is_some_and(|self_name| self_name == dependency_name) {
continue;
}
let names = ResolutionDependencyNames {
let edge = ResolutionDependencyEdge {
from: self_name.cloned(),
to: dependency_name.clone(),
};
let versions = ResolutionDependencyVersions {
from_version: self_version.clone(),
from_extra: self_extra.cloned(),
from_dev: self_dev.cloned(),
to: dependency_name.clone(),
to_version: dependency_version.clone(),
to_extra: None,
to_dev: Some(dependency_dev.clone()),
marker: dependency_marker.clone(),
};
dependencies.entry(names).or_default().insert(versions);
dependencies.insert(edge);
}

_ => {}
Expand Down Expand Up @@ -2365,40 +2351,48 @@ impl ForkState {
.collect();

Resolution {
packages,
dependencies,
nodes: packages,
edges: dependencies,
pins: self.pins,
}
}
}

/// 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.
#[derive(Debug, Default)]
pub(crate) struct Resolution {
pub(crate) packages: FxHashMap<ResolutionPackage, FxHashSet<Version>>,
pub(crate) dependencies:
FxHashMap<ResolutionDependencyNames, FxHashSet<ResolutionDependencyVersions>>,
pub(crate) nodes: FxHashMap<ResolutionPackage, FxHashSet<Version>>,
/// If `foo` requires `bar>=3` and `foo` requires `bar <3` in another fork, we'd store it as
/// `(foo, bar) -> {>=3, <3}`.
pub(crate) edges: FxHashSet<ResolutionDependencyEdge>,
/// Map each package name, version tuple from `packages` to a distribution.
pub(crate) pins: FilePins,
}

/// Package representation we used during resolution where each extra and also the dev-dependencies
/// group are their own package.
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub(crate) struct ResolutionPackage {
pub(crate) name: PackageName,
pub(crate) extra: Option<ExtraName>,
pub(crate) dev: Option<GroupName>,
/// For index packages, this is `None`.
pub(crate) url: Option<VerbatimParsedUrl>,
}

/// The `from_` fields and the `to_` fields allow mapping to the originating and target
/// [`ResolutionPackage`] respectively.
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub(crate) struct ResolutionDependencyNames {
pub(crate) struct ResolutionDependencyEdge {
/// This value is `None` if the dependency comes from the root package.
pub(crate) from: Option<PackageName>,
pub(crate) to: PackageName,
}

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub(crate) struct ResolutionDependencyVersions {
pub(crate) from_version: Version,
pub(crate) from_extra: Option<ExtraName>,
pub(crate) from_dev: Option<GroupName>,
pub(crate) to: PackageName,
pub(crate) to_version: Version,
pub(crate) to_extra: Option<ExtraName>,
pub(crate) to_dev: Option<GroupName>,
Expand All @@ -2407,15 +2401,13 @@ pub(crate) struct ResolutionDependencyVersions {

impl Resolution {
fn union(&mut self, other: Resolution) {
for (other_package, other_versions) in other.packages {
self.packages
for (other_package, other_versions) in other.nodes {
self.nodes
.entry(other_package)
.or_default()
.extend(other_versions);
}
for (names, versions) in other.dependencies {
self.dependencies.entry(names).or_default().extend(versions);
}
self.edges.extend(other.edges);
self.pins.union(other.pins);
}
}
Expand Down

0 comments on commit 0a95b7a

Please sign in to comment.