Create distinct graph nodes for each package extra (#3908)

## Summary

Today, we represent each package as a single node in the graph, and
combine all the extras. This is helpful for the `requirements.txt`-style
resolution, in which we want to show each a single line for each package
with the extras combined into a single array.

This PR modifies the representation to instead use a separate node for
each (package, extra) pair. We then reduce into the previous format when
printing in the `requirements.txt`-style format, so there shouldn't be
any user-facing changes here.

crates/uv-resolver/src/resolution/display.rs

@@ -1,12 +1,15 @@
 use std::collections::BTreeSet;
+use std::hash::BuildHasherDefault;
 
 use owo_colors::OwoColorize;
 use petgraph::visit::EdgeRef;
 use petgraph::Direction;
+use rustc_hash::FxHashMap;
 
 use distribution_types::{Name, SourceAnnotations};
 use uv_normalize::PackageName;
 
+use crate::resolution::AnnotatedDist;
 use crate::ResolutionGraph;
 
 /// A [`std::fmt::Display`] implementation for the resolution graph.
@@ -77,13 +80,53 @@ impl<'a> DisplayResolutionGraph<'a> {
 /// Write the graph in the `{name}=={version}` format of requirements.txt that pip uses.
 impl std::fmt::Display for DisplayResolutionGraph<'_> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        // 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.
+        let petgraph = {
+            let mut petgraph =
+                petgraph::graph::Graph::<AnnotatedDist, (), petgraph::Directed>::with_capacity(
+                    self.resolution.petgraph.node_count(),
+                    self.resolution.petgraph.edge_count(),
+                );
+            let mut inverse = FxHashMap::with_capacity_and_hasher(
+                self.resolution.petgraph.node_count(),
+                BuildHasherDefault::default(),
+            );
+
+            // Re-add the nodes to the reduced graph.
+            for index in self.resolution.petgraph.node_indices() {
+                let dist = &self.resolution.petgraph[index];
+
+                if let Some(index) = inverse.get(dist.name()) {
+                    let node: &mut AnnotatedDist = &mut petgraph[*index];
+                    node.extras.extend(dist.extras.iter().cloned());
+                    node.extras.sort_unstable();
+                    node.extras.dedup();
+                } else {
+                    let index = petgraph.add_node(dist.clone());
+                    inverse.insert(dist.name(), index);
+                }
+            }
+
+            // Re-add the edges to the reduced graph.
+            for edge in self.resolution.petgraph.edge_indices() {
+                let (source, target) = self.resolution.petgraph.edge_endpoints(edge).unwrap();
+                let source = inverse[self.resolution.petgraph[source].name()];
+                let target = inverse[self.resolution.petgraph[target].name()];
+                petgraph.update_edge(source, target, ());
+            }
+
+            petgraph
+        };
+
         // Collect all packages.
-        let mut nodes = self
+        let mut nodes = petgraph
-            .resolution
-            .petgraph
             .node_indices()
             .filter_map(|index| {
-                let dist = &self.resolution.petgraph[index];
+                let dist = &petgraph[index];
                 let name = dist.name();
                 if self.no_emit_packages.contains(name) {
                     return None;
@@ -119,11 +162,9 @@ impl std::fmt::Display for DisplayResolutionGraph<'_> {
             // package (e.g., `# via mypy`).
             if self.include_annotations {
                 // Display all dependencies.
-                let mut edges = self
+                let mut edges = petgraph
-                    .resolution
-                    .petgraph
                     .edges_directed(index, Direction::Incoming)
-                    .map(|edge| &self.resolution.petgraph[edge.source()])
+                    .map(|edge| &petgraph[edge.source()])
                     .collect::<Vec<_>>();
                 edges.sort_unstable_by_key(|package| package.name());
 

crates/uv-resolver/src/resolution/graph.rs

@@ -46,102 +46,17 @@ impl ResolutionGraph {
         state: &State<UvDependencyProvider>,
         preferences: &Preferences,
     ) -> anyhow::Result<Self, ResolveError> {
-        // Collect and validate the extras.
-        let mut extras = FxHashMap::default();
-        let mut diagnostics = Vec::new();
-        for (package, version) in selection {
-            match &**package {
-                PubGrubPackageInner::Package {
-                    name,
-                    extra: Some(extra),
-                    marker: None,
-                    url: None,
-                } => {
-                    let dist = PubGrubDistribution::from_registry(name, version);
-
-                    let response = distributions.get(&dist.version_id()).unwrap_or_else(|| {
-                        panic!(
-                            "Every package should have metadata: {:?}",
-                            dist.version_id()
-                        )
-                    });
-
-                    let MetadataResponse::Found(archive) = &*response else {
-                        panic!(
-                            "Every package should have metadata: {:?}",
-                            dist.version_id()
-                        )
-                    };
-
-                    if archive.metadata.provides_extras.contains(extra) {
-                        extras
-                            .entry(name.clone())
-                            .or_insert_with(Vec::new)
-                            .push(extra.clone());
-                    } else {
-                        let dist = pins
-                            .get(name, version)
-                            .unwrap_or_else(|| panic!("Every package should be pinned: {name:?}"))
-                            .clone();
-
-                        diagnostics.push(ResolutionDiagnostic::MissingExtra {
-                            dist,
-                            extra: extra.clone(),
-                        });
-                    }
-                }
-                PubGrubPackageInner::Package {
-                    name,
-                    extra: Some(extra),
-                    marker: None,
-                    url: Some(url),
-                } => {
-                    let dist = PubGrubDistribution::from_url(name, url);
-
-                    let response = distributions.get(&dist.version_id()).unwrap_or_else(|| {
-                        panic!(
-                            "Every package should have metadata: {:?}",
-                            dist.version_id()
-                        )
-                    });
-
-                    let MetadataResponse::Found(archive) = &*response else {
-                        panic!(
-                            "Every package should have metadata: {:?}",
-                            dist.version_id()
-                        )
-                    };
-
-                    if archive.metadata.provides_extras.contains(extra) {
-                        extras
-                            .entry(name.clone())
-                            .or_insert_with(Vec::new)
-                            .push(extra.clone());
-                    } else {
-                        let dist = Dist::from_url(name.clone(), url_to_precise(url.clone()))?;
-
-                        diagnostics.push(ResolutionDiagnostic::MissingExtra {
-                            dist: dist.into(),
-                            extra: extra.clone(),
-                        });
-                    }
-                }
-                _ => {}
-            };
-        }
-
         // Add every package to the graph.
-        // TODO(charlie): petgraph is a really heavy and unnecessary dependency here. We should
-        // write our own graph, given that our requirements are so simple.
         let mut petgraph = petgraph::graph::Graph::with_capacity(selection.len(), selection.len());
         let mut inverse =
             FxHashMap::with_capacity_and_hasher(selection.len(), BuildHasherDefault::default());
+        let mut diagnostics = Vec::new();
 
         for (package, version) in selection {
             match &**package {
                 PubGrubPackageInner::Package {
                     name,
-                    extra: None,
+                    extra,
                     marker: None,
                     url: None,
                 } => {
@@ -213,8 +128,18 @@ impl ResolutionGraph {
                         archive.metadata.clone()
                     };
 
+                    // Validate the extra.
+                    if let Some(extra) = extra {
+                        if !metadata.provides_extras.contains(extra) {
+                            diagnostics.push(ResolutionDiagnostic::MissingExtra {
+                                dist: dist.clone(),
+                                extra: extra.clone(),
+                            });
+                        }
+                    }
+
                     // Extract the extras.
-                    let extras = extras.get(name).cloned().unwrap_or_default();
+                    let extras = extra.clone().map(|extra| vec![extra]).unwrap_or_default();
 
                     // Add the distribution to the graph.
                     let index = petgraph.add_node(AnnotatedDist {
@@ -223,16 +148,16 @@ impl ResolutionGraph {
                         hashes,
                         metadata,
                     });
-                    inverse.insert(name, index);
+                    inverse.insert((name, extra), index);
                 }
+
                 PubGrubPackageInner::Package {
                     name,
-                    extra: None,
+                    extra,
                     marker: None,
                     url: Some(url),
                 } => {
                     // Create the distribution.
-
                     let dist = Dist::from_url(name.clone(), url_to_precise(url.clone()))?;
 
                     // Extract the hashes, preserving those that were already present in the
@@ -275,8 +200,18 @@ impl ResolutionGraph {
                         archive.metadata.clone()
                     };
 
+                    // Validate the extra.
+                    if let Some(extra) = extra {
+                        if !metadata.provides_extras.contains(extra) {
+                            diagnostics.push(ResolutionDiagnostic::MissingExtra {
+                                dist: dist.clone().into(),
+                                extra: extra.clone(),
+                            });
+                        }
+                    }
+
                     // Extract the extras.
-                    let extras = extras.get(name).cloned().unwrap_or_default();
+                    let extras = extra.clone().map(|extra| vec![extra]).unwrap_or_default();
 
                     // Add the distribution to the graph.
                     let index = petgraph.add_node(AnnotatedDist {
@@ -285,8 +220,9 @@ impl ResolutionGraph {
                         hashes,
                         metadata,
                     });
-                    inverse.insert(name, index);
+                    inverse.insert((name, extra), index);
                 }
+
                 _ => {}
             };
         }
@@ -305,14 +241,22 @@ impl ResolutionGraph {
                         continue;
                     }
 
+                    if !self_version.contains(version) {
+                        continue;
+                    }
+
                     let PubGrubPackageInner::Package {
-                        name: self_name, ..
+                        name: self_name,
+                        extra: self_extra,
+                        ..
                     } = &**self_package
                     else {
                         continue;
                     };
+
                     let PubGrubPackageInner::Package {
                         name: dependency_name,
+                        extra: dependency_extra,
                         ..
                     } = &**dependency_package
                     else {
@@ -324,14 +268,12 @@ impl ResolutionGraph {
                         continue;
                     }
 
-                    if self_version.contains(version) {
+                    let self_index = &inverse[&(self_name, self_extra)];
-                        let self_index = &inverse[self_name];
+                    let dependency_index = &inverse[&(dependency_name, dependency_extra)];
-                        let dependency_index = &inverse[dependency_name];
                     petgraph.update_edge(*self_index, *dependency_index, ());
                 }
             }
         }
-        }
 
         Ok(Self {
             petgraph,
@@ -339,14 +281,18 @@ impl ResolutionGraph {
         })
     }
 
-    /// Return the number of packages in the graph.
+    /// Return the number of distinct packages in the graph.
     pub fn len(&self) -> usize {
-        self.petgraph.node_count()
+        self.petgraph
+            .node_indices()
+            .map(|index| &self.petgraph[index])
+            .filter(|dist| dist.extras.is_empty())
+            .count()
     }
 
     /// Return `true` if there are no packages in the graph.
     pub fn is_empty(&self) -> bool {
-        self.petgraph.node_count() == 0
+        self.len() == 0
     }
 
     /// Returns `true` if the graph contains the given package.