Use version in RegistryIndex (#543)

When building up the `RegistryIndex`, index by both package name and
version to fix #537.

crates/puffin-installer/src/plan.rs

@@ -87,10 +87,22 @@ impl InstallPlan {
 
             // Identify any locally-available distributions that satisfy the requirement.
             match requirement.version_or_url.as_ref() {
-                None | Some(VersionOrUrl::VersionSpecifier(_)) => {
-                    if let Some(distribution) = registry_index
-                        .get(&requirement.name)
-                        .filter(|dist| requirement.is_satisfied_by(&dist.filename.version))
+                None => {
+                    if let Some((_version, distribution)) =
+                        registry_index.by_name(&requirement.name).next()
+                    {
+                        debug!("Requirement already cached: {distribution}");
+                        local.push(CachedDist::Registry(distribution.clone()));
+                        continue;
+                    }
+                }
+                Some(VersionOrUrl::VersionSpecifier(specifier)) => {
+                    if let Some((_version, distribution)) = registry_index
+                        .by_name(&requirement.name)
+                        .find(|(version, dist)| {
+                            specifier.contains(version)
+                                && requirement.is_satisfied_by(&dist.filename.version)
+                        })
                     {
                         debug!("Requirement already cached: {distribution}");
                         local.push(CachedDist::Registry(distribution.clone()));

crates/puffin-installer/src/registry_index.rs

@@ -1,9 +1,10 @@
-use std::collections::HashMap;
+use std::collections::{BTreeMap, HashMap};
 
 use fs_err as fs;
 use tracing::warn;
 
 use distribution_types::{CachedRegistryDist, Metadata};
+use pep440_rs::Version;
 use platform_tags::Tags;
 use puffin_cache::{Cache, CacheBucket, WheelCache};
 use puffin_normalize::PackageName;
@@ -11,12 +12,12 @@ use pypi_types::IndexUrls;
 
 /// A local index of distributions that originate from a registry, like `PyPI`.
 #[derive(Debug, Default)]
-pub struct RegistryIndex(HashMap<PackageName, CachedRegistryDist>);
+pub struct RegistryIndex(HashMap<PackageName, BTreeMap<Version, CachedRegistryDist>>);
 
 impl RegistryIndex {
     /// Build an index of cached distributions from a directory.
     pub fn try_from_directory(cache: &Cache, tags: &Tags, index_urls: &IndexUrls) -> Self {
-        let mut index: HashMap<PackageName, CachedRegistryDist> = HashMap::new();
+        let mut index: HashMap<PackageName, BTreeMap<Version, CachedRegistryDist>> = HashMap::new();
 
         for index_url in index_urls {
             let wheel_dir = cache
@@ -28,34 +29,36 @@ impl RegistryIndex {
             };
 
             for entry in dir {
-                let (path, file_type) =
-                    match entry.and_then(|entry| Ok((entry.path(), entry.file_type()?))) {
-                        Ok((path, file_type)) => (path, file_type),
-                        Err(err) => {
-                            warn!(
-                                "Failed to read entry of cache at {}: {}",
-                                cache.root().display(),
-                                err
-                            );
-                            continue;
-                        }
-                    };
-                if !file_type.is_dir() {
-                    continue;
-                }
+                let path = match entry.map(|entry| entry.path()) {
+                    Ok(path) => path,
+                    Err(err) => {
+                        warn!(
+                            "Failed to read entry of cache at {}: {}",
+                            cache.root().display(),
+                            err
+                        );
+                        continue;
+                    }
+                };
 
                 match CachedRegistryDist::try_from_path(&path) {
                     Ok(None) => {}
                     Ok(Some(dist_info)) => {
                         // Pick the wheel with the highest priority
                         let compatibility = dist_info.filename.compatibility(tags);
-                        if let Some(existing) = index.get_mut(dist_info.name()) {
+                        if let Some(existing) = index
+                            .get_mut(dist_info.name())
+                            .and_then(|package| package.get_mut(&dist_info.filename.version))
+                        {
                             // Override if we have better compatibility
                             if compatibility > existing.filename.compatibility(tags) {
                                 *existing = dist_info;
                             }
                         } else if compatibility.is_some() {
-                            index.insert(dist_info.name().clone(), dist_info);
+                            index
+                                .entry(dist_info.name().clone())
+                                .or_default()
+                                .insert(dist_info.filename.version.clone(), dist_info);
                         }
                     }
                     Err(err) => {
@@ -76,7 +79,13 @@ impl RegistryIndex {
     }
 
     /// Returns a distribution from the index, if it exists.
-    pub fn get(&self, name: &PackageName) -> Option<&CachedRegistryDist> {
-        self.0.get(name)
+    pub fn by_name(
+        &self,
+        name: &PackageName,
+    ) -> impl Iterator<Item = (&Version, &CachedRegistryDist)> {
+        // Using static to extend the lifetime
+        static DEFAULT_MAP: BTreeMap<Version, CachedRegistryDist> = BTreeMap::new();
+        // We should only query this
+        self.0.get(name).unwrap_or(&DEFAULT_MAP).iter().rev()
     }
 }