Avoid iteration for singleton selections (#7195)

## Summary

If we have a singleton `Range`, we don't need to iterate over the map of
available ranges; instead, we can just get the singleton directly.

Closes #6131.

crates/pep440-rs/src/version.rs

@@ -329,6 +329,12 @@ impl Version {
         self.is_pre() || self.is_dev()
     }
 
+    /// Whether this is a stable version (i.e., _not_ an alpha/beta/rc or dev version)
+    #[inline]
+    pub fn is_stable(&self) -> bool {
+        !self.is_pre() && !self.is_dev()
+    }
+
     /// Whether this is an alpha/beta/rc version
     #[inline]
     pub fn is_pre(&self) -> bool {

crates/uv-resolver/src/candidate_selector.rs

@@ -324,7 +324,13 @@ impl CandidateSelector {
             version_maps.iter().map(VersionMap::len).sum::<usize>(),
         );
         let highest = self.use_highest_version(package_name);
-        let allow_prerelease = self.prerelease_strategy.allows(package_name, markers);
+
+        let allow_prerelease = match self.prerelease_strategy.allows(package_name, markers) {
+            AllowPrerelease::Yes => true,
+            AllowPrerelease::No => false,
+            // Allow pre-releases if there are no stable versions available.
+            AllowPrerelease::IfNecessary => !version_maps.iter().any(VersionMap::stable),
+        };
 
         if self.index_strategy == IndexStrategy::UnsafeBestMatch {
             if highest {
@@ -333,7 +339,10 @@ impl CandidateSelector {
                         .iter()
                         .enumerate()
                         .map(|(map_index, version_map)| {
-                            version_map.iter().rev().map(move |item| (map_index, item))
+                            version_map
+                                .iter(range)
+                                .rev()
+                                .map(move |item| (map_index, item))
                         })
                         .kmerge_by(
                             |(index1, (version1, _)), (index2, (version2, _))| match version1
@@ -355,7 +364,7 @@ impl CandidateSelector {
                         .iter()
                         .enumerate()
                         .map(|(map_index, version_map)| {
-                            version_map.iter().map(move |item| (map_index, item))
+                            version_map.iter(range).map(move |item| (map_index, item))
                         })
                         .kmerge_by(
                             |(index1, (version1, _)), (index2, (version2, _))| match version1
@@ -376,7 +385,7 @@ impl CandidateSelector {
             if highest {
                 version_maps.iter().find_map(|version_map| {
                     Self::select_candidate(
-                        version_map.iter().rev(),
+                        version_map.iter(range).rev(),
                         package_name,
                         range,
                         allow_prerelease,
@@ -385,7 +394,7 @@ impl CandidateSelector {
             } else {
                 version_maps.iter().find_map(|version_map| {
                     Self::select_candidate(
-                        version_map.iter(),
+                        version_map.iter(range),
                         package_name,
                         range,
                         allow_prerelease,
@@ -413,83 +422,24 @@ impl CandidateSelector {
         versions: impl Iterator<Item = (&'a Version, VersionMapDistHandle<'a>)>,
         package_name: &'a PackageName,
         range: &Range<Version>,
-        allow_prerelease: AllowPrerelease,
+        allow_prerelease: bool,
     ) -> Option<Candidate<'a>> {
-        #[derive(Debug)]
-        enum PrereleaseCandidate<'a> {
-            NotNecessary,
-            IfNecessary(&'a Version, &'a PrioritizedDist),
-        }
-
-        let mut prerelease = None;
         let mut steps = 0usize;
         for (version, maybe_dist) in versions {
             steps += 1;
-            let candidate = if version.any_prerelease() {
-                if range.contains(version) {
-                    match allow_prerelease {
-                        AllowPrerelease::Yes => {
-                            let Some(dist) = maybe_dist.prioritized_dist() else {
-                                continue;
-                            };
-                            tracing::trace!(
-                                "found candidate for package {:?} with range {:?} \
-                                 after {} steps: {:?} version",
-                                package_name,
-                                range,
-                                steps,
-                                version,
-                            );
-                            // If pre-releases are allowed, treat them equivalently
-                            // to stable distributions.
-                            Candidate::new(
-                                package_name,
-                                version,
-                                dist,
-                                VersionChoiceKind::Compatible,
-                            )
-                        }
-                        AllowPrerelease::IfNecessary => {
-                            let Some(dist) = maybe_dist.prioritized_dist() else {
-                                continue;
-                            };
-                            // If pre-releases are allowed as a fallback, store the
-                            // first-matching prerelease.
-                            if prerelease.is_none() {
-                                prerelease = Some(PrereleaseCandidate::IfNecessary(version, dist));
-                            }
-                            continue;
-                        }
-                        AllowPrerelease::No => {
-                            continue;
-                        }
-                    }
-                } else {
+
+            let candidate = {
+                if version.any_prerelease() && !allow_prerelease {
                     continue;
                 }
-            } else {
-                // If we have at least one stable release, we shouldn't allow the "if-necessary"
-                // pre-release strategy, regardless of whether that stable release satisfies the
-                // current range.
-                prerelease = Some(PrereleaseCandidate::NotNecessary);
-
-                // Return the first-matching stable distribution.
-                if range.contains(version) {
-                    let Some(dist) = maybe_dist.prioritized_dist() else {
-                        continue;
-                    };
-                    tracing::trace!(
-                        "found candidate for package {:?} with range {:?} \
-                         after {} steps: {:?} version",
-                        package_name,
-                        range,
-                        steps,
-                        version,
-                    );
-                    Candidate::new(package_name, version, dist, VersionChoiceKind::Compatible)
-                } else {
+                if !range.contains(version) {
                     continue;
-                }
+                };
+                let Some(dist) = maybe_dist.prioritized_dist() else {
+                    continue;
+                };
+                trace!("found candidate for package {package_name:?} with range {range:?} after {steps} steps: {version:?} version");
+                Candidate::new(package_name, version, dist, VersionChoiceKind::Compatible)
             };
 
             // If candidate is not compatible due to exclude newer, continue searching.
@@ -510,20 +460,8 @@ impl CandidateSelector {
 
             return Some(candidate);
         }
-        trace!(
-            "Exhausted all candidates for package {package_name} with range {range} \
-             after {steps} steps",
-        );
-        match prerelease {
-            None => None,
-            Some(PrereleaseCandidate::NotNecessary) => None,
-            Some(PrereleaseCandidate::IfNecessary(version, dist)) => Some(Candidate::new(
-                package_name,
-                version,
-                dist,
-                VersionChoiceKind::Compatible,
-            )),
-        }
+        trace!("Exhausted all candidates for package {package_name} with range {range} after {steps} steps");
+        None
     }
 }
 

crates/uv-resolver/src/resolver/mod.rs

@@ -1955,7 +1955,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
                         available_versions
                             .entry(name.clone())
                             .or_insert_with(BTreeSet::new)
-                            .extend(version_map.iter().map(|(version, _)| version.clone()));
+                            .extend(version_map.versions().cloned());
                     }
                 }
             }