diff --git a/crates/pep440-rs/src/version.rs b/crates/pep440-rs/src/version.rs index 0172902b7..2ce410d0c 100644 --- a/crates/pep440-rs/src/version.rs +++ b/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 { diff --git a/crates/uv-resolver/src/candidate_selector.rs b/crates/uv-resolver/src/candidate_selector.rs index b4ace8ae0..4deddf414 100644 --- a/crates/uv-resolver/src/candidate_selector.rs +++ b/crates/uv-resolver/src/candidate_selector.rs @@ -324,7 +324,13 @@ impl CandidateSelector { version_maps.iter().map(VersionMap::len).sum::(), ); 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)>, package_name: &'a PackageName, range: &Range, - allow_prerelease: AllowPrerelease, + allow_prerelease: bool, ) -> Option> { - #[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 { - 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 { + let candidate = { + if version.any_prerelease() && !allow_prerelease { continue; } + 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 } } diff --git a/crates/uv-resolver/src/resolver/mod.rs b/crates/uv-resolver/src/resolver/mod.rs index 1fd384ee9..95084bfa3 100644 --- a/crates/uv-resolver/src/resolver/mod.rs +++ b/crates/uv-resolver/src/resolver/mod.rs @@ -1955,7 +1955,7 @@ impl ResolverState, build_options: &BuildOptions, ) -> Self { + let mut stable = false; let mut map = BTreeMap::new(); // Create stubs for each entry in simple metadata. The full conversion // from a `VersionFiles` to a PrioritizedDist for each version @@ -59,6 +61,7 @@ impl VersionMap { .version .deserialize(&mut SharedDeserializeMap::new()) .expect("archived version always deserializes"); + stable |= version.is_stable(); map.insert( version, LazyPrioritizedDist::OnlySimple(SimplePrioritizedDist { @@ -70,6 +73,7 @@ impl VersionMap { // If a set of flat distributions have been given, we need to add those // to our map of entries as well. for (version, prioritized_dist) in flat_index.into_iter().flatten() { + stable |= version.is_stable(); match map.entry(version) { Entry::Vacant(e) => { e.insert(LazyPrioritizedDist::OnlyFlat(prioritized_dist)); @@ -95,6 +99,7 @@ impl VersionMap { Self { inner: VersionMapInner::Lazy(VersionMapLazy { map, + stable, simple_metadata, no_binary: build_options.no_binary_package(package_name), no_build: build_options.no_build_package(package_name), @@ -125,11 +130,19 @@ impl VersionMap { version: &Version, ) -> Option<(&Version, &PrioritizedDist)> { match self.inner { - VersionMapInner::Eager(ref map) => map.get_key_value(version), + VersionMapInner::Eager(ref eager) => eager.map.get_key_value(version), VersionMapInner::Lazy(ref lazy) => lazy.get_with_version(version), } } + /// Return an iterator over the versions in this map. + pub(crate) fn versions(&self) -> impl Iterator { + match &self.inner { + VersionMapInner::Eager(eager) => either::Either::Left(eager.map.keys()), + VersionMapInner::Lazy(lazy) => either::Either::Right(lazy.map.keys()), + } + } + /// Return an iterator over the versions and distributions. /// /// Note that the value returned in this iterator is a [`VersionMapDist`], @@ -138,31 +151,60 @@ impl VersionMap { /// for each version. pub(crate) fn iter( &self, + range: &Range, ) -> impl DoubleEndedIterator + ExactSizeIterator { - match self.inner { - VersionMapInner::Eager(ref map) => { - either::Either::Left(map.iter().map(|(version, dist)| { - let version_map_dist = VersionMapDistHandle { - inner: VersionMapDistHandleInner::Eager(dist), - }; - (version, version_map_dist) - })) - } - VersionMapInner::Lazy(ref lazy) => { - either::Either::Right(lazy.map.iter().map(|(version, dist)| { - let version_map_dist = VersionMapDistHandle { - inner: VersionMapDistHandleInner::Lazy { lazy, dist }, - }; - (version, version_map_dist) - })) - } + // Performance optimization: If we only have a single version, return that version directly. + if let Some(version) = range.as_singleton() { + either::Either::Left(match self.inner { + VersionMapInner::Eager(ref eager) => { + either::Either::Left(eager.map.get_key_value(version).into_iter().map( + move |(version, dist)| { + let version_map_dist = VersionMapDistHandle { + inner: VersionMapDistHandleInner::Eager(dist), + }; + (version, version_map_dist) + }, + )) + } + VersionMapInner::Lazy(ref lazy) => { + either::Either::Right(lazy.map.get_key_value(version).into_iter().map( + move |(version, dist)| { + let version_map_dist = VersionMapDistHandle { + inner: VersionMapDistHandleInner::Lazy { lazy, dist }, + }; + (version, version_map_dist) + }, + )) + } + }) + } else { + either::Either::Right(match self.inner { + VersionMapInner::Eager(ref eager) => { + either::Either::Left(eager.map.iter().map(|(version, dist)| { + let version_map_dist = VersionMapDistHandle { + inner: VersionMapDistHandleInner::Eager(dist), + }; + (version, version_map_dist) + })) + } + VersionMapInner::Lazy(ref lazy) => { + either::Either::Right(lazy.map.iter().map(|(version, dist)| { + let version_map_dist = VersionMapDistHandle { + inner: VersionMapDistHandleInner::Lazy { lazy, dist }, + }; + (version, version_map_dist) + })) + } + }) } } /// Return the [`Hashes`] for the given version, if any. pub(crate) fn hashes(&self, version: &Version) -> Option> { match self.inner { - VersionMapInner::Eager(ref map) => map.get(version).map(|file| file.hashes().to_vec()), + VersionMapInner::Eager(ref eager) => { + eager.map.get(version).map(|file| file.hashes().to_vec()) + } VersionMapInner::Lazy(ref lazy) => lazy.get(version).map(|file| file.hashes().to_vec()), } } @@ -173,33 +215,26 @@ impl VersionMap { /// usable in the current environment. pub(crate) fn len(&self) -> usize { match self.inner { - VersionMapInner::Eager(ref map) => map.len(), + VersionMapInner::Eager(VersionMapEager { ref map, .. }) => map.len(), VersionMapInner::Lazy(VersionMapLazy { ref map, .. }) => map.len(), } } -} -impl Default for VersionMap { - /// Create an empty version map. - fn default() -> Self { - Self { - inner: VersionMapInner::Eager(BTreeMap::default()), + /// Returns `true` if the map contains at least one stable (non-pre-release) version. + pub(crate) fn stable(&self) -> bool { + match self.inner { + VersionMapInner::Eager(ref map) => map.stable, + VersionMapInner::Lazy(ref map) => map.stable, } } } impl From for VersionMap { fn from(flat_index: FlatDistributions) -> Self { + let stable = flat_index.iter().any(|(version, _)| version.is_stable()); + let map = flat_index.into(); Self { - inner: VersionMapInner::Eager(flat_index.into()), - } - } -} - -impl From> for VersionMap { - fn from(value: BTreeMap) -> Self { - Self { - inner: VersionMapInner::Eager(value), + inner: VersionMapInner::Eager(VersionMapEager { map, stable }), } } } @@ -245,7 +280,7 @@ enum VersionMapInner { /// /// This usually happens when one needs a `VersionMap` from a /// `FlatDistributions`. - Eager(BTreeMap), + Eager(VersionMapEager), /// Some distributions might be fully materialized (i.e., by initializing /// a `VersionMap` with a `FlatDistributions`), but some distributions /// might still be in their "raw" `SimpleMetadata` format. In this case, a @@ -254,6 +289,15 @@ enum VersionMapInner { Lazy(VersionMapLazy), } +/// A map from versions to distributions that are fully materialized in memory. +#[derive(Debug)] +struct VersionMapEager { + /// A map from version to distribution. + map: BTreeMap, + /// Whether the version map contains at least one stable (non-pre-release) version. + stable: bool, +} + /// A map that lazily materializes some prioritized distributions upon access. /// /// The idea here is that some packages have a lot of versions published, and @@ -266,6 +310,8 @@ enum VersionMapInner { struct VersionMapLazy { /// A map from version to possibly-initialized distribution. map: BTreeMap, + /// Whether the version map contains at least one stable (non-pre-release) version. + stable: bool, /// The raw simple metadata from which `PrioritizedDist`s should /// be constructed. simple_metadata: OwnedArchive,