Only store compatible wheels in the resolver (#109)

Rather than constantly iterating over all files and testing their
compatibility with the current platform, just store wheels we can
actually consider in the solver cache.

crates/platform-tags/src/lib.rs

@@ -2,7 +2,7 @@ use platform_host::{Arch, Os, Platform, PlatformError};
 
 /// A set of compatible tags for a given Python version and platform, in
 /// (`python_tag`, `abi_tag`, `platform_tag`) format.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct Tags(Vec<(String, String, String)>);
 
 impl Tags {

crates/puffin-resolver/src/error.rs

@@ -1,6 +1,7 @@
 use thiserror::Error;
 
 use pep508_rs::Requirement;
+use puffin_package::package_name::PackageName;
 
 use crate::pubgrub::package::PubGrubPackage;
 use crate::pubgrub::version::PubGrubVersion;
@@ -13,6 +14,9 @@ pub enum ResolveError {
     #[error("The request stream terminated unexpectedly")]
     StreamTermination,
 
+    #[error("No platform-compatible distributions found for: {0}")]
+    NoCompatibleDistributions(PackageName),
+
     #[error(transparent)]
     Client(#[from] puffin_client::PypiClientError),
 

crates/puffin-resolver/src/resolver.rs

@@ -63,23 +63,20 @@ impl<'a> Resolver<'a> {
         // metadata (e.g., given `flask==1.0.0`, fetch the metadata for that version).
         let (request_sink, request_stream) = futures::channel::mpsc::unbounded();
         let requests_fut = tokio::spawn({
+            let tags = self.tags.clone();
             let cache = cache.clone();
             let client = client.clone();
             async move {
                 let mut response_stream = request_stream
                     .map({
                         |request: Request| match request {
-                            Request::Package(package_name) => Either::Left(
+                            Request::Package(package_name) => {
-                                client
+                                Either::Left(client.simple(package_name.clone()).map_ok(
-                                    // TODO(charlie): Remove this clone.
+                                    move |metadata| Response::Package(package_name, metadata),
-                                    .simple(package_name.clone())
+                                ))
-                                    .map_ok(move |metadata| {
+                            }
-                                        Response::Package(package_name, metadata)
-                                    }),
-                            ),
                             Request::Version(file) => Either::Right(
                                 client
-                                    // TODO(charlie): Remove this clone.
                                     .file(file.clone())
                                     .map_ok(move |metadata| Response::Version(file, metadata)),
                             ),
@@ -93,7 +90,49 @@ impl<'a> Resolver<'a> {
                         match response? {
                             Response::Package(package_name, metadata) => {
                                 trace!("Received package metadata for {}", package_name);
-                                cache.packages.insert(package_name.clone(), metadata);
+
+                                // Only bother storing platform-compatible wheels.
+                                let wheels: Vec<Wheel> = metadata
+                                    .files
+                                    .into_iter()
+                                    .filter_map(|file| {
+                                        let Ok(filename) =
+                                            WheelFilename::from_str(file.filename.as_str())
+                                        else {
+                                            debug!("Ignoring non-wheel: {}", file.filename);
+                                            return None;
+                                        };
+
+                                        let Ok(version) =
+                                            pep440_rs::Version::from_str(&filename.version)
+                                        else {
+                                            debug!("Ignoring invalid version: {}", file.filename);
+                                            return None;
+                                        };
+
+                                        if !filename.is_compatible(&tags) {
+                                            debug!(
+                                                "Ignoring wheel with incompatible tags: {}",
+                                                file.filename
+                                            );
+                                            return None;
+                                        }
+
+                                        Some(Wheel {
+                                            name: PackageName::normalize(&filename.distribution),
+                                            version,
+                                            file,
+                                        })
+                                    })
+                                    .collect();
+
+                                if wheels.is_empty() {
+                                    return Err(ResolveError::NoCompatibleDistributions(
+                                        package_name,
+                                    ));
+                                }
+
+                                cache.packages.insert(package_name.clone(), wheels);
                             }
                             Response::Version(file, metadata) => {
                                 trace!("Received file metadata for {}", file.filename);
@@ -184,6 +223,8 @@ impl<'a> Resolver<'a> {
             // Pick the next compatible version.
             let version = match decision.1 {
                 None => {
+                    debug!("No compatible version found for: {}", next);
+
                     let term_intersection = state
                         .partial_solution
                         .term_intersection_for_package(&next)
@@ -300,35 +341,18 @@ impl<'a> Resolver<'a> {
             };
 
             // Find a compatible version.
-            let simple_json = entry.value();
+            let wheels = entry.value();
-            let Some(file) = simple_json.files.iter().rev().find(|file| {
+            let Some(wheel) = wheels.iter().rev().find(|wheel| {
-                let Ok(name) = WheelFilename::from_str(file.filename.as_str()) else {
+                range
-                    return false;
-                };
-
-                let Ok(version) = pep440_rs::Version::from_str(&name.version) else {
-                    return false;
-                };
-
-                if !name.is_compatible(self.tags) {
-                    return false;
-                }
-
-                if !range
                     .borrow()
-                    .contains(&PubGrubVersion::from(version.clone()))
+                    .contains(&PubGrubVersion::from(wheel.version.clone()))
-                {
-                    return false;
-                };
-
-                true
             }) else {
                 continue;
             };
 
             // Emit a request to fetch the metadata for this version.
-            if in_flight.insert(file.hashes.sha256.clone()) {
+            if in_flight.insert(wheel.file.hashes.sha256.clone()) {
-                request_sink.unbounded_send(Request::Version(file.clone()))?;
+                request_sink.unbounded_send(Request::Version(wheel.file.clone()))?;
             }
 
             selection = index;
@@ -346,49 +370,45 @@ impl<'a> Resolver<'a> {
                 // TODO(charlie): Ideally, we'd choose the first package for which metadata is
                 // available.
                 let entry = cache.packages.wait(package_name).await.unwrap();
-                let simple_json = entry.value();
+                let wheels = entry.value();
+
+                debug!(
+                    "Searching for a compatible version of {} ({})",
+                    package_name,
+                    range.borrow()
+                );
 
                 // Find a compatible version.
-                let name_version_file = simple_json.files.iter().rev().find_map(|file| {
+                let wheel = wheels.iter().rev().find(|wheel| {
-                    let Ok(name) = WheelFilename::from_str(file.filename.as_str()) else {
+                    if range
-                        return None;
-                    };
-
-                    let Ok(version) = pep440_rs::Version::from_str(&name.version) else {
-                        return None;
-                    };
-
-                    if !name.is_compatible(self.tags) {
-                        return None;
-                    }
-
-                    if !range
                         .borrow()
-                        .contains(&PubGrubVersion::from(version.clone()))
+                        .contains(&PubGrubVersion::from(wheel.version.clone()))
                     {
-                        return None;
+                        true
-                    };
+                    } else {
-
+                        debug!("Ignoring non-satisfying version: {}", wheel.version);
-                    Some((package_name.clone(), version.clone(), file.clone()))
+                        false
+                    }
                 });
 
-                if let Some((name, version, file)) = name_version_file {
+                if let Some(wheel) = wheel {
-                    debug!("Selecting: {}=={} ({})", name, version, file.filename);
+                    debug!(
+                        "Selecting: {}=={} ({})",
+                        wheel.name, wheel.version, wheel.file.filename
+                    );
 
                     // We want to return a package pinned to a specific version; but we _also_ want to
                     // store the exact file that we selected to satisfy that version.
-                    pins.entry(name)
+                    pins.entry(wheel.name.clone())
                         .or_default()
-                        .insert(version.clone(), file.clone());
+                        .insert(wheel.version.clone(), wheel.file.clone());
 
                     // Emit a request to fetch the metadata for this version.
-                    if cache.versions.get(&file.hashes.sha256).is_none() {
+                    if in_flight.insert(wheel.file.hashes.sha256.clone()) {
-                        if in_flight.insert(file.hashes.sha256.clone()) {
+                        request_sink.unbounded_send(Request::Version(wheel.file.clone()))?;
-                            request_sink.unbounded_send(Request::Version(file.clone()))?;
-                        }
                     }
 
-                    Ok((package, Some(PubGrubVersion::from(version))))
+                    Ok((package, Some(PubGrubVersion::from(wheel.version.clone()))))
                 } else {
                     // We have metadata for the package, but no compatible version.
                     Ok((package, None))
@@ -426,9 +446,9 @@ impl<'a> Resolver<'a> {
             }
             PubGrubPackage::Package(package_name, extra) => {
                 if let Some(extra) = extra.as_ref() {
-                    debug!("Fetching dependencies for {}[{:?}]", package_name, extra);
+                    debug!("Fetching dependencies for: {}[{:?}]", package_name, extra);
                 } else {
-                    debug!("Fetching dependencies for {}", package_name);
+                    debug!("Fetching dependencies for: {}", package_name);
                 }
 
                 // Wait for the metadata to be available.
@@ -497,9 +517,19 @@ enum Response {
     Version(File, Metadata21),
 }
 
+#[derive(Debug, Clone)]
+struct Wheel {
+    /// The underlying [`File`] for this wheel.
+    file: File,
+    /// The normalized name of the package.
+    name: PackageName,
+    /// The version of the package.
+    version: pep440_rs::Version,
+}
+
 struct SolverCache {
-    /// A map from package name to the metadata for that package.
+    /// A map from package name to the wheels available for that package.
-    packages: WaitMap<PackageName, SimpleJson>,
+    packages: WaitMap<PackageName, Vec<Wheel>>,
 
     /// A map from wheel SHA to the metadata for that wheel.
     versions: WaitMap<String, Metadata21>,