diff --git a/crates/uv-resolver/src/candidate_selector.rs b/crates/uv-resolver/src/candidate_selector.rs index b674ac81b3af..1d76f12e4acc 100644 --- a/crates/uv-resolver/src/candidate_selector.rs +++ b/crates/uv-resolver/src/candidate_selector.rs @@ -1,5 +1,6 @@ use std::fmt::{Display, Formatter}; +use either::Either; use itertools::Itertools; use pubgrub::Range; use tracing::{debug, trace}; @@ -178,22 +179,51 @@ impl CandidateSelector { index: Option<&'a IndexUrl>, env: &ResolverEnvironment, ) -> Option> { - // 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::>(); + 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, diff --git a/crates/uv-resolver/src/preferences.rs b/crates/uv-resolver/src/preferences.rs index feb14ce21d98..7f418fd3a944 100644 --- a/crates/uv-resolver/src/preferences.rs +++ b/crates/uv-resolver/src/preferences.rs @@ -121,12 +121,29 @@ impl Preference { } #[derive(Debug, Clone)] -struct Entry { +pub(crate) struct Entry { marker: UniversalMarker, index: Option, 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, &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 } }