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.

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<PrioritizedDistributionInner>);
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Default, Clone)]
 struct PrioritizedDistributionInner {
     /// An arbitrary source distribution for the package version.
     source: Option<DistRequiresPython>,
@@ -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)]

crates/puffin-client/src/flat_index.rs

@@ -279,6 +279,19 @@ impl FlatDistributions {
     pub fn iter(&self) -> impl Iterator<Item = (&Version, &PrioritizedDistribution)> {
         self.0.iter()
     }
+
+    pub fn remove(&mut self, version: &Version) -> Option<PrioritizedDistribution> {
+        self.0.remove(version)
+    }
+}
+
+impl IntoIterator for FlatDistributions {
+    type IntoIter = std::collections::btree_map::IntoIter<Version, PrioritizedDistribution>;
+    type Item = (Version, PrioritizedDistribution);
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.0.into_iter()
+    }
 }
 
 impl From<FlatDistributions> for BTreeMap<Version, PrioritizedDistribution> {

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<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)
     }