Skip to content

Commit

Permalink
Sort preferences by environment, then index
Browse files Browse the repository at this point in the history
  • Loading branch information
charliermarsh committed Jan 17, 2025
1 parent 3344609 commit 22873a9
Show file tree
Hide file tree
Showing 2 changed files with 69 additions and 26 deletions.
62 changes: 46 additions & 16 deletions crates/uv-resolver/src/candidate_selector.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
use std::fmt::{Display, Formatter};

use either::Either;
use itertools::Itertools;
use pubgrub::Range;
use tracing::{debug, trace};
Expand Down Expand Up @@ -178,22 +179,51 @@ 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())))
}
[..] => {
let mut preferences = preferences.iter().collect::<Vec<_>>();
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,
Expand Down
33 changes: 23 additions & 10 deletions crates/uv-resolver/src/preferences.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
}
}
Expand Down

0 comments on commit 22873a9

Please sign in to comment.