puffin-resolver: simplify version map construction (#1267)

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.
This commit is contained in:
Andrew Gallant 2024-02-08 15:33:33 -05:00 committed by GitHub
parent fc2ab611d5
commit 96276d9e3e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 40 additions and 41 deletions

View file

@ -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<Utc>>,
flat_index: Option<FlatDistributions>,
mut flat_index: Option<FlatDistributions>,
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<Version, PrioritizedDistribution> =
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)
}