diff --git a/Cargo.lock b/Cargo.lock index 7cdc815a1849..3efc03f58fa8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2673,7 +2673,7 @@ dependencies = [ [[package]] name = "pubgrub" version = "0.2.1" -source = "git+https://github.com/astral-sh/pubgrub?rev=57afc831bf2551f164617a10383cf288bf5d190d#57afc831bf2551f164617a10383cf288bf5d190d" +source = "git+https://github.com/astral-sh/pubgrub?rev=9cd9049a64c7352de2ff3b525b9ae36421b0cc18#9cd9049a64c7352de2ff3b525b9ae36421b0cc18" dependencies = [ "indexmap", "log", @@ -5666,7 +5666,7 @@ checksum = "830b7e5d4d90034032940e4ace0d9a9a057e7a45cd94e6c007832e39edb82f6d" [[package]] name = "version-ranges" version = "0.1.1" -source = "git+https://github.com/astral-sh/pubgrub?rev=57afc831bf2551f164617a10383cf288bf5d190d#57afc831bf2551f164617a10383cf288bf5d190d" +source = "git+https://github.com/astral-sh/pubgrub?rev=9cd9049a64c7352de2ff3b525b9ae36421b0cc18#9cd9049a64c7352de2ff3b525b9ae36421b0cc18" dependencies = [ "smallvec", ] diff --git a/Cargo.toml b/Cargo.toml index 59d80aebd180..a94db3161e9a 100644 --- a/Cargo.toml +++ b/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" } diff --git a/crates/uv-resolver/src/resolver/batch_prefetch.rs b/crates/uv-resolver/src/resolver/batch_prefetch.rs index ec2852109969..f62ca97f7c7b 100644 --- a/crates/uv-resolver/src/resolver/batch_prefetch.rs +++ b/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, - last_prefetch: FxHashMap, + tried_versions: FxHashMap, usize>, + last_prefetch: FxHashMap, 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, 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,15 +222,15 @@ 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, package: &PubGrubPackage) { // Only track base packages, no virtual packages from extras. if matches!( - &*package, + &**package, PubGrubPackageInner::Package { extra: None, dev: None, @@ -236,16 +238,16 @@ impl BatchPrefetcher { .. } ) { - *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) -> (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) { 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() diff --git a/crates/uv-resolver/src/resolver/derivation.rs b/crates/uv-resolver/src/resolver/derivation.rs index c0218d06dafc..42473ea2ae01 100644 --- a/crates/uv-resolver/src/resolver/derivation.rs +++ b/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, version: &Version, state: &State, ) -> Option { /// Find a path from the current package to the root package. fn find_path( - package: &PubGrubPackage, + id: Id, version: &Version, state: &State, - solution: &SelectedDependencies, + solution: &FxHashMap, Version>, path: &mut Vec, ) -> 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(); diff --git a/crates/uv-resolver/src/resolver/mod.rs b/crates/uv-resolver/src/resolver/mod.rs index cc93b08d971c..3200b0cc2caa 100644 --- a/crates/uv-resolver/src/resolver/mod.rs +++ b/crates/uv-resolver/src/resolver/mod.rs @@ -13,7 +13,7 @@ use dashmap::DashMap; use either::Either; use futures::{FutureExt, StreamExt}; use itertools::Itertools; -use pubgrub::{Incompatibility, Range, State}; +use pubgrub::{Id, Incompatibility, Range, State}; use rustc_hash::{FxHashMap, FxHashSet}; use tokio::sync::mpsc::{self, Receiver, Sender}; use tokio::sync::oneshot; @@ -308,13 +308,9 @@ impl ResolverState ResolverState ResolverState ResolverState ResolverState ResolverState ResolverState { - debug!("No compatible version found for: {next}", next = state.next); + debug!("No compatible version found for: {next_package}"); let term_intersection = state .pubgrub .partial_solution - .term_intersection_for_package(&state.next) + .term_intersection_for_package(next_id) .expect("a package was chosen but we don't have a term"); - if let PubGrubPackageInner::Package { ref name, .. } = &*state.next { + if let PubGrubPackageInner::Package { ref name, .. } = &**next_package { // Check if the decision was due to the package being unavailable if let Some(entry) = self.unavailable_packages.get(name) { state .pubgrub .add_incompatibility(Incompatibility::custom_term( - state.next.clone(), + next_id, term_intersection.clone(), UnavailableReason::Package(entry.clone()), )); @@ -465,7 +472,7 @@ impl ResolverState ResolverState ResolverState ResolverState ResolverState { state.add_package_version_dependencies( + next_id, &version, &self.urls, &self.indexes, @@ -675,7 +685,7 @@ impl ResolverState impl Iterator> + 'a { debug!( "Splitting resolution on {}=={} over {} into {} resolution{} with separate markers", - current_state.next, + current_state.pubgrub.package_store[current_state.next], version, diverging_packages .iter() @@ -690,6 +700,7 @@ impl ResolverState ResolverState ResolverState, package: &PubGrubPackage, version: &Version, fork_urls: &ForkUrls, @@ -1213,6 +1226,7 @@ impl ResolverState, ) -> Result { let result = self.get_dependencies( + id, package, version, fork_urls, @@ -1236,6 +1250,7 @@ impl ResolverState, package: &PubGrubPackage, version: &Version, fork_urls: &ForkUrls, @@ -1381,7 +1396,7 @@ impl ResolverState { - let chain = DerivationChainBuilder::from_state(package, version, pubgrub) + let chain = DerivationChainBuilder::from_state(id, version, pubgrub) .unwrap_or_default(); return Err(match &**dist { Dist::Built(built_dist @ BuiltDist::Path(_)) => ResolveError::Read( @@ -2160,7 +2175,7 @@ pub(crate) struct ForkState { /// assignments (to packages) from this state's "partial solution." pubgrub: State, /// The next package on which to run unit propagation. - next: PubGrubPackage, + next: Id, /// The set of pinned versions we accrue throughout resolution. /// /// The key of this map is a package name, and each package name maps to @@ -2190,7 +2205,7 @@ pub(crate) struct ForkState { priorities: PubGrubPriorities, /// This keeps track of the set of versions for each package that we've /// already visited during resolution. This avoids doing redundant work. - added_dependencies: FxHashMap>, + added_dependencies: FxHashMap, FxHashSet>, /// The marker expression that created this state. /// /// The root state always corresponds to a marker expression that is always @@ -2224,13 +2239,12 @@ pub(crate) struct ForkState { impl ForkState { fn new( pubgrub: State, - root: PubGrubPackage, env: ResolverEnvironment, python_requirement: PythonRequirement, ) -> Self { Self { + next: pubgrub.root_package, pubgrub, - next: root, pins: FilePins::default(), fork_urls: ForkUrls::default(), fork_indexes: ForkIndexes::default(), @@ -2245,6 +2259,7 @@ impl ForkState { /// self-dependencies and handling URLs. fn add_package_version_dependencies( &mut self, + for_package: Id, for_version: &Version, urls: &Urls, indexes: &Indexes, @@ -2277,7 +2292,7 @@ impl ForkState { } } - if let Some(name) = self.next.name_no_root() { + if let Some(name) = self.pubgrub.package_store[for_package].name_no_root() { debug!( "Adding transitive dependency for {name}=={for_version}: {package}{version}" ); @@ -2308,7 +2323,7 @@ impl ForkState { } self.pubgrub.add_package_version_dependencies( - self.next.clone(), + self.next, for_version.clone(), dependencies.into_iter().map(|dependency| { let PubGrubDependency { @@ -2333,26 +2348,26 @@ impl ForkState { ) = reason { let package = &self.next; + let python = self.pubgrub.package_store.alloc(PubGrubPackage::from( + PubGrubPackageInner::Python(match kind { + PythonRequirementKind::Installed => PubGrubPython::Installed, + PythonRequirementKind::Target => PubGrubPython::Target, + }), + )); self.pubgrub .add_incompatibility(Incompatibility::from_dependency( - package.clone(), + *package, Range::singleton(version.clone()), - ( - PubGrubPackage::from(PubGrubPackageInner::Python(match kind { - PythonRequirementKind::Installed => PubGrubPython::Installed, - PythonRequirementKind::Target => PubGrubPython::Target, - })), - release_specifiers_to_ranges(requires_python), - ), + (python, release_specifiers_to_ranges(requires_python)), )); self.pubgrub .partial_solution - .add_decision(self.next.clone(), version); + .add_decision(self.next, version); return; }; self.pubgrub .add_incompatibility(Incompatibility::custom_version( - self.next.clone(), + self.next, version.clone(), UnavailableReason::Version(reason), )); @@ -2374,32 +2389,35 @@ impl ForkState { } fn into_resolution(self) -> Resolution { - let solution = self.pubgrub.partial_solution.extract_solution(); + let solution: FxHashMap<_, _> = self.pubgrub.partial_solution.extract_solution().collect(); let mut edges: FxHashSet = FxHashSet::default(); for (package, self_version) in &solution { for id in &self.pubgrub.incompatibilities[package] { let pubgrub::Kind::FromDependencyOf( - ref self_package, + self_package, ref self_range, - ref dependency_package, + dependency_package, ref dependency_range, ) = self.pubgrub.incompatibility_store[*id].kind else { continue; }; - if package != self_package { + if *package != self_package { continue; } if !self_range.contains(self_version) { continue; } - let Some(dependency_version) = solution.get(dependency_package) else { + let Some(dependency_version) = solution.get(&dependency_package) else { continue; }; if !dependency_range.contains(dependency_version) { continue; } + let self_package = &self.pubgrub.package_store[self_package]; + let dependency_package = &self.pubgrub.package_store[dependency_package]; + let (self_name, self_extra, self_dev) = match &**self_package { PubGrubPackageInner::Package { name: self_name, @@ -2554,7 +2572,7 @@ impl ForkState { extra, dev, marker: None, - } = &*package + } = &*self.pubgrub.package_store[package] { Some(( ResolutionPackage {