From 96276d9e3e2fba179c2d5ff57e7617b50a4d5dfb Mon Sep 17 00:00:00 2001 From: Andrew Gallant Date: Thu, 8 Feb 2024 15:33:33 -0500 Subject: [PATCH] puffin-resolver: simplify version map construction (#1267) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit In the process of making VersionMap construction lazy, I realized this refactoring would be useful to me. It also simplifies a fair bit of case analysis and does fewer BTreeMap lookups during construction. With that said, this doesn't seem to matter for perf: ``` $ hyperfine -w10 --runs 50 \ "puffin-main pip compile --cache-dir ~/astral/tmp/cache-main ~/astral/tmp/reqs/home-assistant-reduced.in -o /dev/null" \ "puffin-test pip compile --cache-dir ~/astral/tmp/cache-test ~/astral/tmp/reqs/home-assistant-reduced.in -o /dev/null" Benchmark 1: puffin-main pip compile --cache-dir ~/astral/tmp/cache-main ~/astral/tmp/reqs/home-assistant-reduced.in -o /dev/null Time (mean ± σ): 146.8 ms ± 4.1 ms [User: 350.1 ms, System: 314.2 ms] Range (min … max): 140.7 ms … 158.0 ms 50 runs Benchmark 2: puffin-test pip compile --cache-dir ~/astral/tmp/cache-test ~/astral/tmp/reqs/home-assistant-reduced.in -o /dev/null Time (mean ± σ): 146.8 ms ± 4.5 ms [User: 359.8 ms, System: 308.3 ms] Range (min … max): 138.2 ms … 160.1 ms 50 runs Summary puffin-main pip compile --cache-dir ~/astral/tmp/cache-main ~/astral/tmp/reqs/home-assistant-reduced.in -o /dev/null ran 1.00 ± 0.04 times faster than puffin-test pip compile --cache-dir ~/astral/tmp/cache-test ~/astral/tmp/reqs/home-assistant-reduced.in -o /dev/null ``` But the simplification is still nice, and will decrease the delta between what we have now and a lazy version map. --- .../src/prioritized_distribution.rs | 12 +++- crates/puffin-client/src/flat_index.rs | 13 +++++ crates/puffin-resolver/src/version_map.rs | 56 ++++++------------- 3 files changed, 40 insertions(+), 41 deletions(-) diff --git a/crates/distribution-types/src/prioritized_distribution.rs b/crates/distribution-types/src/prioritized_distribution.rs index 164bf8488..eb8c0a8a1 100644 --- a/crates/distribution-types/src/prioritized_distribution.rs +++ b/crates/distribution-types/src/prioritized_distribution.rs @@ -13,10 +13,10 @@ pub struct DistRequiresPython { } // Boxed because `Dist` is large. -#[derive(Debug, Clone)] +#[derive(Debug, Default, Clone)] pub struct PrioritizedDistribution(Box); -#[derive(Debug, Clone)] +#[derive(Debug, Default, Clone)] struct PrioritizedDistributionInner { /// An arbitrary source distribution for the package version. source: Option, @@ -178,6 +178,14 @@ impl PrioritizedDistribution { pub fn hashes(&self) -> &[Hashes] { &self.0.hashes } + + /// Returns true if and only if this distribution does not contain any + /// source distributions or wheels. + pub fn is_empty(&self) -> bool { + self.0.source.is_none() + && self.0.compatible_wheel.is_none() + && self.0.incompatible_wheel.is_none() + } } #[derive(Debug, Clone)] diff --git a/crates/puffin-client/src/flat_index.rs b/crates/puffin-client/src/flat_index.rs index 816fbd2bb..3d99f0373 100644 --- a/crates/puffin-client/src/flat_index.rs +++ b/crates/puffin-client/src/flat_index.rs @@ -279,6 +279,19 @@ impl FlatDistributions { pub fn iter(&self) -> impl Iterator { self.0.iter() } + + pub fn remove(&mut self, version: &Version) -> Option { + self.0.remove(version) + } +} + +impl IntoIterator for FlatDistributions { + type IntoIter = std::collections::btree_map::IntoIter; + type Item = (Version, PrioritizedDistribution); + + fn into_iter(self) -> Self::IntoIter { + self.0.into_iter() + } } impl From for BTreeMap { diff --git a/crates/puffin-resolver/src/version_map.rs b/crates/puffin-resolver/src/version_map.rs index 8d32f0048..8c5572719 100644 --- a/crates/puffin-resolver/src/version_map.rs +++ b/crates/puffin-resolver/src/version_map.rs @@ -1,4 +1,3 @@ -use std::collections::btree_map::Entry; use std::collections::BTreeMap; use chrono::{DateTime, Utc}; @@ -33,7 +32,7 @@ impl VersionMap { python_requirement: &PythonRequirement, allowed_yanks: &AllowedYanks, exclude_newer: Option<&DateTime>, - flat_index: Option, + mut flat_index: Option, no_binary: &NoBinary, ) -> Self { // NOTE: We should experiment with refactoring the code @@ -45,9 +44,7 @@ impl VersionMap { // up-front. let metadata = OwnedArchive::deserialize(&raw_metadata); - // If we have packages of the same name from find links, gives them priority, otherwise start empty - let mut version_map: BTreeMap = - flat_index.map(Into::into).unwrap_or_default(); + let mut version_map = BTreeMap::new(); // Check if binaries are allowed for this package let no_binary = match no_binary { @@ -58,6 +55,12 @@ impl VersionMap { // Collect compatible distributions. for SimpleMetadatum { version, files } in metadata { + // If we have packages of the same name from find links, give them + // priority, otherwise start with an empty priority dist. + let mut priority_dist = flat_index + .as_mut() + .and_then(|flat_index| flat_index.remove(&version)) + .unwrap_or_default(); for (filename, file) in files.all() { // Support resolving as if it were an earlier timestamp, at least as long files have // upload time information. @@ -94,7 +97,7 @@ impl VersionMap { // If pre-built binaries are disabled, skip this wheel if no_binary { continue; - }; + } // To be compatible, the wheel must both have compatible tags _and_ have a // compatible Python requirement. @@ -110,24 +113,7 @@ impl VersionMap { file, index.clone(), ); - match version_map.entry(version.clone()) { - Entry::Occupied(mut entry) => { - entry.get_mut().insert_built( - dist, - requires_python, - Some(hash), - priority, - ); - } - Entry::Vacant(entry) => { - entry.insert(PrioritizedDistribution::from_built( - dist, - requires_python, - Some(hash), - priority, - )); - } - } + priority_dist.insert_built(dist, requires_python, Some(hash), priority); } DistFilename::SourceDistFilename(filename) => { let dist = Dist::from_registry( @@ -135,25 +121,17 @@ impl VersionMap { file, index.clone(), ); - match version_map.entry(version.clone()) { - Entry::Occupied(mut entry) => { - entry - .get_mut() - .insert_source(dist, requires_python, Some(hash)); - } - Entry::Vacant(entry) => { - entry.insert(PrioritizedDistribution::from_source( - dist, - requires_python, - Some(hash), - )); - } - } + priority_dist.insert_source(dist, requires_python, Some(hash)); } } } + version_map.insert(version, priority_dist); + } + // Add any left over packages from the version map that we didn't visit + // above via `SimpleMetadata`. + if let Some(flat_index) = flat_index { + version_map.extend(flat_index.into_iter()); } - Self(version_map) }