From b7ff4500b7527dd7c90e841419fd7fd583da7137 Mon Sep 17 00:00:00 2001 From: konstin Date: Thu, 5 Sep 2024 18:33:11 +0200 Subject: [PATCH] Remove propagate_markers --- crates/uv-resolver/src/graph_ops.rs | 82 +-------- .../uv-resolver/src/lock/requirements_txt.rs | 61 +++---- crates/uv-resolver/src/resolution/display.rs | 155 +++++++----------- .../src/resolution/requirements_txt.rs | 8 +- 4 files changed, 88 insertions(+), 218 deletions(-) diff --git a/crates/uv-resolver/src/graph_ops.rs b/crates/uv-resolver/src/graph_ops.rs index 7d75b97f7f62..7a88674acf68 100644 --- a/crates/uv-resolver/src/graph_ops.rs +++ b/crates/uv-resolver/src/graph_ops.rs @@ -1,88 +1,10 @@ use pep508_rs::MarkerTree; -use petgraph::algo::greedy_feedback_arc_set; use petgraph::graph::NodeIndex; -use petgraph::visit::{EdgeRef, Topo}; -use petgraph::{Directed, Direction, Graph}; +use petgraph::visit::EdgeRef; +use petgraph::{Direction, Graph}; use rustc_hash::FxHashMap; use std::collections::hash_map::Entry; -/// A trait for a graph node that can be annotated with a [`MarkerTree`]. -pub(crate) trait Markers { - fn set_markers(&mut self, markers: MarkerTree); -} - -/// Propagate the [`MarkerTree`] qualifiers across the graph. -/// -/// The graph is directed, so if any edge contains a marker, we need to propagate it to all -/// downstream nodes. -pub(crate) fn propagate_markers( - mut graph: Graph, -) -> Graph { - // Remove any cycles. By absorption, it should be fine to ignore cycles. - // - // Imagine a graph: `A -> B -> C -> A`. Assume that `A` has weight `1`, `B` has weight `2`, - // and `C` has weight `3`. The weights are the marker trees. - // - // When propagating, we'd return to `A` when we hit the cycle, to create `1 or (1 and 2 and 3)`, - // which resolves to `1`. - // - // TODO(charlie): The above reasoning could be incorrect. Consider using a graph algorithm that - // can handle weight propagation with cycles. - let edges = { - let mut fas = greedy_feedback_arc_set(&graph) - .map(|edge| edge.id()) - .collect::>(); - fas.sort_unstable(); - let mut edges = Vec::with_capacity(fas.len()); - for edge_id in fas.into_iter().rev() { - edges.push(graph.edge_endpoints(edge_id).unwrap()); - graph.remove_edge(edge_id); - } - edges - }; - - let mut topo = Topo::new(&graph); - while let Some(index) = topo.next(&graph) { - let marker_tree = { - // Fold over the edges to combine the marker trees. If any edge is `None`, then - // the combined marker tree is `None`. - let mut edges = graph.edges_directed(index, Direction::Incoming); - - edges - .next() - .and_then(|edge| graph.edge_weight(edge.id()).cloned()) - .and_then(|initial| { - edges.try_fold(initial, |mut acc, edge| { - acc.or(graph.edge_weight(edge.id())?.clone()); - Some(acc) - }) - }) - .unwrap_or_default() - }; - - // Propagate the marker tree to all downstream nodes. - let mut walker = graph - .neighbors_directed(index, Direction::Outgoing) - .detach(); - while let Some((outgoing, _)) = walker.next(&graph) { - if let Some(weight) = graph.edge_weight_mut(outgoing) { - weight.and(marker_tree.clone()); - } - } - - let node = &mut graph[index]; - node.set_markers(marker_tree); - } - - // Re-add the removed edges. We no longer care about the edge _weights_, but we do want the - // edges to be present, to power the `# via` annotations. - for (source, target) in edges { - graph.add_edge(source, target, MarkerTree::TRUE); - } - - graph -} - /// Determine the markers under which a package is reachable in the dependency tree. /// /// The algorithm is a variant of Dijkstra's algorithm for not totally ordered distances: diff --git a/crates/uv-resolver/src/lock/requirements_txt.rs b/crates/uv-resolver/src/lock/requirements_txt.rs index 19107f78e422..516cc4470393 100644 --- a/crates/uv-resolver/src/lock/requirements_txt.rs +++ b/crates/uv-resolver/src/lock/requirements_txt.rs @@ -4,6 +4,8 @@ use std::fmt::Formatter; use std::path::{Path, PathBuf}; use either::Either; +use petgraph::graph::NodeIndex; +use petgraph::visit::IntoNodeReferences; use petgraph::{Directed, Graph}; use rustc_hash::{FxHashMap, FxHashSet}; use url::Url; @@ -16,16 +18,17 @@ use uv_fs::Simplified; use uv_git::GitReference; use uv_normalize::{ExtraName, GroupName, PackageName}; -use crate::graph_ops::{propagate_markers, Markers}; +use crate::graph_ops::marker_reachability; use crate::lock::{Package, PackageId, Source}; use crate::{Lock, LockError}; -type LockGraph<'lock> = Graph, Edge, Directed>; +type LockGraph<'lock> = Graph<&'lock Package, Edge, Directed>; /// An export of a [`Lock`] that renders in `requirements.txt` format. #[derive(Debug)] pub struct RequirementsTxtExport<'lock> { graph: LockGraph<'lock>, + reachability: FxHashMap, hashes: bool, } @@ -68,7 +71,7 @@ impl<'lock> RequirementsTxtExport<'lock> { } // Add the root package to the graph. - inverse.insert(&root.id, petgraph.add_node(Node::from_package(root))); + inverse.insert(&root.id, petgraph.add_node(root)); // Create all the relevant nodes. let mut seen = FxHashSet::default(); @@ -97,7 +100,7 @@ impl<'lock> RequirementsTxtExport<'lock> { // Add the dependency to the graph. if let Entry::Vacant(entry) = inverse.entry(&dep.package_id) { - entry.insert(petgraph.add_node(Node::from_package(dep_dist))); + entry.insert(petgraph.add_node(dep_dist)); } // Add the edge. @@ -120,31 +123,28 @@ impl<'lock> RequirementsTxtExport<'lock> { } } - let graph = propagate_markers(petgraph); + let reachability = marker_reachability(&petgraph, &[]); - Ok(Self { graph, hashes }) + Ok(Self { + graph: petgraph, + reachability, + hashes, + }) } } impl std::fmt::Display for RequirementsTxtExport<'_> { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { // Collect all packages. - let mut nodes = self - .graph - .raw_nodes() - .iter() - .map(|node| &node.weight) - .collect::>(); + let mut nodes = self.graph.node_references().collect::>(); // Sort the nodes, such that unnamed URLs (editables) appear at the top. - nodes.sort_unstable_by(|a, b| { - NodeComparator::from(a.package).cmp(&NodeComparator::from(b.package)) + nodes.sort_unstable_by(|(_, a), (_, b)| { + NodeComparator::from(**a).cmp(&NodeComparator::from(**b)) }); - // Write out each node. - for node in nodes { - let Node { package, markers } = node; - + // Write out each package. + for (node_index, package) in nodes { match &package.id.source { Source::Registry(_) => { write!(f, "{}=={}", package.id.name, package.id.version)?; @@ -201,7 +201,7 @@ impl std::fmt::Display for RequirementsTxtExport<'_> { } } - if let Some(contents) = markers.contents() { + if let Some(contents) = self.reachability[&node_index].contents() { write!(f, " ; {contents}")?; } @@ -223,29 +223,6 @@ impl std::fmt::Display for RequirementsTxtExport<'_> { } } -/// The nodes of the [`LockGraph`]. -#[derive(Debug)] -struct Node<'lock> { - package: &'lock Package, - markers: MarkerTree, -} - -impl<'lock> Node<'lock> { - /// Construct a [`Node`] from a [`Package`]. - fn from_package(package: &'lock Package) -> Self { - Self { - package, - markers: MarkerTree::default(), - } - } -} - -impl Markers for Node<'_> { - fn set_markers(&mut self, markers: MarkerTree) { - self.markers = markers; - } -} - /// The edges of the [`LockGraph`]. type Edge = MarkerTree; diff --git a/crates/uv-resolver/src/resolution/display.rs b/crates/uv-resolver/src/resolution/display.rs index 705d1239f228..a9766b87af8b 100644 --- a/crates/uv-resolver/src/resolution/display.rs +++ b/crates/uv-resolver/src/resolution/display.rs @@ -1,6 +1,7 @@ use std::collections::BTreeSet; use owo_colors::OwoColorize; +use petgraph::graph::NodeIndex; use petgraph::visit::EdgeRef; use petgraph::Direction; use rustc_hash::{FxBuildHasher, FxHashMap}; @@ -9,14 +10,9 @@ use distribution_types::{DistributionMetadata, Name, SourceAnnotation, SourceAnn use pep508_rs::MarkerTree; use uv_normalize::PackageName; -use crate::graph_ops::{propagate_markers, Markers}; use crate::resolution::{RequirementsTxtDist, ResolutionGraphNode}; use crate::{ResolutionGraph, ResolverMarkers}; -static UNIVERSAL_MARKERS: ResolverMarkers = ResolverMarkers::Universal { - fork_preferences: vec![], -}; - /// A [`std::fmt::Display`] implementation for the resolution graph. #[derive(Debug)] #[allow(clippy::struct_excessive_bools)] @@ -49,30 +45,6 @@ enum DisplayResolutionGraphNode { Dist(RequirementsTxtDist), } -impl Markers for DisplayResolutionGraphNode { - fn set_markers(&mut self, markers: MarkerTree) { - if let DisplayResolutionGraphNode::Dist(node) = self { - node.markers = markers; - } - } -} - -impl<'a> From<&'a ResolutionGraph> for DisplayResolutionGraph<'a> { - fn from(resolution: &'a ResolutionGraph) -> Self { - Self::new( - resolution, - &UNIVERSAL_MARKERS, - &[], - false, - false, - false, - true, - false, - AnnotationStyle::default(), - ) - } -} - impl<'a> DisplayResolutionGraph<'a> { /// Create a new [`DisplayResolutionGraph`] for the given graph. #[allow(clippy::fn_params_excessive_bools)] @@ -156,15 +128,30 @@ impl std::fmt::Display for DisplayResolutionGraph<'_> { SourceAnnotations::default() }; - // Convert from `AnnotatedDist` to `RequirementsTxtDist`. - let petgraph = to_requirements_txt_graph(&self.resolution.petgraph); - - // Propagate markers across the graph. - let petgraph = propagate_markers(petgraph); + // Convert a [`petgraph::graph::Graph`] based on [`ResolutionGraphNode`] to a graph based on + // [`DisplayResolutionGraphNode`]. In other words: converts from [`AnnotatedDist`] to + // [`RequirementsTxtDist`]. + // + // We assign each package its propagated markers: In `requirements.txt`, we want a flat list + // that for each package tells us if it should be installed on the current platform, without + // looking at which packages depend on it. + let petgraph = self.resolution.petgraph.map( + |index, node| match node { + ResolutionGraphNode::Root => DisplayResolutionGraphNode::Root, + ResolutionGraphNode::Dist(dist) => { + let reachability = self.resolution.reachability[&index].clone(); + let dist = RequirementsTxtDist::from_annotated_dist(dist, reachability); + DisplayResolutionGraphNode::Dist(dist) + } + }, + // We can drop the edge markers, while retaining their existence and direction for the + // annotations. + |_index, _edge| (), + ); // Reduce the graph, such that all nodes for a single package are combined, regardless of // the extras. - let petgraph = combine_extras(&petgraph); + let petgraph = combine_extras(&petgraph, &self.resolution.reachability); // Collect all packages. let mut nodes = petgraph @@ -318,47 +305,11 @@ pub enum AnnotationStyle { Split, } -type ResolutionPetGraph = - petgraph::graph::Graph; - +/// We don't need the edge markers anymore since we switched to propagated markers. type IntermediatePetGraph = - petgraph::graph::Graph; + petgraph::graph::Graph; -type RequirementsTxtGraph = - petgraph::graph::Graph; - -/// Convert a [`petgraph::graph::Graph`] based on [`ResolutionGraphNode`] to a graph based on -/// [`DisplayResolutionGraphNode`]. -/// -/// In other words: converts from [`AnnotatedDist`] to [`RequirementsTxtDist`]. -fn to_requirements_txt_graph(graph: &ResolutionPetGraph) -> IntermediatePetGraph { - let mut next = IntermediatePetGraph::with_capacity(graph.node_count(), graph.edge_count()); - let mut inverse = FxHashMap::with_capacity_and_hasher(graph.node_count(), FxBuildHasher); - - // Re-add the nodes to the reduced graph. - for index in graph.node_indices() { - match &graph[index] { - ResolutionGraphNode::Root => { - inverse.insert(index, next.add_node(DisplayResolutionGraphNode::Root)); - } - ResolutionGraphNode::Dist(dist) => { - let dist = RequirementsTxtDist::from(dist); - inverse.insert(index, next.add_node(DisplayResolutionGraphNode::Dist(dist))); - } - } - } - - // Re-add the edges to the reduced graph. - for edge in graph.edge_indices() { - let (source, target) = graph.edge_endpoints(edge).unwrap(); - let weight = graph[edge].clone(); - let source = inverse[&source]; - let target = inverse[&target]; - next.update_edge(source, target, weight); - } - - next -} +type RequirementsTxtGraph = petgraph::graph::Graph; /// Reduce the graph, such that all nodes for a single package are combined, regardless of /// the extras. @@ -367,7 +318,10 @@ fn to_requirements_txt_graph(graph: &ResolutionPetGraph) -> IntermediatePetGraph /// node. /// /// We also remove the root node, to simplify the graph structure. -fn combine_extras(graph: &IntermediatePetGraph) -> RequirementsTxtGraph { +fn combine_extras( + graph: &IntermediatePetGraph, + reachability: &FxHashMap, +) -> RequirementsTxtGraph { let mut next = RequirementsTxtGraph::with_capacity(graph.node_count(), graph.edge_count()); let mut inverse = FxHashMap::with_capacity_and_hasher(graph.node_count(), FxBuildHasher); @@ -377,14 +331,45 @@ fn combine_extras(graph: &IntermediatePetGraph) -> RequirementsTxtGraph { continue; }; + // In the `requirements.txt` output, we want a flat installation list, so we need to use + // the reachability markers instead of the edge markers. + // We use the markers of the base package: We know that each virtual extra package has an + // edge to the base package, so we know that base package markers are more general than the + // extra package markers (the extra package markers are a subset of the base package + // markers). if let Some(index) = inverse.get(&dist.version_id()) { let node: &mut RequirementsTxtDist = &mut next[*index]; node.extras.extend(dist.extras.iter().cloned()); node.extras.sort_unstable(); node.extras.dedup(); + if dist.extras.is_empty() { + node.markers = dist.markers.clone(); + } } else { - let index = next.add_node(dist.clone()); - inverse.insert(dist.version_id(), index); + let version_id = dist.version_id(); + let dist = dist.clone(); + let index = next.add_node(dist); + inverse.insert(version_id, index); + } + } + + // Verify that the package markers are more general than the extra markers. + if cfg!(debug_assertions) { + for index in graph.node_indices() { + let DisplayResolutionGraphNode::Dist(dist) = &graph[index] else { + continue; + }; + let combined_markers = next[inverse[&dist.version_id()]].markers.clone(); + let mut package_markers = combined_markers.clone(); + package_markers.or(reachability[&index].clone()); + assert_eq!( + package_markers, + combined_markers, + "{} {:?} {:?}", + dist.version_id(), + dist.extras, + dist.markers.try_to_string() + ); } } @@ -397,24 +382,10 @@ fn combine_extras(graph: &IntermediatePetGraph) -> RequirementsTxtGraph { let DisplayResolutionGraphNode::Dist(target_node) = &graph[target] else { continue; }; - let weight = graph[edge].clone(); let source = inverse[&source_node.version_id()]; let target = inverse[&target_node.version_id()]; - // If either the existing marker or new marker is `true`, then the dependency is - // included unconditionally, and so the combined marker should be `true`. - if let Some(edge) = next - .find_edge(source, target) - .and_then(|edge| next.edge_weight_mut(edge)) - { - if edge.is_true() || weight.is_true() { - *edge = MarkerTree::TRUE; - } else { - edge.and(weight.clone()); - } - } else { - next.update_edge(source, target, weight); - } + next.update_edge(source, target, ()); } next diff --git a/crates/uv-resolver/src/resolution/requirements_txt.rs b/crates/uv-resolver/src/resolution/requirements_txt.rs index 73ef1d7f0cbf..d80f683ebb10 100644 --- a/crates/uv-resolver/src/resolution/requirements_txt.rs +++ b/crates/uv-resolver/src/resolution/requirements_txt.rs @@ -22,6 +22,8 @@ pub(crate) struct RequirementsTxtDist { pub(crate) version: Version, pub(crate) extras: Vec, pub(crate) hashes: Vec, + /// Propagated markers that determine whether this package should be installed on the current + /// platform, without looking at which packages depend on it. pub(crate) markers: MarkerTree, } @@ -162,10 +164,8 @@ impl RequirementsTxtDist { } } } -} -impl From<&AnnotatedDist> for RequirementsTxtDist { - fn from(annotated: &AnnotatedDist) -> Self { + pub(crate) fn from_annotated_dist(annotated: &AnnotatedDist, markers: MarkerTree) -> Self { Self { dist: annotated.dist.clone(), version: annotated.version.clone(), @@ -175,7 +175,7 @@ impl From<&AnnotatedDist> for RequirementsTxtDist { vec![] }, hashes: annotated.hashes.clone(), - markers: MarkerTree::default(), + markers, } } }