Add an extra struct around the package-to-flat index map (#923)

## Summary

`FlatIndex` is now the thing that's keyed on `PackageName`, while
`FlatDistributions` is what used to be called `FlatIndex` (a map from
version to `PrioritizedDistribution`, for a single package). I find this
a bit clearer, since we can also remove the `from_files` that doesn't
return `Self`, which I had trouble following.

crates/puffin-resolver/src/finder.rs

@@ -10,7 +10,7 @@ use distribution_filename::DistFilename;
 use distribution_types::{Dist, IndexUrl, Resolution};
 use pep508_rs::{Requirement, VersionOrUrl};
 use platform_tags::Tags;
-use puffin_client::{FlatIndex, RegistryClient, SimpleMetadata};
+use puffin_client::{FlatDistributions, FlatIndex, RegistryClient, SimpleMetadata};
 use puffin_interpreter::Interpreter;
 use puffin_normalize::PackageName;
 
@@ -21,7 +21,7 @@ pub struct DistFinder<'a> {
     client: &'a RegistryClient,
     reporter: Option<Box<dyn Reporter>>,
     interpreter: &'a Interpreter,
-    flat_index: &'a FxHashMap<PackageName, FlatIndex>,
+    flat_index: &'a FlatIndex,
 }
 
 impl<'a> DistFinder<'a> {
@@ -30,7 +30,7 @@ impl<'a> DistFinder<'a> {
         tags: &'a Tags,
         client: &'a RegistryClient,
         interpreter: &'a Interpreter,
-        flat_index: &'a FxHashMap<PackageName, FlatIndex>,
+        flat_index: &'a FlatIndex,
     ) -> Self {
         Self {
             tags,
@@ -56,7 +56,7 @@ impl<'a> DistFinder<'a> {
     async fn resolve_requirement(
         &self,
         requirement: &Requirement,
-        flat_index: Option<&FlatIndex>,
+        flat_index: Option<&FlatDistributions>,
     ) -> Result<(PackageName, Dist), ResolveError> {
         match requirement.version_or_url.as_ref() {
             None | Some(VersionOrUrl::VersionSpecifier(_)) => {
@@ -118,7 +118,7 @@ impl<'a> DistFinder<'a> {
         requirement: &Requirement,
         metadata: SimpleMetadata,
         index: &IndexUrl,
-        flat_index: Option<&FlatIndex>,
+        flat_index: Option<&FlatDistributions>,
     ) -> Option<Dist> {
         // Prioritize the flat index by initializing the "best" matches with its entries.
         let matching_override = if let Some(flat_index) = flat_index {

crates/puffin-resolver/src/resolver/mod.rs

@@ -24,7 +24,7 @@ use distribution_types::{
 use pep440_rs::{Version, VersionSpecifiers, MIN_VERSION};
 use pep508_rs::{MarkerEnvironment, Requirement};
 use platform_tags::Tags;
-use puffin_client::RegistryClient;
+use puffin_client::{FlatIndex, RegistryClient};
 use puffin_distribution::DistributionDatabase;
 use puffin_interpreter::Interpreter;
 use puffin_normalize::PackageName;
@@ -89,6 +89,7 @@ impl<'a, Context: BuildContext + Send + Sync> Resolver<'a, DefaultResolverProvid
         let provider = DefaultResolverProvider::new(
             client,
             DistributionDatabase::new(build_context.cache(), tags, client, build_context),
+            FlatIndex::from_files(client.flat_index().await?, tags),
             tags,
             PythonRequirement::new(interpreter, markers),
             options.exclude_newer,
@@ -97,8 +98,7 @@ impl<'a, Context: BuildContext + Send + Sync> Resolver<'a, DefaultResolverProvid
                 .iter()
                 .chain(manifest.constraints.iter())
                 .collect(),
-        )
-        .await?;
+        );
         Ok(Self::new_custom_io(
             manifest,
             options,

crates/puffin-resolver/src/resolver/provider.rs

@@ -3,7 +3,6 @@ use std::future::Future;
 use anyhow::Result;
 use chrono::{DateTime, Utc};
 use futures::FutureExt;
-use rustc_hash::FxHashMap;
 use url::Url;
 
 use distribution_types::Dist;
@@ -46,10 +45,12 @@ pub trait ResolverProvider: Send + Sync {
 /// The main IO backend for the resolver, which does cached requests network requests using the
 /// [`RegistryClient`] and [`DistributionDatabase`].
 pub struct DefaultResolverProvider<'a, Context: BuildContext + Send + Sync> {
+    /// The [`RegistryClient`] used to query the index.
     client: &'a RegistryClient,
-    /// These are the entries from `--find-links` that act as overrides for index responses.
-    flat_index: FxHashMap<PackageName, FlatIndex>,
+    /// The [`DistributionDatabase`] used to build source distributions.
     fetcher: DistributionDatabase<'a, Context>,
+    /// These are the entries from `--find-links` that act as overrides for index responses.
+    flat_index: FlatIndex,
     tags: &'a Tags,
     python_requirement: PythonRequirement<'a>,
     exclude_newer: Option<DateTime<Utc>>,
@@ -58,26 +59,24 @@ pub struct DefaultResolverProvider<'a, Context: BuildContext + Send + Sync> {
 
 impl<'a, Context: BuildContext + Send + Sync> DefaultResolverProvider<'a, Context> {
     /// Reads the flat index entries and builds the provider.
-    pub async fn new(
+    pub fn new(
         client: &'a RegistryClient,
         fetcher: DistributionDatabase<'a, Context>,
+        flat_index: FlatIndex,
         tags: &'a Tags,
         python_requirement: PythonRequirement<'a>,
         exclude_newer: Option<DateTime<Utc>>,
         allowed_yanks: AllowedYanks,
-    ) -> Result<Self, puffin_client::Error> {
-        let flat_index_dists = client.flat_index().await?;
-        let flat_index = FlatIndex::from_files(flat_index_dists, tags);
-
-        Ok(Self {
+    ) -> Self {
+        Self {
             client,
-            flat_index,
             fetcher,
+            flat_index,
             tags,
             python_requirement,
             exclude_newer,
             allowed_yanks,
-        })
+        }
     }
 }
 
@@ -88,7 +87,6 @@ impl<'a, Context: BuildContext + Send + Sync> ResolverProvider
         &'io self,
         package_name: &'io PackageName,
     ) -> impl Future<Output = VersionMapResponse> + Send + 'io {
-        let flat_index_override = self.flat_index.get(package_name).cloned();
         self.client
             .simple(package_name)
             .map(move |result| match result {
@@ -100,10 +98,10 @@ impl<'a, Context: BuildContext + Send + Sync> ResolverProvider
                     &self.python_requirement,
                     &self.allowed_yanks,
                     self.exclude_newer.as_ref(),
-                    flat_index_override,
+                    self.flat_index.get(package_name).cloned(),
                 )),
                 Err(err @ puffin_client::Error::PackageNotFound(_)) => {
-                    if let Some(flat_index) = flat_index_override {
+                    if let Some(flat_index) = self.flat_index.get(package_name).cloned() {
                         Ok(VersionMap::from(flat_index))
                     } else {
                         Err(err)

crates/puffin-resolver/src/version_map.rs

@@ -8,7 +8,7 @@ use distribution_filename::DistFilename;
 use distribution_types::{Dist, IndexUrl, PrioritizedDistribution, ResolvableDist};
 use pep440_rs::Version;
 use platform_tags::Tags;
-use puffin_client::{FlatIndex, SimpleMetadata};
+use puffin_client::{FlatDistributions, SimpleMetadata};
 use puffin_normalize::PackageName;
 use puffin_warnings::warn_user_once;
 use pypi_types::{Hashes, Yanked};
@@ -32,11 +32,11 @@ impl VersionMap {
         python_requirement: &PythonRequirement,
         allowed_yanks: &AllowedYanks,
         exclude_newer: Option<&DateTime<Utc>>,
-        flat_index: Option<FlatIndex>,
+        flat_index: Option<FlatDistributions>,
     ) -> Self {
         // 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(|overrides| overrides.0).unwrap_or_default();
+            flat_index.map(Into::into).unwrap_or_default();
 
         // Collect compatible distributions.
         for (version, files) in metadata {
@@ -155,8 +155,8 @@ impl VersionMap {
     }
 }
 
-impl From<FlatIndex> for VersionMap {
-    fn from(flat_index: FlatIndex) -> Self {
-        Self(flat_index.0)
+impl From<FlatDistributions> for VersionMap {
+    fn from(flat_index: FlatDistributions) -> Self {
+        Self(flat_index.into())
     }
 }