mirrors/uv
1
0
Fork 0
mirror of https://github.com/astral-sh/uv.git synced 2025-08-04 10:58:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

Use Hashbrown's raw entry API to reduce hashes and clone in priority (#10881)

Browse source 
## Summary

I'm open to not merging this -- I was kind of just interested in what
the API looked like. But the idea is: we can avoid hashing values twice
and unnecessarily cloning within the priority map by using the raw entry
API.
This commit is contained in:
Charlie Marsh 2025-01-23 09:34:37 -05:00 • committed by GitHub
parent db4ab9dc8a
commit b2dac9979d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 48 additions and 39 deletions
Download patch file Download diff file Expand all files Collapse all files

2
crates/uv-resolver/src/lib.rs
View file

@ -34,6 +34,8 @@ pub use yanks::AllowedYanks;
/// `ConflictItemRef`. i.e., We can avoid allocs on lookups.
type FxHashbrownSet<T> = hashbrown::HashSet<T, rustc_hash::FxBuildHasher>;
type FxHashbrownMap<K, V> = hashbrown::HashMap<K, V, rustc_hash::FxBuildHasher>;
mod candidate_selector;
mod dependency_mode;
mod dependency_provider;

6
crates/uv-resolver/src/pubgrub/package.rs
View file

@ -354,3 +354,9 @@ impl std::fmt::Display for PubGrubPackageInner {
}
}
}
impl From<&PubGrubPackage> for PubGrubPackage {
fn from(package: &PubGrubPackage) -> Self {
package.clone()
}
}

79
crates/uv-resolver/src/pubgrub/priority.rs
View file

@ -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: <https://github.com/pypa/pip/blob/ef78c129b1a966dbbbdb8ebfffc43723e89110d1/src/pip/_internal/resolution/resolvelib/provider.py#L120>
#[derive(Clone, Debug, Default)]
pub(crate) struct PubGrubPriorities {
package_priority: FxHashMap<PackageName, PubGrubPriority>,
virtual_package_tiebreaker: FxHashMap<PubGrubPackage, PubGrubTiebreaker>,
package_priority: FxHashbrownMap<PackageName, PubGrubPriority>,
virtual_package_tiebreaker: FxHashbrownMap<PubGrubPackage, PubGrubTiebreaker>,
}
impl PubGrubPriorities {
@ -35,27 +35,24 @@ impl PubGrubPriorities {
version: &Range<Version>,
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<PackageName, PubGrubPriority>) -> Option<usize> {
fn get_index(
entry: &OccupiedEntry<'_, PackageName, PubGrubPriority, FxBuildHasher>,
) -> Option<usize> {
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
}
}
}
}