Sort preferences by environment, then index (#10700)

## Summary

This has a few effects:

1. We only call `preferences` once, which should be more efficient.
2. We collect `preferences` into a vector when there are multiple. Less
efficient, but pretty rare?
3. We now correctly prefer preferences from the same index.
This commit is contained in:
Charlie Marsh 2025-01-17 12:27:44 -05:00 committed by GitHub
parent 1f29165796
commit e02a7bb75d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 73 additions and 27 deletions

View file

@ -1,7 +1,9 @@
use std::fmt::{Display, Formatter};
use either::Either;
use itertools::Itertools;
use pubgrub::Range;
use smallvec::SmallVec;
use tracing::{debug, trace};
use uv_configuration::IndexStrategy;
@ -11,7 +13,7 @@ use uv_normalize::PackageName;
use uv_pep440::Version;
use uv_types::InstalledPackagesProvider;
use crate::preferences::Preferences;
use crate::preferences::{Entry, Preferences};
use crate::prerelease::{AllowPrerelease, PrereleaseStrategy};
use crate::resolution_mode::ResolutionStrategy;
use crate::universal_marker::UniversalMarker;
@ -178,22 +180,53 @@ impl CandidateSelector {
index: Option<&'a IndexUrl>,
env: &ResolverEnvironment,
) -> Option<Candidate<'a>> {
// In the branches, we "sort" the preferences by marker-matching through an iterator that
// first has the matching half and then the mismatching half.
let preferences_match = preferences
.get(package_name)
.filter(|(marker, _index, _version)| env.included_by_marker(marker.pep508()));
let preferences_mismatch = preferences
.get(package_name)
.filter(|(marker, _index, _version)| !env.included_by_marker(marker.pep508()));
let preferences = preferences_match.chain(preferences_mismatch).filter_map(
|(marker, source, version)| {
// Ignore preferences that are associated with conflicting indexes.
index
.is_none_or(|index| source.is_none_or(|source| source == index))
.then_some((marker, version))
},
);
let preferences = preferences.get(package_name);
// If there are multiple preferences for the same package, we need to sort them by priority.
let preferences = match preferences {
[] => return None,
[entry] => {
// Filter out preferences that map to a conflicting index.
if index.is_some_and(|index| {
entry
.index()
.is_some_and(|entry_index| entry_index != index)
}) {
return None;
};
Either::Left(std::iter::once((entry.marker(), entry.pin().version())))
}
[..] => {
type Entries<'a> = SmallVec<[&'a Entry; 3]>;
let mut preferences = preferences.iter().collect::<Entries>();
preferences.retain(|entry| {
// Filter out preferences that map to a conflicting index.
!index.is_some_and(|index| {
entry
.index()
.is_some_and(|entry_index| entry_index != index)
})
});
preferences.sort_by_key(|entry| {
let marker = entry.marker();
// Prefer preferences that match the current environment.
let matches_env = env.included_by_marker(marker.pep508());
// Prefer preferences that match the current index.
let matches_index = index == entry.index();
std::cmp::Reverse((matches_env, matches_index))
});
Either::Right(
preferences
.into_iter()
.map(|entry| (entry.marker(), entry.pin().version())),
)
}
};
self.get_preferred_from_iter(
preferences,
package_name,

View file

@ -121,12 +121,29 @@ impl Preference {
}
#[derive(Debug, Clone)]
struct Entry {
pub(crate) struct Entry {
marker: UniversalMarker,
index: Option<IndexUrl>,
pin: Pin,
}
impl Entry {
/// Return the [`UniversalMarker`] associated with the entry.
pub(crate) fn marker(&self) -> &UniversalMarker {
&self.marker
}
/// Return the [`IndexUrl`] associated with the entry, if any.
pub(crate) fn index(&self) -> Option<&IndexUrl> {
self.index.as_ref()
}
/// Return the pinned data associated with the entry.
pub(crate) fn pin(&self) -> &Pin {
&self.pin
}
}
/// A set of pinned packages that should be preserved during resolution, if possible.
///
/// The marker is the marker of the fork that resolved to the pin, if any.
@ -232,15 +249,11 @@ impl Preferences {
}
/// Return the pinned version for a package, if any.
pub(crate) fn get(
&self,
package_name: &PackageName,
) -> impl Iterator<Item = (&UniversalMarker, Option<&IndexUrl>, &Version)> {
pub(crate) fn get(&self, package_name: &PackageName) -> &[Entry] {
self.0
.get(package_name)
.into_iter()
.flatten()
.map(|entry| (&entry.marker, entry.index.as_ref(), entry.pin.version()))
.map(Vec::as_slice)
.unwrap_or_default()
}
/// Return the hashes for a package, if the version matches that of the pin.
@ -273,12 +286,12 @@ pub(crate) struct Pin {
impl Pin {
/// Return the version of the pinned package.
fn version(&self) -> &Version {
pub(crate) fn version(&self) -> &Version {
&self.version
}
/// Return the hashes of the pinned package.
fn hashes(&self) -> &[HashDigest] {
pub(crate) fn hashes(&self) -> &[HashDigest] {
&self.hashes
}
}