Resolver module improvements (#9773)

Further small refactorings for the resolver.

crates/uv-resolver/src/resolver/batch_prefetch.rs

@@ -39,13 +39,31 @@ enum BatchPrefetchStrategy {
 /// have to fetch the metadata for a lot of versions.
 ///
 /// Note that these all heuristics that could totally prefetch lots of irrelevant versions.
-#[derive(Default)]
 pub(crate) struct BatchPrefetcher {
+    // Internal types.
     tried_versions: FxHashMap<PackageName, usize>,
     last_prefetch: FxHashMap<PackageName, usize>,
+    // Shared (e.g., `Arc`) types.
+    capabilities: IndexCapabilities,
+    index: InMemoryIndex,
+    request_sink: Sender<Request>,
 }
 
 impl BatchPrefetcher {
+    pub(crate) fn new(
+        capabilities: IndexCapabilities,
+        index: InMemoryIndex,
+        request_sink: Sender<Request>,
+    ) -> Self {
+        Self {
+            tried_versions: FxHashMap::default(),
+            last_prefetch: FxHashMap::default(),
+            capabilities,
+            index,
+            request_sink,
+        }
+    }
+
     /// Prefetch a large number of versions if we already unsuccessfully tried many versions.
     pub(crate) fn prefetch_batches(
         &mut self,
@@ -55,9 +73,6 @@ impl BatchPrefetcher {
         current_range: &Range<Version>,
         unchangeable_constraints: Option<&Term<Range<Version>>>,
         python_requirement: &PythonRequirement,
-        request_sink: &Sender<Request>,
-        in_memory: &InMemoryIndex,
-        capabilities: &IndexCapabilities,
         selector: &CandidateSelector,
         env: &ResolverEnvironment,
     ) -> anyhow::Result<(), ResolveError> {
@@ -79,12 +94,12 @@ impl BatchPrefetcher {
 
         // This is immediate, we already fetched the version map.
         let versions_response = if let Some(index) = index {
-            in_memory
+            self.index
                 .explicit()
                 .wait_blocking(&(name.clone(), index.clone()))
                 .ok_or_else(|| ResolveError::UnregisteredTask(name.to_string()))?
         } else {
-            in_memory
+            self.index
                 .implicit()
                 .wait_blocking(name)
                 .ok_or_else(|| ResolveError::UnregisteredTask(name.to_string()))?
@@ -166,7 +181,7 @@ impl BatchPrefetcher {
             // Avoid prefetching built distributions that don't support _either_ PEP 658 (`.metadata`)
             // or range requests.
             if !(wheel.file.dist_info_metadata
-                || capabilities.supports_range_requests(&wheel.index))
+                || self.capabilities.supports_range_requests(&wheel.index))
             {
                 debug!("Abandoning prefetch for {wheel} due to missing registry capabilities");
                 return Ok(());
@@ -214,9 +229,9 @@ impl BatchPrefetcher {
             );
             prefetch_count += 1;
 
-            if in_memory.distributions().register(candidate.version_id()) {
+            if self.index.distributions().register(candidate.version_id()) {
                 let request = Request::from(dist);
-                request_sink.blocking_send(request)?;
+                self.request_sink.blocking_send(request)?;
             }
         }
 

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

@@ -306,7 +306,11 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
 
         let root = PubGrubPackage::from(PubGrubPackageInner::Root(self.project.clone()));
         let pubgrub = State::init(root.clone(), MIN_VERSION.clone());
-        let mut prefetcher = BatchPrefetcher::default();
+        let mut prefetcher = BatchPrefetcher::new(
+            self.capabilities.clone(),
+            self.index.clone(),
+            request_sink.clone(),
+        );
         let state = ForkState::new(pubgrub, self.env.clone(), self.python_requirement.clone());
         let mut preferences = self.preferences.clone();
         let mut forked_states = self.env.initial_forked_states(state);
@@ -497,9 +501,6 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
                             .partial_solution
                             .unchanging_term_for_package(next_id),
                         &state.python_requirement,
-                        &request_sink,
-                        &self.index,
-                        &self.capabilities,
                         &self.selector,
                         &state.env,
                     )?;
@@ -1078,57 +1079,8 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
             }
         };
 
-        // Validate the Python requirement.
-        let requires_python = match dist {
-            CompatibleDist::InstalledDist(_) => None,
-            CompatibleDist::SourceDist { sdist, .. }
-            | CompatibleDist::IncompatibleWheel { sdist, .. } => {
-                sdist.file.requires_python.as_ref()
-            }
-            CompatibleDist::CompatibleWheel { wheel, .. } => wheel.file.requires_python.as_ref(),
-        };
-        let incompatibility = requires_python.and_then(|requires_python| {
-            if python_requirement.installed() == python_requirement.target() {
-                if !python_requirement
-                    .installed()
-                    .is_contained_by(requires_python)
-                {
-                    return if matches!(dist, CompatibleDist::CompatibleWheel { .. }) {
-                        Some(IncompatibleDist::Wheel(IncompatibleWheel::RequiresPython(
-                            requires_python.clone(),
-                            PythonRequirementKind::Installed,
-                        )))
-                    } else {
-                        Some(IncompatibleDist::Source(
-                            IncompatibleSource::RequiresPython(
-                                requires_python.clone(),
-                                PythonRequirementKind::Installed,
-                            ),
-                        ))
-                    };
-                }
-            } else {
-                if !python_requirement.target().is_contained_by(requires_python) {
-                    return if matches!(dist, CompatibleDist::CompatibleWheel { .. }) {
-                        Some(IncompatibleDist::Wheel(IncompatibleWheel::RequiresPython(
-                            requires_python.clone(),
-                            PythonRequirementKind::Target,
-                        )))
-                    } else {
-                        Some(IncompatibleDist::Source(
-                            IncompatibleSource::RequiresPython(
-                                requires_python.clone(),
-                                PythonRequirementKind::Target,
-                            ),
-                        ))
-                    };
-                }
-            }
-            None
-        });
-
-        // The version is incompatible due to its Python requirement.
-        if let Some(incompatibility) = incompatibility {
+        // Check whether the version is incompatible due to its Python requirement.
+        if let Some(incompatibility) = Self::check_requires_python(dist, python_requirement) {
             return Ok(Some(ResolverVersion::Unavailable(
                 candidate.version().clone(),
                 UnavailableVersion::IncompatibleDist(incompatibility),
@@ -1180,6 +1132,59 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
         Ok(Some(ResolverVersion::Available(version)))
     }
 
+    /// Check if the distribution is incompatible with the Python requirement, and if so, return
+    /// the incompatibility.
+    fn check_requires_python(
+        dist: &CompatibleDist,
+        python_requirement: &PythonRequirement,
+    ) -> Option<IncompatibleDist> {
+        let requires_python = match dist {
+            CompatibleDist::InstalledDist(_) => None,
+            CompatibleDist::SourceDist { sdist, .. }
+            | CompatibleDist::IncompatibleWheel { sdist, .. } => {
+                sdist.file.requires_python.as_ref()
+            }
+            CompatibleDist::CompatibleWheel { wheel, .. } => wheel.file.requires_python.as_ref(),
+        }?;
+        if python_requirement.installed() == python_requirement.target() {
+            if !python_requirement
+                .installed()
+                .is_contained_by(requires_python)
+            {
+                return if matches!(dist, CompatibleDist::CompatibleWheel { .. }) {
+                    Some(IncompatibleDist::Wheel(IncompatibleWheel::RequiresPython(
+                        requires_python.clone(),
+                        PythonRequirementKind::Installed,
+                    )))
+                } else {
+                    Some(IncompatibleDist::Source(
+                        IncompatibleSource::RequiresPython(
+                            requires_python.clone(),
+                            PythonRequirementKind::Installed,
+                        ),
+                    ))
+                };
+            }
+        } else {
+            if !python_requirement.target().is_contained_by(requires_python) {
+                return if matches!(dist, CompatibleDist::CompatibleWheel { .. }) {
+                    Some(IncompatibleDist::Wheel(IncompatibleWheel::RequiresPython(
+                        requires_python.clone(),
+                        PythonRequirementKind::Target,
+                    )))
+                } else {
+                    Some(IncompatibleDist::Source(
+                        IncompatibleSource::RequiresPython(
+                            requires_python.clone(),
+                            PythonRequirementKind::Target,
+                        ),
+                    ))
+                };
+            }
+        }
+        None
+    }
+
     /// Given a candidate package and version, return its dependencies.
     #[instrument(skip_all, fields(%package, %version))]
     fn get_dependencies_forking(