pioduino/uv
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 1 Wiki Activity

Avoid iteration for singleton selections (#7195)

Browse Source 
## 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.
This commit is contained in:
Charlie Marsh
2024-09-09 09:43:57 -04:00
committed by GitHub
parent dafa3596c5
commit dcbaa1f486
4 changed files with 118 additions and 128 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/pep440-rs/src/version.rs
+6
View File
@@ -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
+27 -89
View File
@@ -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 {
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
}
}
crates/uv-resolver/src/resolver/mod.rs
+1 -1
View File
@@ -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());
}
}
}
crates/uv-resolver/src/version_map.rs
+84 -38
View File
@@ -1,7 +1,7 @@
use pubgrub::Range;
use rkyv::{de::deserializers::SharedDeserializeMap, Deserialize};
use std::collections::btree_map::{BTreeMap, Entry};
use std::sync::OnceLock;
use rkyv::{de::deserializers::SharedDeserializeMap, Deserialize};
use tracing::instrument;
use distribution_filename::{DistFilename, WheelFilename};
@@ -24,6 +24,7 @@ use crate::{yanks::AllowedYanks, ExcludeNewer, RequiresPython};
/// A map from versions to distributions.
#[derive(Debug)]
pub struct VersionMap {
/// The inner representation of the version map.
inner: VersionMapInner,
}
@@ -50,6 +51,7 @@ impl VersionMap {
flat_index: Option<FlatDistributions>,
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<Item = &Version> {
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<Version>,
) -> impl DoubleEndedIterator<Item = (&Version, VersionMapDistHandle)> + 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<Vec<HashDigest>> {
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<FlatDistributions> 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<BTreeMap<Version, PrioritizedDist>> for VersionMap {
fn from(value: BTreeMap<Version, PrioritizedDist>) -> 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<Version, PrioritizedDist>),
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<Version, PrioritizedDist>,
/// 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<Version, LazyPrioritizedDist>,
/// 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<SimpleMetadata>,