Migrate to PubGrub's arena for package names

Cargo.toml

@@ -130,7 +130,7 @@ petgraph = { version = "0.6.5" }
 platform-info = { version = "2.0.3" }
 proc-macro2 = { version = "1.0.86" }
 procfs = { version = "0.17.0", default-features = false, features = ["flate2"] }
-pubgrub = { git = "https://github.com/astral-sh/pubgrub", rev = "57afc831bf2551f164617a10383cf288bf5d190d" }
+pubgrub = { git = "https://github.com/astral-sh/pubgrub", rev = "9cd9049a64c7352de2ff3b525b9ae36421b0cc18" }
 quote = { version = "1.0.37" }
 rayon = { version = "1.10.0" }
 reflink-copy = { version = "0.1.19" }
@@ -175,7 +175,7 @@ unicode-width = { version = "0.1.13" }
 unscanny = { version = "0.1.0" }
 url = { version = "2.5.2", features = ["serde"] }
 urlencoding = { version = "2.1.3" }
-version-ranges = { git = "https://github.com/astral-sh/pubgrub", rev = "57afc831bf2551f164617a10383cf288bf5d190d" }
+version-ranges = { git = "https://github.com/astral-sh/pubgrub", rev = "9cd9049a64c7352de2ff3b525b9ae36421b0cc18" }
 walkdir = { version = "2.5.0" }
 which = { version = "7.0.0", features = ["regex"] }
 windows-registry = { version = "0.3.0" }

crates/uv-resolver/src/resolver/batch_prefetch.rs

@@ -1,12 +1,13 @@
 use std::cmp::min;
 
 use itertools::Itertools;
-use pubgrub::{Range, Term};
+use pubgrub::{Id, Range, State, Term};
 use rustc_hash::FxHashMap;
 use tokio::sync::mpsc::Sender;
 use tracing::{debug, trace};
 
 use crate::candidate_selector::CandidateSelector;
