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.

crates/uv-resolver/src/candidate_selector.rs

@@ -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
+        let preferences = preferences.get(package_name);
-        // first has the matching half and then the mismatching half.
+
-        let preferences_match = preferences
+        // If there are multiple preferences for the same package, we need to sort them by priority.
-            .get(package_name)
+        let preferences = match preferences {
-            .filter(|(marker, _index, _version)| env.included_by_marker(marker.pep508()));
+            [] => return None,
-        let preferences_mismatch = preferences
+            [entry] => {
-            .get(package_name)
+                // Filter out preferences that map to a conflicting index.
-            .filter(|(marker, _index, _version)| !env.included_by_marker(marker.pep508()));
+                if index.is_some_and(|index| {
-        let preferences = preferences_match.chain(preferences_mismatch).filter_map(
+                    entry
-            |(marker, source, version)| {
+                        .index()
-                // Ignore preferences that are associated with conflicting indexes.
+                        .is_some_and(|entry_index| entry_index != index)
-                index
+                }) {
-                    .is_none_or(|index| source.is_none_or(|source| source == index))
+                    return None;
-                    .then_some((marker, version))
+                };
-            },
+                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,

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<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(
+    pub(crate) fn get(&self, package_name: &PackageName) -> &[Entry] {
-        &self,
-        package_name: &PackageName,
-    ) -> impl Iterator<Item = (&UniversalMarker, Option<&IndexUrl>, &Version)> {
         self.0
             .get(package_name)
-            .into_iter()
+            .map(Vec::as_slice)
-            .flatten()
+            .unwrap_or_default()
-            .map(|entry| (&entry.marker, entry.index.as_ref(), entry.pin.version()))
     }
 
     /// 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
     }
 }