Use a single requirements iterator in sync (#71)

crates/puffin-cli/src/commands/compile.rs

@@ -52,7 +52,7 @@ pub(crate) async fn compile(src: &Path, cache: Option<&Path>) -> Result<ExitStat
 
     // Resolve the dependencies.
     let resolution = puffin_resolver::resolve(
-        &requirements,
+        requirements.iter(),
         markers,
         &tags,
         &client,

crates/puffin-cli/src/commands/sync.rs

@@ -8,7 +8,7 @@ use tracing::{debug, info};
 use platform_host::Platform;
 use platform_tags::Tags;
 use puffin_client::PypiClientBuilder;
-use puffin_installer::{Distribution, RemoteDistribution};
+use puffin_installer::{Distribution, LocalDistribution, LocalIndex, RemoteDistribution};
 use puffin_interpreter::{PythonExecutable, SitePackages};
 use puffin_package::package_name::PackageName;
 use puffin_package::requirements::Requirements;
@@ -42,25 +42,53 @@ pub(crate) async fn sync(src: &Path, cache: Option<&Path>, flags: SyncFlags) ->
         python.executable().display()
     );
 
-    // Remove any already-installed packages.
+    // Determine the current environment markers.
-    let requirements = if flags.intersects(SyncFlags::IGNORE_INSTALLED) {
+    let markers = python.markers();
-        requirements
+    let tags = Tags::from_env(python.platform(), python.simple_version())?;
+
+    // Index all the already-installed packages in site-packages.
+    let site_packages = if flags.intersects(SyncFlags::IGNORE_INSTALLED) {
+        SitePackages::default()
     } else {
-        let site_packages = SitePackages::from_executable(&python).await?;
+        SitePackages::from_executable(&python).await?
-        requirements.filter(|requirement| {
-            let package = PackageName::normalize(&requirement.name);
-            if let Some(version) = site_packages.get(&package) {
-                #[allow(clippy::print_stdout)]
-                {
-                    info!("Requirement already satisfied: {package} ({version})");
-                }
-                false
-            } else {
-                true
-            }
-        })
     };
 
+    // Index all the already-downloaded wheels in the cache.
+    let local_index = if let Some(cache) = cache {
+        LocalIndex::from_directory(cache).await?
+    } else {
+        LocalIndex::default()
+    };
+
+    let requirements = requirements
+        .iter()
+        .filter_map(|requirement| {
+            let package = PackageName::normalize(&requirement.name);
+
+            // Filter out already-installed packages.
+            if let Some(version) = site_packages.get(&package) {
+                info!("Requirement already satisfied: {package} ({version})");
+                return None;
+            }
+
+            // Identify any locally-available distributions that satisfy the requirement.
+            if let Some(distribution) = local_index
+                .get(&package)
+                .filter(|dist| requirement.is_satisfied_by(dist.version()))
+            {
+                debug!(
+                    "Requirement already cached: {} ({})",
+                    distribution.name(),
+                    distribution.version()
+                );
+                return Some(Requirement::Local(distribution.clone()));
+            }
+
+            debug!("Identified uncached requirement: {}", requirement);
+            Some(Requirement::Remote(requirement.clone()))
+        })
+        .collect::<Vec<_>>();
+
     if requirements.is_empty() {
         let s = if initial_requirements == 1 { "" } else { "s" };
         info!(
@@ -72,42 +100,7 @@ pub(crate) async fn sync(src: &Path, cache: Option<&Path>, flags: SyncFlags) ->
         return Ok(ExitStatus::Success);
     }
 
-    // Detect any cached wheels.
+    // Resolve the dependencies.
-    let (uncached, cached) = if let Some(cache) = cache {
-        let mut cached = Vec::with_capacity(requirements.len());
-        let mut uncached = Vec::with_capacity(requirements.len());
-
-        let index = puffin_installer::LocalIndex::from_directory(cache).await?;
-        for requirement in requirements {
-            let package = PackageName::normalize(&requirement.name);
-            if let Some(distribution) = index
-                .get(&package)
-                .filter(|dist| requirement.is_satisfied_by(dist.version()))
-            {
-                debug!(
-                    "Requirement already cached: {} ({})",
-                    distribution.name(),
-                    distribution.version()
-                );
-                cached.push(distribution.clone());
-            } else {
-                debug!("Identified uncached requirement: {}", requirement);
-                uncached.push(requirement);
-            }
-        }
-
-        (Requirements::new(uncached), cached)
-    } else {
-        (requirements, Vec::new())
-    };
-
-    // Determine the current environment markers.
-    let markers = python.markers();
-
-    // Determine the compatible platform tags.
-    let tags = Tags::from_env(python.platform(), python.simple_version())?;
-
-    // Instantiate a client.
     let client = {
         let mut pypi_client = PypiClientBuilder::default();
         if let Some(cache) = cache {
@@ -115,25 +108,27 @@ pub(crate) async fn sync(src: &Path, cache: Option<&Path>, flags: SyncFlags) ->
         }
         pypi_client.build()
     };
+    let resolution = puffin_resolver::resolve(
+        requirements
+            .iter()
+            .filter_map(|requirement| match requirement {
+                Requirement::Remote(requirement) => Some(requirement),
+                Requirement::Local(_) => None,
+            }),
+        markers,
+        &tags,
+        &client,
+        puffin_resolver::ResolveFlags::NO_DEPS,
+    )
+    .await?;
 
-    // Resolve the dependencies.
+    // Install the resolved distributions.
-    let resolution = if uncached.is_empty() {
+    let wheels = requirements
-        puffin_resolver::Resolution::empty()
-    } else {
-        puffin_resolver::resolve(
-            &uncached,
-            markers,
-            &tags,
-            &client,
-            puffin_resolver::ResolveFlags::NO_DEPS,
-        )
-        .await?
-    };
-
-    // Install into the current environment.
-    let wheels = cached
         .into_iter()
-        .map(|local| Ok(Distribution::Local(local)))
+        .filter_map(|requirement| match requirement {
+            Requirement::Remote(_) => None,
+            Requirement::Local(distribution) => Some(Ok(Distribution::Local(distribution))),
+        })
         .chain(
             resolution
                 .into_files()
@@ -152,3 +147,11 @@ pub(crate) async fn sync(src: &Path, cache: Option<&Path>, flags: SyncFlags) ->
 
     Ok(ExitStatus::Success)
 }
+
+#[derive(Debug)]
+enum Requirement {
+    /// A requirement that must be downloaded from PyPI.
+    Remote(pep508_rs::Requirement),
+    /// A requirement that is already available locally.
+    Local(LocalDistribution),
+}

crates/puffin-installer/src/index.rs

@@ -1,21 +1,21 @@
-use std::collections::BTreeMap;
+use std::collections::HashMap;
 use std::path::Path;
 
 use anyhow::Result;
 
-use crate::cache::WheelCache;
 use puffin_package::package_name::PackageName;
 
+use crate::cache::WheelCache;
 use crate::distribution::LocalDistribution;
 
 /// A local index of cached distributions.
-#[derive(Debug)]
+#[derive(Debug, Default)]
-pub struct LocalIndex(BTreeMap<PackageName, LocalDistribution>);
+pub struct LocalIndex(HashMap<PackageName, LocalDistribution>);
 
 impl LocalIndex {
     /// Build an index of cached distributions from a directory.
     pub async fn from_directory(path: &Path) -> Result<Self> {
-        let mut index = BTreeMap::new();
+        let mut index = HashMap::new();
 
         let cache = WheelCache::new(path);
         let Ok(mut dir) = cache.read_dir().await else {

crates/puffin-installer/src/lib.rs

@@ -1,4 +1,4 @@
-pub use distribution::{Distribution, RemoteDistribution};
+pub use distribution::{Distribution, LocalDistribution, RemoteDistribution};
 pub use index::LocalIndex;
 pub use install::install;
 

crates/puffin-interpreter/src/site_packages.rs

@@ -9,7 +9,7 @@ use puffin_package::package_name::PackageName;
 
 use crate::PythonExecutable;
 
-#[derive(Debug)]
+#[derive(Debug, Default)]
 pub struct SitePackages(BTreeMap<PackageName, Version>);
 
 impl SitePackages {

crates/puffin-resolver/src/lib.rs

@@ -14,17 +14,12 @@ use platform_tags::Tags;
 use puffin_client::{File, PypiClient, SimpleJson};
 use puffin_package::metadata::Metadata21;
 use puffin_package::package_name::PackageName;
-use puffin_package::requirements::Requirements;
 use wheel_filename::WheelFilename;
 
-#[derive(Debug)]
+#[derive(Debug, Default)]
 pub struct Resolution(HashMap<PackageName, PinnedPackage>);
 
 impl Resolution {
-    pub fn empty() -> Self {
-        Self(HashMap::new())
-    }
-
     /// Iterate over the pinned packages in this resolution.
     pub fn iter(&self) -> impl Iterator<Item = (&PackageName, &PinnedPackage)> {
         self.0.iter()
@@ -86,7 +81,7 @@ impl<T> From<futures::channel::mpsc::TrySendError<T>> for ResolveError {
 
 /// Resolve a set of requirements into a set of pinned versions.
 pub async fn resolve(
-    requirements: &Requirements,
+    requirements: impl Iterator<Item = &Requirement>,
     markers: &MarkerEnvironment,
     tags: &Tags,
     client: &PypiClient,
@@ -116,16 +111,20 @@ pub async fn resolve(
         .ready_chunks(32);
 
     // Push all the requirements into the package sink.
-    let mut in_flight: HashSet<PackageName> = HashSet::with_capacity(requirements.len());
+    let mut in_flight: HashSet<PackageName> = HashSet::new();
-    for requirement in requirements.iter() {
+    for requirement in requirements {
         debug!("--> adding root dependency: {}", requirement);
         package_sink.unbounded_send(Request::Package(requirement.clone()))?;
         in_flight.insert(PackageName::normalize(&requirement.name));
     }
 
+    if in_flight.is_empty() {
+        return Ok(Resolution::default());
+    }
+
     // Resolve the requirements.
     let mut resolution: HashMap<PackageName, PinnedPackage> =
-        HashMap::with_capacity(requirements.len());
+        HashMap::with_capacity(in_flight.len());
 
     while let Some(chunk) = package_stream.next().await {
         for result in chunk {