+use crate::dependency_provider::UvDependencyProvider;
 use crate::pubgrub::{PubGrubPackage, PubGrubPackageInner};
 use crate::resolver::Request;
 use crate::{
@@ -39,14 +40,15 @@ enum BatchPrefetchStrategy {
 /// Note that these all heuristics that could totally prefetch lots of irrelevant versions.
 #[derive(Default)]
 pub(crate) struct BatchPrefetcher {
-    tried_versions: FxHashMap<PubGrubPackage, usize>,
-    last_prefetch: FxHashMap<PubGrubPackage, usize>,
+    tried_versions: FxHashMap<Id<PubGrubPackage>, usize>,
+    last_prefetch: FxHashMap<Id<PubGrubPackage>, usize>,
 }
 
 impl BatchPrefetcher {
     /// Prefetch a large number of versions if we already unsuccessfully tried many versions.
     pub(crate) fn prefetch_batches(
         &mut self,
+        id: Id<PubGrubPackage>,
         next: &PubGrubPackage,
         index: Option<&IndexUrl>,
         version: &Version,
@@ -69,7 +71,7 @@ impl BatchPrefetcher {
             return Ok(());
         };
 
-        let (num_tried, do_prefetch) = self.should_prefetch(next);
+        let (num_tried, do_prefetch) = self.should_prefetch(id);
         if !do_prefetch {
             return Ok(());
         }
@@ -220,32 +222,32 @@ impl BatchPrefetcher {
 
         debug!("Prefetching {prefetch_count} {name} versions");
 
-        self.last_prefetch.insert(next.clone(), num_tried);
+        self.last_prefetch.insert(id, num_tried);
         Ok(())
     }
 
     /// Each time we tried a version for a package, we register that here.
-    pub(crate) fn version_tried(&mut self, package: PubGrubPackage) {
+    pub(crate) fn version_tried(&mut self, id: Id<PubGrubPackage>, package: &PubGrubPackage) {
         // Only track base packages, no virtual packages from extras.
         if matches!(
-            &*package,
+            &**package,
             PubGrubPackageInner::Package {
                 extra: None,
                 dev: None,
                 marker: None,
                 ..
             }
         ) {
-            *self.tried_versions.entry(package).or_default() += 1;
+            *self.tried_versions.entry(id).or_default() += 1;
         }
     }
 
     /// After 5, 10, 20, 40 tried versions, prefetch that many versions to start early but not
     /// too aggressive. Later we schedule the prefetch of 50 versions every 20 versions, this gives
     /// us a good buffer until we see prefetch again and is high enough to saturate the task pool.
-    fn should_prefetch(&self, next: &PubGrubPackage) -> (usize, bool) {
-        let num_tried = self.tried_versions.get(next).copied().unwrap_or_default();
-        let previous_prefetch = self.last_prefetch.get(next).copied().unwrap_or_default();
+    fn should_prefetch(&self, id: Id<PubGrubPackage>) -> (usize, bool) {
+        let num_tried = self.tried_versions.get(&id).copied().unwrap_or_default();
+        let previous_prefetch = self.last_prefetch.get(&id).copied().unwrap_or_default();
         let do_prefetch = (num_tried >= 5 && previous_prefetch < 5)
             || (num_tried >= 10 && previous_prefetch < 10)
             || (num_tried >= 20 && previous_prefetch < 20)
@@ -257,9 +259,13 @@ impl BatchPrefetcher {
     ///
     /// Note that they may be inflated when we count the same version repeatedly during
     /// backtracking.
-    pub(crate) fn log_tried_versions(&self) {
+    pub(crate) fn log_tried_versions(&self, state: &State<UvDependencyProvider>) {
         let total_versions: usize = self.tried_versions.values().sum();
-        let mut tried_versions: Vec<_> = self.tried_versions.iter().collect();
+        let mut tried_versions: Vec<_> = self
+            .tried_versions
+            .iter()
+            .map(|(id, count)| (&state.package_store[*id], *count))
+            .collect();
         tried_versions.sort_by(|(p1, c1), (p2, c2)| {
             c1.cmp(c2)
                 .reverse()

crates/uv-resolver/src/resolver/derivation.rs

@@ -2,8 +2,8 @@ use std::collections::VecDeque;
 
 use petgraph::visit::EdgeRef;
 use petgraph::Direction;
-use pubgrub::{Kind, Range, SelectedDependencies, State};
-use rustc_hash::FxHashSet;
+use pubgrub::{Id, Kind, Range, State};
+use rustc_hash::{FxHashMap, FxHashSet};
 
 use uv_distribution_types::{
     DerivationChain, DerivationStep, DistRef, Edge, Name, Node, Resolution, ResolvedDist,
@@ -87,32 +87,35 @@ impl DerivationChainBuilder {
     ///
     /// This is used to construct a derivation chain upon resolution failure.
     pub(crate) fn from_state(
-        package: &PubGrubPackage,
+        id: Id<PubGrubPackage>,
         version: &Version,
         state: &State<UvDependencyProvider>,
     ) -> Option<DerivationChain> {
         /// Find a path from the current package to the root package.
         fn find_path(
-            package: &PubGrubPackage,
+            id: Id<PubGrubPackage>,
             version: &Version,
             state: &State<UvDependencyProvider>,
-            solution: &SelectedDependencies<UvDependencyProvider>,
+            solution: &FxHashMap<Id<PubGrubPackage>, Version>,
             path: &mut Vec<DerivationStep>,
         ) -> bool {
             // Retrieve the incompatibilities for the current package.
-            let Some(incompatibilities) = state.incompatibilities.get(package) else {
+            let Some(incompatibilities) = state.incompatibilities.get(&id) else {
                 return false;
             };
             for index in incompatibilities {
                 let incompat = &state.incompatibility_store[*index];
 
                 // Find a dependency from a package to the current package.
-                if let Kind::FromDependencyOf(p1, _, p2, v2) = &incompat.kind {
-                    if p2 == package && v2.contains(version) {
-                        if let Some(version) = solution.get(p1) {
+                if let Kind::FromDependencyOf(id1, _, id2, v2) = &incompat.kind {
+                    if id == *id2 && v2.contains(version) {
+                        if let Some(version) = solution.get(id1) {
+                            let p1 = &state.package_store[*id1];
+                            let p2 = &state.package_store[*id2];
+
                             if p1.name_no_root() == p2.name_no_root() {
                                 // Skip proxied dependencies.
-                                if find_path(p1, version, state, solution, path) {
+                                if find_path(*id1, version, state, solution, path) {
                                     return true;
                                 }
                             } else if let Some(name) = p1.name_no_root() {
@@ -126,7 +129,7 @@ impl DerivationChainBuilder {
                                 ));
 
                                 // Recursively search the next package.
-                                if find_path(p1, version, state, solution, path) {
+                                if find_path(*id1, version, state, solution, path) {
                                     return true;
                                 }
 
@@ -143,10 +146,10 @@ impl DerivationChainBuilder {
             false
         }
 
-        let solution = state.partial_solution.extract_solution();
+        let solution: FxHashMap<_, _> = state.partial_solution.extract_solution().collect();
         let path = {
             let mut path = vec![];
-            if !find_path(package, version, state, &solution, &mut path) {
+            if !find_path(id, version, state, &solution, &mut path) {
                 return None;
             }
             path.reverse();