Skip to content

Commit

Permalink
Add root
Browse files Browse the repository at this point in the history
  • Loading branch information
@charliermarsh
charliermarsh committed Jun 25, 2024
1 parent 1ee63c0 commit 82c6ea1
Show file tree
Hide file tree
Showing 12 changed files with 302 additions and 153 deletions.
25 changes: 19 additions & 6 deletions crates/uv-resolver/src/lock.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ use uv_configuration::ExtrasSpecification;
use uv_git::{GitReference, GitSha, RepositoryReference, ResolvedRepositoryReference};
use uv_normalize::{ExtraName, GroupName, PackageName};

use crate::resolution::AnnotatedDist;
use crate::resolution::{AnnotatedDist, ResolutionGraphNode};
use crate::{RequiresPython, ResolutionGraph};

#[derive(Clone, Debug, serde::Deserialize)]
Expand Down Expand Up @@ -62,11 +62,16 @@ impl Lock {

// Lock all base packages.
for node_index in graph.petgraph.node_indices() {
let dist = &graph.petgraph[node_index];
let ResolutionGraphNode::Dist(dist) = &graph.petgraph[node_index] else {
continue;
};
if dist.is_base() {
let mut locked_dist = Distribution::from_annotated_dist(dist)?;
for edge in graph.petgraph.edges(node_index) {
let dependency_dist = &graph.petgraph[edge.target()];
let ResolutionGraphNode::Dist(dependency_dist) = &graph.petgraph[edge.target()]
else {
continue;
};
let marker = edge.weight().as_ref();
locked_dist.add_dependency(dependency_dist, marker);
}
Expand All @@ -82,7 +87,9 @@ impl Lock {

// Lock all extras and development dependencies.
for node_index in graph.petgraph.node_indices() {
let dist = &graph.petgraph[node_index];
let ResolutionGraphNode::Dist(dist) = &graph.petgraph[node_index] else {
continue;
};
if let Some(extra) = dist.extra.as_ref() {
let id = DistributionId::from_annotated_dist(dist);
let Some(locked_dist) = locked_dists.get_mut(&id) else {
Expand All @@ -93,7 +100,10 @@ impl Lock {
.into());
};
for edge in graph.petgraph.edges(node_index) {
let dependency_dist = &graph.petgraph[edge.target()];
let ResolutionGraphNode::Dist(dependency_dist) = &graph.petgraph[edge.target()]
else {
continue;
};
let marker = edge.weight().as_ref();
locked_dist.add_optional_dependency(extra.clone(), dependency_dist, marker);
}
Expand All @@ -108,7 +118,10 @@ impl Lock {
.into());
};
for edge in graph.petgraph.edges(node_index) {
let dependency_dist = &graph.petgraph[edge.target()];
let ResolutionGraphNode::Dist(dependency_dist) = &graph.petgraph[edge.target()]
else {
continue;
};
let marker = edge.weight().as_ref();
locked_dist.add_dev_dependency(group.clone(), dependency_dist, marker);
}
Expand Down
62 changes: 46 additions & 16 deletions crates/uv-resolver/src/resolution/display.rs
View file
Original file line number Diff line number Diff line change
@@ -1,17 +1,17 @@
use std::collections::BTreeSet;

use owo_colors::OwoColorize;
use petgraph::Direction;
use petgraph::visit::{EdgeRef, Topo};
use petgraph::Direction;
use rustc_hash::{FxBuildHasher, FxHashMap};

use distribution_types::{Name, SourceAnnotation, SourceAnnotations};
use distribution_types::{DistributionMetadata, Name, SourceAnnotation, SourceAnnotations};
use pep508_rs::MarkerEnvironment;
use pep508_rs::MarkerTree;
use uv_normalize::PackageName;

use crate::resolution::{RequirementsTxtDist, ResolutionGraphNode};
use crate::{marker, ResolutionGraph};
use crate::resolution::RequirementsTxtDist;

/// A [`std::fmt::Display`] implementation for the resolution graph.
#[derive(Debug)]
Expand Down Expand Up @@ -39,6 +39,12 @@ pub struct DisplayResolutionGraph<'a> {
annotation_style: AnnotationStyle,
}

#[derive(Debug)]
enum DisplayResolutionGraphNode {
Root,
Dist(RequirementsTxtDist),
}

impl<'a> From<&'a ResolutionGraph> for DisplayResolutionGraph<'a> {
fn from(resolution: &'a ResolutionGraph) -> Self {
Self::new(
Expand Down Expand Up @@ -139,8 +145,8 @@ impl std::fmt::Display for DisplayResolutionGraph<'_> {

// Convert from `AnnotatedDist` to `RequirementsTxtDist`.
let mut petgraph = {
let mut petgraph = petgraph::graph::Graph::<
RequirementsTxtDist,
let mut nextgraph = petgraph::graph::Graph::<
DisplayResolutionGraphNode,
Option<MarkerTree>,
petgraph::Directed,
>::with_capacity(
Expand All @@ -154,8 +160,18 @@ impl std::fmt::Display for DisplayResolutionGraph<'_> {

// Re-add the nodes to the reduced graph.
for index in self.resolution.petgraph.node_indices() {
let dist = &self.resolution.petgraph[index];
inverse.insert(index, petgraph.add_node(RequirementsTxtDist::from(dist)));
match &self.resolution.petgraph[index] {
ResolutionGraphNode::Root => {
inverse.insert(index, nextgraph.add_node(DisplayResolutionGraphNode::Root));
}
ResolutionGraphNode::Dist(dist) => {
let dist = RequirementsTxtDist::from(dist);
inverse.insert(
index,
nextgraph.add_node(DisplayResolutionGraphNode::Dist(dist)),
);
}
}
}

// Re-add the edges to the reduced graph.
Expand All @@ -164,10 +180,10 @@ impl std::fmt::Display for DisplayResolutionGraph<'_> {
let weight = self.resolution.petgraph[edge].clone();
let source = inverse[&source];
let target = inverse[&target];
petgraph.update_edge(source, target, weight);
nextgraph.update_edge(source, target, weight);
}

petgraph
nextgraph
};

// Propagate markers across the graph: the graph is directed, so if any edge contains a
Expand Down Expand Up @@ -203,22 +219,26 @@ impl std::fmt::Display for DisplayResolutionGraph<'_> {
while let Some((outgoing, _)) = walker.next(&petgraph) {
if let Some(weight) = petgraph.edge_weight_mut(outgoing) {
if let Some(weight) = weight {
weight.or(marker_tree.clone());
weight.and(marker_tree.clone());
} else {
*weight = Some(marker_tree.clone());
}
}
}
}

petgraph[index].markers = marker_tree.and_then(marker::normalize);
if let DisplayResolutionGraphNode::Dist(node) = &mut petgraph[index] {
node.markers = marker_tree.and_then(marker::normalize);
};
}

// Reduce the graph, such that all nodes for a single package are combined, regardless of
// the extras.
//
// For example, `flask` and `flask[dotenv]` should be reduced into a single `flask[dotenv]`
// node.
//
// We also remove the root node, to simplify the graph structure.
let petgraph = {
let mut nextgraph = petgraph::graph::Graph::<
RequirementsTxtDist,
Expand All @@ -232,7 +252,9 @@ impl std::fmt::Display for DisplayResolutionGraph<'_> {

// Re-add the nodes to the reduced graph.
for index in petgraph.node_indices() {
let dist = &petgraph[index];
let DisplayResolutionGraphNode::Dist(dist) = &petgraph[index] else {
continue;
};

if let Some(index) = inverse.get(&dist.version_id()) {
let node: &mut RequirementsTxtDist = &mut nextgraph[*index];
Expand All @@ -248,9 +270,15 @@ impl std::fmt::Display for DisplayResolutionGraph<'_> {
// Re-add the edges to the reduced graph.
for edge in petgraph.edge_indices() {
let (source, target) = petgraph.edge_endpoints(edge).unwrap();
let DisplayResolutionGraphNode::Dist(source_node) = &petgraph[source] else {
continue;
};
let DisplayResolutionGraphNode::Dist(target_node) = &petgraph[target] else {
continue;
};
let weight = petgraph[edge].clone();
let source = inverse[&petgraph[source].version_id()];
let target = inverse[&petgraph[target].version_id()];
let source = inverse[&source_node.version_id()];
let target = inverse[&target_node.version_id()];

// If either the existing marker or new marker is `None`, then the dependency is
// included unconditionally, and so the combined marker should be `None`.
Expand All @@ -259,7 +287,7 @@ impl std::fmt::Display for DisplayResolutionGraph<'_> {
.and_then(|edge| nextgraph.edge_weight_mut(edge))
{
if let (Some(marker), Some(ref version_marker)) = (edge.as_mut(), weight) {
marker.or(version_marker.clone());
marker.and(version_marker.clone());
} else {
*edge = None;
}
Expand Down Expand Up @@ -291,7 +319,9 @@ impl std::fmt::Display for DisplayResolutionGraph<'_> {
// Print out the dependency graph.
for (index, node) in nodes {
// Display the node itself.
let mut line = node.to_requirements_txt(self.include_extras).to_string();
let mut line = node
.to_requirements_txt(self.include_extras, self.include_markers)
.to_string();

// Display the distribution hashes, if any.
let mut has_hashes = false;
Expand Down
83 changes: 45 additions & 38 deletions crates/uv-resolver/src/resolution/graph.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ use crate::{
#[derive(Debug)]
pub struct ResolutionGraph {
/// The underlying graph.
pub(crate) petgraph: Graph<AnnotatedDist, Option<MarkerTree>, Directed>,
pub(crate) petgraph: Graph<ResolutionGraphNode, Option<MarkerTree>, Directed>,
/// The range of supported Python versions.
pub(crate) requires_python: Option<RequiresPython>,
/// Any diagnostics that were encountered while building the graph.
Expand All @@ -44,12 +44,11 @@ pub struct ResolutionGraph {
pub(crate) overrides: Overrides,
}

type NodeKey<'a> = (
&'a PackageName,
&'a Version,
Option<&'a ExtraName>,
Option<&'a GroupName>,
);
#[derive(Debug)]
pub(crate) enum ResolutionGraphNode {
Root,
Dist(AnnotatedDist),
}

impl ResolutionGraph {
/// Create a new graph from the resolved PubGrub state.
Expand All @@ -64,12 +63,22 @@ impl ResolutionGraph {
python: &PythonRequirement,
resolution: Resolution,
) -> Result<Self, ResolveError> {
let mut petgraph: Graph<AnnotatedDist, Option<MarkerTree>, Directed> =
type NodeKey<'a> = (
&'a PackageName,
&'a Version,
Option<&'a ExtraName>,
Option<&'a GroupName>,
);

let mut petgraph: Graph<ResolutionGraphNode, Option<MarkerTree>, Directed> =
Graph::with_capacity(resolution.packages.len(), resolution.packages.len());
let mut inverse: FxHashMap<NodeKey, NodeIndex<u32>> =
FxHashMap::with_capacity_and_hasher(resolution.packages.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 {
let ResolutionPackage {
Expand Down Expand Up @@ -227,37 +236,35 @@ impl ResolutionGraph {
}

// Add the distribution to the graph.
let index = petgraph.add_node(AnnotatedDist {
let index = petgraph.add_node(ResolutionGraphNode::Dist(AnnotatedDist {
dist,
extra: extra.clone(),
dev: dev.clone(),
hashes,
metadata,
});
}));
inverse.insert((name, version, extra.as_ref(), dev.as_ref()), index);
}
}

// Add every edge to the graph.
for (names, version_set) in resolution.dependencies {
for versions in version_set {
let from_index = inverse[&(
&names.from,
&versions.from_version,
versions.from_extra.as_ref(),
versions.from_dev.as_ref(),
)];
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(),
)];

println!("from_index: {:?}", from_index);
println!("to_index: {:?}", to_index);
println!("versions.marker: {:?}", versions.marker);

if let Some(edge) = petgraph
.find_edge(from_index, to_index)
.and_then(|edge| petgraph.edge_weight_mut(edge))
Expand Down Expand Up @@ -295,25 +302,29 @@ impl ResolutionGraph {
})
}

/// Return the number of distinct packages in the graph.
pub fn len(&self) -> usize {
/// Returns an iterator over the distinct packages in the graph.
fn dists(&self) -> impl Iterator<Item = &AnnotatedDist> {
self.petgraph
.node_indices()
.map(|index| &self.petgraph[index])
.filter(|dist| dist.is_base())
.count()
.filter_map(move |index| match &self.petgraph[index] {
ResolutionGraphNode::Root => None,
ResolutionGraphNode::Dist(dist) => Some(dist),
})
}

/// Return the number of distinct packages in the graph.
pub fn len(&self) -> usize {
self.dists().filter(|dist| dist.is_base()).count()
}

/// Return `true` if there are no packages in the graph.
pub fn is_empty(&self) -> bool {
self.len() == 0
self.dists().any(super::AnnotatedDist::is_base)
}

/// Returns `true` if the graph contains the given package.
pub fn contains(&self, name: &PackageName) -> bool {
self.petgraph
.node_indices()
.any(|index| self.petgraph[index].name() == name)
self.dists().any(|dist| dist.name() == name)
}

/// Return the [`ResolutionDiagnostic`]s that were encountered while building the graph.
Expand Down Expand Up @@ -394,7 +405,9 @@ impl ResolutionGraph {

let mut seen_marker_values = IndexSet::default();
for i in self.petgraph.node_indices() {
let dist = &self.petgraph[i];
let ResolutionGraphNode::Dist(dist) = &self.petgraph[i] else {
continue;
};
let version_id = match dist.version_or_url() {
VersionOrUrlRef::Version(version) => {
VersionId::from_registry(dist.name().clone(), version.clone())
Expand Down Expand Up @@ -464,14 +477,8 @@ impl From<ResolutionGraph> for distribution_types::Resolution {
fn from(graph: ResolutionGraph) -> Self {
Self::new(
graph
.petgraph
.node_indices()
.map(|node| {
(
graph.petgraph[node].name().clone(),
graph.petgraph[node].dist.clone(),
)
})
.dists()
.map(|node| (node.name().clone(), node.dist.clone()))
.collect(),
graph.diagnostics,
)
Expand Down
1 change: 1 addition & 0 deletions crates/uv-resolver/src/resolution/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ use uv_normalize::{ExtraName, GroupName, PackageName};

pub use crate::resolution::display::{AnnotationStyle, DisplayResolutionGraph};
pub use crate::resolution::graph::ResolutionGraph;
pub(crate) use crate::resolution::graph::ResolutionGraphNode;
pub(crate) use crate::resolution::requirements_txt::RequirementsTxtDist;

mod display;
Expand Down
Loading

0 comments on commit 82c6ea1

Please sign in to comment.