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.

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

Remove propagate_markers #7076

Merged
merged 1 commit into from
Sep 5, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
82 changes: 2 additions & 80 deletions crates/uv-resolver/src/graph_ops.rs
Original file line number Diff line number Diff line change
@@ -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<T: Markers>(
mut graph: Graph<T, MarkerTree, Directed>,
) -> Graph<T, MarkerTree, Directed> {
// 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::<Vec<_>>();
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:
Expand Down
61 changes: 19 additions & 42 deletions crates/uv-resolver/src/lock/requirements_txt.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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<Node<'lock>, 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<NodeIndex, MarkerTree>,
hashes: bool,
}

Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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.
Expand All @@ -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::<Vec<_>>();
let mut nodes = self.graph.node_references().collect::<Vec<_>>();

// 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)?;
Expand Down Expand Up @@ -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() {
Copy link
Member

Choose a reason for hiding this comment

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

Nice!

write!(f, " ; {contents}")?;
}

Expand All @@ -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;

Expand Down
Loading
Loading