diff --git a/crates/uv-resolver/src/lib.rs b/crates/uv-resolver/src/lib.rs index 629539bac..3e4b5f22c 100644 --- a/crates/uv-resolver/src/lib.rs +++ b/crates/uv-resolver/src/lib.rs @@ -34,6 +34,8 @@ pub use yanks::AllowedYanks; /// `ConflictItemRef`. i.e., We can avoid allocs on lookups. type FxHashbrownSet = hashbrown::HashSet; +type FxHashbrownMap = hashbrown::HashMap; + mod candidate_selector; mod dependency_mode; mod dependency_provider; diff --git a/crates/uv-resolver/src/pubgrub/package.rs b/crates/uv-resolver/src/pubgrub/package.rs index 084f479ed..77b45eccb 100644 --- a/crates/uv-resolver/src/pubgrub/package.rs +++ b/crates/uv-resolver/src/pubgrub/package.rs @@ -354,3 +354,9 @@ impl std::fmt::Display for PubGrubPackageInner { } } } + +impl From<&PubGrubPackage> for PubGrubPackage { + fn from(package: &PubGrubPackage) -> Self { + package.clone() + } +} diff --git a/crates/uv-resolver/src/pubgrub/priority.rs b/crates/uv-resolver/src/pubgrub/priority.rs index d08cb34ff..32da9b0a2 100644 --- a/crates/uv-resolver/src/pubgrub/priority.rs +++ b/crates/uv-resolver/src/pubgrub/priority.rs @@ -1,8 +1,8 @@ use std::cmp::Reverse; -use std::collections::hash_map::OccupiedEntry; +use hashbrown::hash_map::{EntryRef, OccupiedEntry}; use pubgrub::Range; -use rustc_hash::FxHashMap; +use rustc_hash::FxBuildHasher; use uv_normalize::PackageName; use uv_pep440::Version; @@ -10,7 +10,7 @@ use uv_pep440::Version; use crate::fork_urls::ForkUrls; use crate::pubgrub::package::PubGrubPackage; use crate::pubgrub::PubGrubPackageInner; -use crate::SentinelRange; +use crate::{FxHashbrownMap, SentinelRange}; /// A prioritization map to guide the PubGrub resolution process. /// @@ -23,8 +23,8 @@ use crate::SentinelRange; /// See: #[derive(Clone, Debug, Default)] pub(crate) struct PubGrubPriorities { - package_priority: FxHashMap, - virtual_package_tiebreaker: FxHashMap, + package_priority: FxHashbrownMap, + virtual_package_tiebreaker: FxHashbrownMap, } impl PubGrubPriorities { @@ -35,27 +35,24 @@ impl PubGrubPriorities { version: &Range, urls: &ForkUrls, ) { - if !self.virtual_package_tiebreaker.contains_key(package) { - self.virtual_package_tiebreaker.insert( - package.clone(), - PubGrubTiebreaker::from( - u32::try_from(self.virtual_package_tiebreaker.len()) - .expect("Less than 2**32 packages"), - ), - ); - } + let len = self.virtual_package_tiebreaker.len(); + self.virtual_package_tiebreaker + .entry_ref(package) + .or_insert_with(|| { + PubGrubTiebreaker::from(u32::try_from(len).expect("Less than 2**32 packages")) + }); - let next = self.package_priority.len(); // The root package and Python constraints have no explicit priority, the root package is // always first and the Python version (range) is fixed. let Some(name) = package.name_no_root() else { return; }; - match self.package_priority.entry(name.clone()) { - std::collections::hash_map::Entry::Occupied(mut entry) => { + let len = self.package_priority.len(); + match self.package_priority.entry_ref(name) { + EntryRef::Occupied(mut entry) => { // Preserve the original index. - let index = Self::get_index(&entry).unwrap_or(next); + let index = Self::get_index(&entry).unwrap_or(len); // Compute the priority. let priority = if urls.get(name).is_some() { @@ -81,16 +78,16 @@ impl PubGrubPriorities { entry.insert(priority); } } - std::collections::hash_map::Entry::Vacant(entry) => { + EntryRef::Vacant(entry) => { // Compute the priority. let priority = if urls.get(name).is_some() { - PubGrubPriority::DirectUrl(Reverse(next)) + PubGrubPriority::DirectUrl(Reverse(len)) } else if version.as_singleton().is_some() || SentinelRange::from(version).is_sentinel() { - PubGrubPriority::Singleton(Reverse(next)) + PubGrubPriority::Singleton(Reverse(len)) } else { - PubGrubPriority::Unspecified(Reverse(next)) + PubGrubPriority::Unspecified(Reverse(len)) }; // Insert the priority. @@ -99,7 +96,9 @@ impl PubGrubPriorities { } } - fn get_index(entry: &OccupiedEntry) -> Option { + fn get_index( + entry: &OccupiedEntry<'_, PackageName, PubGrubPriority, FxBuildHasher>, + ) -> Option { match entry.get() { PubGrubPriority::ConflictLate(Reverse(index)) | PubGrubPriority::Unspecified(Reverse(index)) @@ -133,7 +132,6 @@ impl PubGrubPriorities { /// Returns whether the priority was changed, i.e., it's the first time we hit this condition /// for the package. pub(crate) fn mark_conflict_early(&mut self, package: &PubGrubPackage) -> bool { - let next = self.package_priority.len(); let Some(name) = package.name_no_root() else { // Not a correctness bug if cfg!(debug_assertions) { @@ -142,20 +140,22 @@ impl PubGrubPriorities { return false; } }; - match self.package_priority.entry(name.clone()) { - std::collections::hash_map::Entry::Occupied(mut entry) => { + + let len = self.package_priority.len(); + match self.package_priority.entry_ref(name) { + EntryRef::Vacant(entry) => { + entry.insert(PubGrubPriority::ConflictEarly(Reverse(len))); + true + } + EntryRef::Occupied(mut entry) => { if matches!(entry.get(), PubGrubPriority::ConflictEarly(_)) { // Already in the right category return false; }; - let index = Self::get_index(&entry).unwrap_or(next); + let index = Self::get_index(&entry).unwrap_or(len); entry.insert(PubGrubPriority::ConflictEarly(Reverse(index))); true } - std::collections::hash_map::Entry::Vacant(entry) => { - entry.insert(PubGrubPriority::ConflictEarly(Reverse(next))); - true - } } } @@ -165,7 +165,6 @@ impl PubGrubPriorities { /// Returns whether the priority was changed, i.e., it's the first time this package was /// marked as conflicting above the threshold. pub(crate) fn mark_conflict_late(&mut self, package: &PubGrubPackage) -> bool { - let next = self.package_priority.len(); let Some(name) = package.name_no_root() else { // Not a correctness bug if cfg!(debug_assertions) { @@ -174,8 +173,14 @@ impl PubGrubPriorities { return false; } }; - match self.package_priority.entry(name.clone()) { - std::collections::hash_map::Entry::Occupied(mut entry) => { + + let len = self.package_priority.len(); + match self.package_priority.entry_ref(name) { + EntryRef::Vacant(entry) => { + entry.insert(PubGrubPriority::ConflictLate(Reverse(len))); + true + } + EntryRef::Occupied(mut entry) => { // The ConflictEarly` match avoids infinite loops. if matches!( entry.get(), @@ -184,14 +189,10 @@ impl PubGrubPriorities { // Already in the right category return false; }; - let index = Self::get_index(&entry).unwrap_or(next); + let index = Self::get_index(&entry).unwrap_or(len); entry.insert(PubGrubPriority::ConflictLate(Reverse(index))); true } - std::collections::hash_map::Entry::Vacant(entry) => { - entry.insert(PubGrubPriority::ConflictLate(Reverse(next))); - true - } } } }