perf: fix performance around extras

Cargo.toml

@@ -22,3 +22,7 @@ rust-version = "1.70"
 allow-branch = ["main"]
 consolidate-commits = true
 tag-prefix = ""
+
+#[patch.crates-io]
+#resolvo = { path="../resolvo" }
+#resolvo = { git="https://github.com/mamba-org/resolvo.git", branch="main" }

crates/rattler_installs_packages/Cargo.toml

@@ -27,6 +27,7 @@ futures = "0.3.28"
 http = "0.2.9"
 http-cache-semantics = { version = "1.0.1", default-features = false, features = ["with_serde", "reqwest"] }
 indexmap = "2.0.1"
+itertools = "0.11.0"
 miette = "5.10.0"
 mime = "0.3.17"
 once_cell = "1.18.0"
@@ -50,7 +51,6 @@ tracing = { version = "0.1.37", default-features = false, features = ["attribute
 url = { version = "2.4.1", features = ["serde"] }
 zip = "0.6.6"
 resolvo = { version = "0.1.0", optional = true }
-dirs = "5.0.1"
 
 [dev-dependencies]
 criterion = "0.5"

crates/rattler_installs_packages/src/resolve.rs

@@ -2,6 +2,7 @@ use crate::{
     CompareOp, Extra, NormalizedPackageName, PackageDb, PackageName, Requirement, Specifier,
     Specifiers, UserRequirement, Version, Wheel,
 };
+use itertools::Itertools;
 use resolvo::{
     Candidates, DefaultSolvableDisplay, Dependencies, DependencyProvider, NameId, Pool, SolvableId,
     Solver, SolverCache, VersionSet,
@@ -120,7 +121,7 @@ impl<'db> DependencyProvider<PypiVersionSet, PypiPackageName> for PypiDependency
 
     fn get_candidates(&self, name: NameId) -> Option<Candidates> {
         let package_name = self.pool.resolve_package_name(name);
-        tracing::info!("Fetching metadata for {}", package_name);
+        tracing::info!("collecting {}", package_name);
 
         // Get all the metadata for this package
         let result = task::block_in_place(move || {
@@ -139,6 +140,7 @@ impl<'db> DependencyProvider<PypiVersionSet, PypiPackageName> for PypiDependency
             }
         };
         let mut candidates = Candidates::default();
+        let mut no_wheels = Vec::new();
         for (version, artifacts) in artifacts.iter() {
             // Filter only artifacts we can work with
             let available_artifacts = artifacts
@@ -154,7 +156,7 @@ impl<'db> DependencyProvider<PypiVersionSet, PypiPackageName> for PypiDependency
             // Check if there are wheel artifacts for this version
             if available_artifacts.is_empty() {
                 // If there are no wheel artifacts, we're just gonna skip it
-                tracing::warn!("No available wheel artifact {package_name} {version} (skipping)");
+                no_wheels.push(version);
                 continue;
             }
 
@@ -165,21 +167,36 @@ impl<'db> DependencyProvider<PypiVersionSet, PypiPackageName> for PypiDependency
                 .collect::<Vec<_>>();
 
             if non_yanked_artifacts.is_empty() {
-                tracing::info!("{package_name} {version} was yanked (skipping)");
                 continue;
             }
             let solvable_id = self
                 .pool
                 .intern_solvable(name, PypiVersion(version.clone()));
             candidates.candidates.push(solvable_id);
         }
+
+        // Print some information about skipped packages
+        if !no_wheels.is_empty() && package_name.extra().is_none() {
+            tracing::warn!(
+                "Not considering {} {} because there are no wheel artifacts available",
+                package_name,
+                no_wheels.iter().format(", ")
+            );
+        }
+
         Some(candidates)
     }
 
     fn get_dependencies(&self, solvable: SolvableId) -> Dependencies {
         let solvable = self.pool.resolve_solvable(solvable);
         let package_name = self.pool.resolve_package_name(solvable.name_id());
 
+        tracing::info!(
+            "obtaining dependency information from {}={}",
+            package_name,
+            solvable.inner()
+        );
+
         // TODO: https://peps.python.org/pep-0508/#environment-markers
         let env = HashMap::from_iter([
             // TODO: We should add some proper values here.
@@ -257,6 +274,24 @@ impl<'db> DependencyProvider<PypiVersionSet, PypiPackageName> for PypiDependency
                 .unwrap()
         });
 
+        // Add constraints that restrict that the extra packages are set to the same version.
+        if let PypiPackageName::Base(package_name) = package_name {
+            // Add constraints on the extras of a package
+            for extra in metadata.extras {
+                let extra_name_id = self
+                    .pool
+                    .intern_package_name(PypiPackageName::Extra(package_name.clone(), extra));
+                let specifiers = Specifiers(vec![Specifier {
+                    op: CompareOp::Equal,
+                    value: solvable.inner().0.to_string(),
+                }]);
+                let version_set_id = self
+                    .pool
+                    .intern_version_set(extra_name_id, specifiers.into());
+                dependencies.constrains.push(version_set_id);
+            }
+        }
+
         for requirement in metadata.requires_dist {
             // Evaluate environment markers
             if let Some(env_marker) = &requirement.env_marker_expr {